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