ys memos

Blog

Rustで簡単なJWTライブラリを実装してみた ~6. JWT~


rust

2024/08/31


  1. 概要と使用例
  2. プロジェクト設定
  3. base64関連
  4. 署名アルゴリズム関連
  5. ヘッダ定義
  6. JWT本体の実装
  7. JWT本体のTest

本Partでは、JWT本体の実装を行う。


src/jwt/mod.rs
mod base64;
pub mod header;
#[cfg(test)]
mod tests;

use base64::Base64;
use header::Header;

use crate::error::{Error, Result};

#[derive(Debug, Clone)]
pub struct Jwt<P>
where
    P: serde::ser::Serialize + serde::de::DeserializeOwned,
{
    header: Header,
    payload: P,
    signature: Vec<u8>,
}

impl<P> Jwt<P>
where
    P: serde::ser::Serialize + serde::de::DeserializeOwned,
{
    pub fn new(alg: header::Alg, payload: P) -> Self {
        Self {
            header: Header::new(alg),
            payload,
            signature: Vec::new(),
        }
    }

    pub fn decode<S: AsRef<str>>(token: S) -> Result<Self> {
        let parts = token.as_ref().split('.').collect::<Vec<_>>();
        if parts.len() != 3 {
            return Err(Error::DecodeInvalidParts);
        }
        Ok(Self {
            header: Base64::deserialize(parts[0])?,
            payload: Base64::deserialize(parts[1])?,
            signature: Base64::decode(parts[2])?,
        })
    }

    pub fn encode(self) -> Result<String> {
        if self.signature.is_empty() {
            return Err(Error::EncodeUnsigned);
        }
        Ok(dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
            Base64::encode(self.signature),
        ]))
    }

    pub fn payload(&self) -> &P {
        &self.payload
    }

    pub fn sign(&mut self, secret: &[u8]) -> Result<()> {
        let data = dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
        ]);
        self.signature = self.header.alg.sign(secret, data)?;
        Ok(())
    }

    pub fn verify(&self, secret: &[u8]) -> Result<bool> {
        let data = dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
        ]);
        self.header.alg.verify(secret, data, &self.signature)
    }
}

#[inline(always)]
fn dot_join(parts: &[String]) -> String {
    parts.join(".")
}


各種モジュールを読み込む。

mod base64;
pub mod header;
#[cfg(test)]
mod tests;

use base64::Base64;
use header::Header;

use crate::error::{Error, Result};

Payloadは、ジェネリクスで受けることで、入れるものはライブラリ利用側に委ねる。 その中でも、serdeSerialize/DeserializeOwnedを実装しているものを受け入れるようにしている。

今回は、 signatureVec<u8>で持つことにしているが、Option<Vec<u8>>という選択もあり得る。(より明示的な署名の有無に寄与する)

#[derive(Debug, Clone)]
pub struct Jwt<P>
where
    P: serde::ser::Serialize + serde::de::DeserializeOwned,
{
    header: Header,
    payload: P,
    signature: Vec<u8>,
}

new()はおなじみの初期化関数。algを受けるようにしているのは、現時点でTypがJWT固定なので、アルゴリズムのみを受け取るようにしている。これは、利用側に何を委ねるかの考え方によって変えることになりそう。 こちらは、絶対に成功する。

decode()は、トークンをデコードしてデシリアライズするための関数。base64デコードやデシリアライズに失敗する可能性がある。

impl<P> Jwt<P>
where
    P: serde::ser::Serialize + serde::de::DeserializeOwned,
{
    pub fn new(alg: header::Alg, payload: P) -> Self {
        Self {
            header: Header::new(alg),
            payload,
            signature: Vec::new(),
        }
    }

    pub fn decode<S: AsRef<str>>(token: S) -> Result<Self> {
        let parts = token.as_ref().split('.').collect::<Vec<_>>();
        if parts.len() != 3 {
            return Err(Error::DecodeInvalidParts);
        }
        Ok(Self {
            header: Base64::deserialize(parts[0])?,
            payload: Base64::deserialize(parts[1])?,
            signature: Base64::decode(parts[2])?,
        })
    }

encode()は、署名がない場合はエラーを返す。署名がある場合は、JWT(=ヘッダ、ペイロード、署名のbase64encode結果を.で結合したもの)を返す。

今回は、エンコードというシチュエーションを考えて、Jwtを消費するように書いたが、ここは、&selfで受け取るようにしてもいいかもしれない。

    pub fn encode(self) -> Result<String> {
        if self.signature.is_empty() {
            return Err(Error::EncodeUnsigned);
        }
        Ok(dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
            Base64::encode(self.signature),
        ]))
    }

トークンをデコードしたら、通常はペイロードを利用することが多いと考え、ペイロードを取得するメソッドを用意している。

(変更時にsignatureをリセットするなどを含む)セッターを作ってもいいかもしれない。

    pub fn payload(&self) -> &P {
        &self.payload
    }

sign()でヘッダ+ペイロードに署名を付与、verify()で署名を検証する。

    pub fn sign(&mut self, secret: &[u8]) -> Result<()> {
        let data = dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
        ]);
        self.signature = self.header.alg.sign(secret, data)?;
        Ok(())
    }

    pub fn verify(&self, secret: &[u8]) -> Result<bool> {
        let data = dot_join(&[
            Base64::serialize(&self.header)?,
            Base64::serialize(&self.payload)?,
        ]);
        self.header.alg.verify(secret, data, &self.signature)
    }
}

JWTでは、ドット結合を多用するので、そのためのヘルパー関数を用意した。

ヘルパーを用いずにformat!("{}.{}.{}", a, b, c)[a, b, c].join(".")を使うこともできるが、ヘルパー関数を用意することで、可読性の向上を図っている。

今回は、小さい一つのヘルパーだったのでJwt実装のおまけとして記述したが、より増えたり、大きくなったりしたらファイル分割してもいいかもしれない。

#[inline(always)]
fn dot_join(parts: &[String]) -> String {
    parts.join(".")
}


関連タグを探す