Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
8.6 kB · 280 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281use crate::compress::decode_scope;
use super::types::{ Claims, Header, SigningAlgorithm, TokenData, TokenDecodeError, TokenScope, TokenType, TokenVerifyError, UnsafeClaims,};use anyhow::{Context, Result, anyhow};use base64::Engine as _;use base64::engine::general_purpose::URL_SAFE_NO_PAD;use chrono::Utc;use hmac::{Hmac, Mac};use k256::ecdsa::{Signature, SigningKey, VerifyingKey, signature::Verifier};use sha2::Sha256;use subtle::ConstantTimeEq;use tranquil_types::{Did, Jti};
type HmacSha256 = Hmac<Sha256>;
pub fn get_did_from_token(token: &str) -> Result<Did, TokenDecodeError> { let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return Err(TokenDecodeError::InvalidFormat); }
let payload_bytes = URL_SAFE_NO_PAD .decode(parts[1]) .map_err(|_| TokenDecodeError::Base64DecodeFailed)?;
let claims: UnsafeClaims = serde_json::from_slice(&payload_bytes).map_err(|_| TokenDecodeError::JsonDecodeFailed)?;
Ok(claims.sub.unwrap_or(claims.iss))}
pub fn get_jti_from_token(token: &str) -> Result<Jti, TokenDecodeError> { let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return Err(TokenDecodeError::InvalidFormat); }
let payload_bytes = URL_SAFE_NO_PAD .decode(parts[1]) .map_err(|_| TokenDecodeError::Base64DecodeFailed)?;
let claims: serde_json::Value = serde_json::from_slice(&payload_bytes).map_err(|_| TokenDecodeError::JsonDecodeFailed)?;
claims .get("jti") .and_then(|j| j.as_str()) .map(Jti::new) .ok_or(TokenDecodeError::MissingClaim)}
pub fn get_algorithm_from_token(token: &str) -> Result<SigningAlgorithm, TokenDecodeError> { let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return Err(TokenDecodeError::InvalidFormat); }
let header_bytes = URL_SAFE_NO_PAD .decode(parts[0]) .map_err(|_| TokenDecodeError::Base64DecodeFailed)?;
let header: Header = serde_json::from_slice(&header_bytes).map_err(|_| TokenDecodeError::JsonDecodeFailed)?;
Ok(header.alg)}
pub fn verify_token(token: &str, key_bytes: &[u8]) -> Result<TokenData<Claims>> { verify_token_es256k(token, key_bytes, None, None).map_err(anyhow::Error::from)}
pub fn verify_access_token(token: &str, key_bytes: &[u8]) -> Result<TokenData<Claims>> { verify_token_es256k( token, key_bytes, Some(TokenType::Access), Some(&[ TokenScope::Access, TokenScope::AppPass, TokenScope::AppPassPrivileged, ]), ) .map_err(anyhow::Error::from)}
pub fn verify_refresh_token(token: &str, key_bytes: &[u8]) -> Result<TokenData<Claims>> { verify_token_es256k( token, key_bytes, Some(TokenType::Refresh), Some(&[TokenScope::Refresh]), ) .map_err(anyhow::Error::from)}
pub fn verify_access_token_hs256(token: &str, secret: &[u8]) -> Result<TokenData<Claims>> { verify_token_hs256_internal( token, secret, Some(TokenType::Access), Some(&[ TokenScope::Access, TokenScope::AppPass, TokenScope::AppPassPrivileged, ]), )}
pub fn verify_refresh_token_hs256(token: &str, secret: &[u8]) -> Result<TokenData<Claims>> { verify_token_hs256_internal( token, secret, Some(TokenType::Refresh), Some(&[TokenScope::Refresh]), )}
pub fn verify_token_es256k( token: &str, key_bytes: &[u8], expected_typ: Option<TokenType>, allowed_scopes: Option<&[TokenScope]>,) -> Result<TokenData<Claims>, TokenVerifyError> { let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return Err(TokenVerifyError::Invalid("Invalid token format")); }
let header_b64 = parts[0]; let claims_b64 = parts[1]; let signature_b64 = parts[2];
let header_bytes = URL_SAFE_NO_PAD .decode(header_b64) .map_err(|_| TokenVerifyError::Invalid("Base64 decode of header failed"))?;
let header: Header = serde_json::from_slice(&header_bytes) .map_err(|_| TokenVerifyError::Invalid("JSON decode of header failed"))?;
if let Some(expected) = expected_typ && header.typ != expected { return Err(TokenVerifyError::Invalid("Invalid token type")); }
let signature_bytes = URL_SAFE_NO_PAD .decode(signature_b64) .map_err(|_| TokenVerifyError::Invalid("Base64 decode of signature failed"))?;
let signature = Signature::from_slice(&signature_bytes) .map_err(|_| TokenVerifyError::Invalid("Invalid signature format"))?;
let signing_key = SigningKey::from_slice(key_bytes) .map_err(|_| TokenVerifyError::Invalid("Invalid signing key"))?; let verifying_key = VerifyingKey::from(&signing_key);
let message = format!("{}.{}", header_b64, claims_b64); verifying_key .verify(message.as_bytes(), &signature) .map_err(|_| TokenVerifyError::Invalid("Signature verification failed"))?;
let claims_bytes = URL_SAFE_NO_PAD .decode(claims_b64) .map_err(|_| TokenVerifyError::Invalid("Base64 decode of claims failed"))?;
let mut claims: Claims = serde_json::from_slice(&claims_bytes) .map_err(|_| TokenVerifyError::Invalid("JSON decode of claims failed"))?;
if let Some(scope) = &claims.scope { claims.scope = Some( decode_scope(scope).map_err(|_| TokenVerifyError::Invalid("Invalid token scope"))?, ); }
let now = Utc::now().timestamp(); if claims.exp < now { return Err(TokenVerifyError::Expired); }
if let Some(scopes) = allowed_scopes { let token_scope: TokenScope = claims .scope .as_deref() .unwrap_or("") .parse() .unwrap_or_else(|e| match e {}); if !scopes.contains(&token_scope) { return Err(TokenVerifyError::Invalid("Invalid token scope")); } }
Ok(TokenData { claims })}
fn verify_token_hs256_internal( token: &str, secret: &[u8], expected_typ: Option<TokenType>, allowed_scopes: Option<&[TokenScope]>,) -> Result<TokenData<Claims>> { let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { return Err(anyhow!("Invalid token format")); }
let header_b64 = parts[0]; let claims_b64 = parts[1]; let signature_b64 = parts[2];
let header_bytes = URL_SAFE_NO_PAD .decode(header_b64) .context("Base64 decode of header failed")?;
let header: Header = serde_json::from_slice(&header_bytes).context("JSON decode of header failed")?;
if header.alg != SigningAlgorithm::HS256 { return Err(anyhow!("Expected HS256 algorithm, got {}", header.alg)); }
if let Some(expected) = expected_typ && header.typ != expected { return Err(anyhow!( "Invalid token type: expected {}, got {}", expected, header.typ )); }
let signature_bytes = URL_SAFE_NO_PAD .decode(signature_b64) .context("Base64 decode of signature failed")?;
let message = format!("{}.{}", header_b64, claims_b64);
let mut mac = HmacSha256::new_from_slice(secret).map_err(|e| anyhow!("Invalid secret: {}", e))?; mac.update(message.as_bytes());
let expected_signature = mac.finalize().into_bytes(); let is_valid: bool = signature_bytes.ct_eq(&expected_signature).into();
if !is_valid { return Err(anyhow!("Signature verification failed")); }
let claims_bytes = URL_SAFE_NO_PAD .decode(claims_b64) .context("Base64 decode of claims failed")?;
let mut claims: Claims = serde_json::from_slice(&claims_bytes).context("JSON decode of claims failed")?;
if let Some(scope) = &claims.scope { claims.scope = Some(decode_scope(scope).context("Invalid scope claim encoding")?); }
let now = Utc::now().timestamp(); if claims.exp < now { return Err(anyhow!("Token expired")); }
if let Some(scopes) = allowed_scopes { let token_scope: TokenScope = claims .scope .as_deref() .unwrap_or("") .parse() .unwrap_or_else(|e| match e {}); if !scopes.contains(&token_scope) { return Err(anyhow!("Invalid token scope: {}", token_scope)); } }
Ok(TokenData { claims })}