use 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 })
}