Something went wrong. Try again.
Microservice to bring 2FA to self hosted PDSes
Something went wrong. Try again.
5.1 kB · 116 lines
Rust
at new-2fa
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117use crate::helpers::json_error_response;use axum::extract::Request;use axum::http::{HeaderMap, StatusCode};use axum::middleware::Next;use axum::response::IntoResponse;use jwt_compact::alg::{Hs256, Hs256Key};use jwt_compact::{AlgorithmExt, Claims, Token, UntrustedToken, ValidationError};use serde::{Deserialize, Serialize};use std::env;use tracing::log;
#[derive(Clone, Debug)]pub struct Did(pub Option<String>);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum AuthScheme { Bearer, DPoP,}
#[derive(Serialize, Deserialize)]pub struct TokenClaims { pub sub: String,}
pub async fn extract_did(mut req: Request, next: Next) -> impl IntoResponse { let auth = extract_auth(req.headers());
match auth { Ok(auth_opt) => { match auth_opt { None => json_error_response(StatusCode::BAD_REQUEST, "TokenRequired", "") .expect("Error creating an error response"), Some((scheme, token_str)) => { // For Bearer, validate JWT and extract DID from `sub`. // For DPoP, we currently only pass through and do not validate here; insert None DID. match scheme { AuthScheme::Bearer => { let token = UntrustedToken::new(&token_str); if token.is_err() { return json_error_response( StatusCode::BAD_REQUEST, "TokenRequired", "", ) .expect("Error creating an error response"); } let parsed_token = token.expect("Already checked for error"); let claims: Result<Claims<TokenClaims>, ValidationError> = parsed_token.deserialize_claims_unchecked(); if claims.is_err() { return json_error_response( StatusCode::BAD_REQUEST, "TokenRequired", "", ) .expect("Error creating an error response"); }
let key = Hs256Key::new( env::var("PDS_JWT_SECRET") .expect("PDS_JWT_SECRET not set in the pds.env"), ); let token: Result<Token<TokenClaims>, ValidationError> = Hs256.validator(&key).validate(&parsed_token); if token.is_err() { return json_error_response( StatusCode::BAD_REQUEST, "InvalidToken", "", ) .expect("Error creating an error response"); } let token = token.expect("Already checked for error,"); req.extensions_mut() .insert(Did(Some(token.claims().custom.sub.clone()))); } AuthScheme::DPoP => { //Not going to worry about oauth email update for now, just always forward to the PDS req.extensions_mut().insert(Did(None)); } }
next.run(req).await } } } Err(err) => { log::error!("Error extracting token: {err}"); json_error_response(StatusCode::BAD_REQUEST, "InvalidToken", "") .expect("Error creating an error response") } }}
fn extract_auth(headers: &HeaderMap) -> Result<Option<(AuthScheme, String)>, String> { match headers.get(axum::http::header::AUTHORIZATION) { None => Ok(None), Some(hv) => { match hv.to_str() { Err(_) => Err("Authorization header is not valid".into()), Ok(s) => { // Accept forms like: "Bearer <token>" or "DPoP <token>" (case-sensitive for the scheme here) let mut parts = s.splitn(2, ' '); match (parts.next(), parts.next()) { (Some("Bearer"), Some(tok)) if !tok.is_empty() => Ok(Some((AuthScheme::Bearer, tok.to_string()))), (Some("DPoP"), Some(tok)) if !tok.is_empty() => Ok(Some((AuthScheme::DPoP, tok.to_string()))), _ => Err("Authorization header must be in format 'Bearer <token>' or 'DPoP <token>'".into()), } } } } }}