Zenn
🍪

Rust | axumにCSRF対応のトークンを組み込む

2024/08/15に公開
Rust
axum
tech

概要

「Rust | axumでモノリシックなWEBサイトを作るための色々」の続き

axumでモノリシックなWEBサイトを作る際に使えそうなネタのメモ。
今回はCSRF対策について
試すクレートは下記２点

  • axum_csrf
  • csrf

axum_csrfクレートでCSRFトークンを追加する

axum_csrfを利用してCSRFトークンを利用できるようにする

Cargo.tomlへ下記を追加

axum_csrf = { version = "0.9.0", features = ["layer"] }

main.rsへaxum_csrf関連のコード追加

use axum_csrf::CsrfConfig;
// 略

#[tokio::main]
async fn main() {

    // axum_csrf
    let csrf_config = CsrfConfig::default();  // 追加

    let app = Router::new()
        .route("/csrf", get(csrf))
        .layer(CsrfLayer::new(csrf_config)); // 追加

    let listener = TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

生成されたトークンをHTML側へわたすようにする

askamaを利用する
（知らない間に、askama_axumとか便利なものがあるので利用）
Cargo.tomlへ下記を追加

askama = {version = "0.12.1"}
askama_axum = "0.4.0"
serde = { version = "1.0.197", features = ["derive"] }

テンプレート用のHTML用意
templates/csrf.html

<!DOCTYPE html>
<html lang="ja">
<head>
    <meta charset="UTF-8">
    <title>axum_csrf</title>
</head>
<body>
    <h2>CSRF</h2>
    <form method="post" action="/csrf">
        <input type="hidden" name="authenticity_token" value="{{ authenticity_token }}"/>
        <input id="button" type="submit" value="Submit" />
    </form>
</body>
</html>

main.rsへ追記

#[derive(Template, Deserialize, Serialize)]
#[template(path = "csrf.html")]
struct Keys {
    authenticity_token: String,
}

async fn csrf(token: CsrfToken) -> impl IntoResponse {
    let keys = Keys {
        authenticity_token: token.authenticity_token().unwrap(),
    };

    (token, keys).into_response()
}

実行して、/csrfにアクセスすると、「Submit」ボタンが表示されるはず。
hiddenでトークンがセットされたものもあるはず。

トークンチェック部分を追加

main.rsを変更

#[tokio::main]
async fn main() {

    // axum_csrf
    let csrf_config = CsrfConfig::default();

    let app = Router::new()
        .route("/csrf", get(csrf).post(csrf_post))  // 変更
        .layer(CsrfLayer::new(csrf_config));

    let listener = TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn csrf_post(token: CsrfToken, Form(payload): Form<Keys>) -> impl IntoResponse {
    // トークンチェック
    if token.verify(&payload.authenticity_token).is_err() {
        "トークンチェック NG".into_response()
    } else {
        "トークンチェック OK".into_response()
    }
}

Submitボタンをクリックした際、「トークンチェック OK」が表示されれば成功。

ただ、トークンの生成がセキュアな関係でちょっと重たい。
せっかくのRustなので、速い方が良いので別の方法も試す。

csrfクレートでCSRFトークンを追加する

csrfをaxumで利用

Cargo.tomlに下記を追加

csrf = "0.4.1"
cookie = "0.18.1"
data-encoding = "2.5.0"
anyhow = "1.0.81"

main.rsで利用の準備

use std::sync::Arc;
use rand::{Rng, thread_rng};
use csrf::{AesGcmCsrfProtection, CsrfProtection};

// 略

// CSRF
let mut rand = thread_rng();
let csrf_protection_key: [u8; 32] = rand.gen();
let protect = Arc::new(AesGcmCsrfProtection::from_key(csrf_protection_key));

Routerにセット。
/csrf_defaultで確認画面が表示できるようにする

let app = Router::new()
    // 追加
    .route("/csrf_default", get(csrf_default))
    .route("/", get(index))
    .route("/redirect", get(redirect))
    .layer(CsrfLayer::new(csrf_config))
    .layer(SessionLayer::new(session_store))
    .layer(
        ServiceBuilder::new()
            .layer(AddExtensionLayer::new(pool))
            .into_inner()
    )
    .nest_service("/assets", ServeDir::new("assets"))
    // 追加
    .layer(
        ServiceBuilder::new()
            .layer(AddExtensionLayer::new(protect))
            .into_inner()
    );

let listener = TcpListener::bind("127.0.0.1:3000").await.unwrap();

axum::serve(listener, app).await.unwrap();

続いて画面まわり

/csrf_defaultにアクセスした際に、tokenを生成するよう実装
[main.rs]

const CSRF_KEY: &str = "axum_csrf";
type ArcCsrProtection = Arc<AesGcmCsrfProtection>;

#[derive(Template, Deserialize, Serialize)]
#[template(path = "csrf.html")]
struct Keys {
    authenticity_token: String,
}

async fn csrf_default(
    Extension(protect): Extension<ArcCsrProtection>,
) -> impl IntoResponse {
    
    let (token, cookie) = protect.generate_token_pair(None, 300)
        .expect("token/cookie pair 生成失敗");

    let token_str = token.b64_string();
    let cookie_str = cookie.b64_string();

    let keys = Keys {
        authenticity_token: token_str.to_string(),
    };

    // Cookieに値設定
    let mut response = (keys).into_response();

    let cookie = Cookie::build((CSRF_KEY, cookie_str.to_string()))
        .http_only(true)
        .build();

    response.headers_mut().append(SET_COOKIE, cookie.encoded().to_string().parse().unwrap());

    response
}

流れとしては、protect.generate_token_pairでPOST用とCookie格納用の値を生成。
Keysの構造体を作り、画面側に値を渡してaskamaでテンプレートに配置。
Cookie格納用はそのままCookieへセット

なお、テンプレートはaxum_csrfで利用したものを流用
[csrf.html]

<!DOCTYPE html>
<html lang="ja">
<head>
    <meta charset="UTF-8">
    <title>axum_csrf</title>
</head>
<body>
    <h2>CSRF</h2>
    <form method="post" action="/csrf_default">
        <input type="hidden" name="authenticity_token" value="{{ authenticity_token }}"/>
        <input id="button" type="submit" value="Submit" />
    </form>
</body>
</html>

実行して/csrf_defaultにアクセスすれば、hidden項目にトークンの値がセットされ、Cookieのaxum_csrfにも値がセットされているはず。

POST時の処理

POSTした際の処理を作成していきます。
まずは、POSTを受け付けるhandlerを追加

let app = Router::new()
    .route("/csrf_default", get(csrf_default).post(csrf_default_post))
    // 略

POSTされてくるトークン受け取り用の構造体を用意

#[derive(Deserialize, Debug)]
struct PostForm {
    authenticity_token: String,
}

POSTされたトークンと、Cookieの値を利用して有効性チェックを実施する

async fn csrf_default_post(
    headers: HeaderMap,
    Extension(protect): Extension<ArcCsrProtection>,
    Form(form): Form<PostForm>,
) -> impl IntoResponse {
    let token_str = form.authenticity_token;
    let token_bytes = BASE64.decode(token_str.as_bytes()).expect("token not base64");

    let cookie_header = headers.get("cookie").unwrap();
    let cookie_str = cookie_header.to_str().unwrap();
    let mut jar = CookieJar::new();

    for cookie_str in cookie_str.split("; ") {
        let cookie = Cookie::parse_encoded(cookie_str).unwrap().into_owned();
        jar.add(cookie);
    }

    let value = jar.get(CSRF_KEY).expect("axum_csrfの取り出し失敗");
    let cookie_byte = BASE64.decode(value.value().as_bytes()).expect("cookieのdecode失敗");

    let parsed_token = protect.parse_token(&token_bytes).expect("token_byteのparse失敗");
    let parsed_cookie = protect.parse_cookie(&cookie_byte).expect("cookie_byteのparse失敗");

    if protect.verify_token_pair(&parsed_token, &parsed_cookie) {
        println!("成功");
    } else {
        println!("失敗")
    }

    "CSRFチェック".into_response()
}

トークンチェックに成功したら、ログに成功が出力されます。

これを各画面に組み込んでいけば、CSRF対策ができるはずだけど、axum_csrfと比べて使い勝手に難あり。

csrfクレートも良い感じでaxumに組み込めるようにしてみる

csrfクレートをaxumに良い感じで組み込めるように調整してみる

Cookieに値を保持する部分の調整

まずは、各handlerでCookieに値を保持する部分について調整を行う。
処理を共通化して呼ぶのもイケてない気がするので、axum_csrfと同様にinto_responseのタイミングで保持するようにする

[my_csrf_token.rs]

#[derive(Clone)]
pub struct MyCSrfToken {
    authenticity_token: String,
}

impl IntoResponseParts for MyCSrfToken {
    type Error = Infallible;

    fn into_response_parts(self, mut res: ResponseParts) -> Result<ResponseParts, Self::Error> {

        let cookie_str = &self.authenticity_cookie;
        let cooke = Cookie::build((CSRF_KEY, cookie_str.to_string()))
            .http_only(true)
            .build();

        res.headers_mut().append(SET_COOKIE, cooke.encoded().to_string().parse().unwrap());
        Ok(res)
    }
}

authenticity_tokenに保持していたトークンをCookieにセットする

これに合わせて、csrf_default側を変更すれば、csrf_default内にてCookieに値設定を行っていた部分を共通的に実施できるようになるので、ちょっとイケてる感じになる

[main.rs]

async fn csrf_default(
    Extension(protect): Extension<ArcCsrProtection>,
) -> impl IntoResponse {

    let (token, cookie) = protect.generate_token_pair(None, 300)
        .expect("token/cookie pair 生成失敗");

    let token_str = token.b64_string();
    let cookie_str = cookie.b64_string();

    let keys = Keys {
        authenticity_token: token_str.to_string(),
    };
    
    let my_csrf_token = MyCSrfToken {
        authenticity_cookie: cookie_str
    };

    (my_csrf_token, keys).into_response();
}

MyCsrTokeにCookieをセットする部分をミドルウェアで行えば更にスッキリできるけど、今回はここまで。

Discussion