Microservice to bring 2FA to self hosted PDSes
at sendmail 5.1 kB view raw
1use crate::helpers::json_error_response; 2use axum::extract::Request; 3use axum::http::{HeaderMap, StatusCode}; 4use axum::middleware::Next; 5use axum::response::IntoResponse; 6use jwt_compact::alg::{Hs256, Hs256Key}; 7use jwt_compact::{AlgorithmExt, Claims, Token, UntrustedToken, ValidationError}; 8use serde::{Deserialize, Serialize}; 9use std::env; 10use tracing::log; 11 12#[derive(Clone, Debug)] 13pub struct Did(pub Option<String>); 14 15#[derive(Clone, Copy, Debug, PartialEq, Eq)] 16pub enum AuthScheme { 17 Bearer, 18 DPoP, 19} 20 21#[derive(Serialize, Deserialize)] 22pub struct TokenClaims { 23 pub sub: String, 24} 25 26pub async fn extract_did(mut req: Request, next: Next) -> impl IntoResponse { 27 let auth = extract_auth(req.headers()); 28 29 match auth { 30 Ok(auth_opt) => { 31 match auth_opt { 32 None => json_error_response(StatusCode::BAD_REQUEST, "TokenRequired", "") 33 .expect("Error creating an error response"), 34 Some((scheme, token_str)) => { 35 // For Bearer, validate JWT and extract DID from `sub`. 36 // For DPoP, we currently only pass through and do not validate here; insert None DID. 37 match scheme { 38 AuthScheme::Bearer => { 39 let token = UntrustedToken::new(&token_str); 40 if token.is_err() { 41 return json_error_response( 42 StatusCode::BAD_REQUEST, 43 "TokenRequired", 44 "", 45 ) 46 .expect("Error creating an error response"); 47 } 48 let parsed_token = token.expect("Already checked for error"); 49 let claims: Result<Claims<TokenClaims>, ValidationError> = 50 parsed_token.deserialize_claims_unchecked(); 51 if claims.is_err() { 52 return json_error_response( 53 StatusCode::BAD_REQUEST, 54 "TokenRequired", 55 "", 56 ) 57 .expect("Error creating an error response"); 58 } 59 60 let key = Hs256Key::new( 61 env::var("PDS_JWT_SECRET") 62 .expect("PDS_JWT_SECRET not set in the pds.env"), 63 ); 64 let token: Result<Token<TokenClaims>, ValidationError> = 65 Hs256.validator(&key).validate(&parsed_token); 66 if token.is_err() { 67 return json_error_response( 68 StatusCode::BAD_REQUEST, 69 "InvalidToken", 70 "", 71 ) 72 .expect("Error creating an error response"); 73 } 74 let token = token.expect("Already checked for error,"); 75 req.extensions_mut() 76 .insert(Did(Some(token.claims().custom.sub.clone()))); 77 } 78 AuthScheme::DPoP => { 79 //Not going to worry about oauth email update for now, just always forward to the PDS 80 req.extensions_mut().insert(Did(None)); 81 } 82 } 83 84 next.run(req).await 85 } 86 } 87 } 88 Err(err) => { 89 log::error!("Error extracting token: {err}"); 90 json_error_response(StatusCode::BAD_REQUEST, "InvalidToken", "") 91 .expect("Error creating an error response") 92 } 93 } 94} 95 96fn extract_auth(headers: &HeaderMap) -> Result<Option<(AuthScheme, String)>, String> { 97 match headers.get(axum::http::header::AUTHORIZATION) { 98 None => Ok(None), 99 Some(hv) => { 100 match hv.to_str() { 101 Err(_) => Err("Authorization header is not valid".into()), 102 Ok(s) => { 103 // Accept forms like: "Bearer <token>" or "DPoP <token>" (case-sensitive for the scheme here) 104 let mut parts = s.splitn(2, ' '); 105 match (parts.next(), parts.next()) { 106 (Some("Bearer"), Some(tok)) if !tok.is_empty() => 107 Ok(Some((AuthScheme::Bearer, tok.to_string()))), 108 (Some("DPoP"), Some(tok)) if !tok.is_empty() => 109 Ok(Some((AuthScheme::DPoP, tok.to_string()))), 110 _ => Err("Authorization header must be in format 'Bearer <token>' or 'DPoP <token>'".into()), 111 } 112 } 113 } 114 } 115 } 116}