aa_auth/
jwt.rs

1//! JWT signing and verification using HMAC-SHA256.
2
3use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
4use serde::{Deserialize, Serialize};
5use thiserror::Error;
6
7use super::scope::Scope;
8
9/// JWT token expiry duration: 24 hours in seconds.
10const TOKEN_EXPIRY_SECS: u64 = 24 * 60 * 60;
11
12/// JWT claims payload.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct Claims {
15    /// Subject: the API key ID that this token was issued for.
16    pub sub: String,
17    /// Issued-at timestamp (Unix epoch seconds).
18    pub iat: u64,
19    /// Expiry timestamp (Unix epoch seconds).
20    pub exp: u64,
21    /// Scopes granted to this token.
22    pub scope: Vec<Scope>,
23    /// AAASM-3139 — the team this token is scoped to. When present, a
24    /// non-admin caller is confined to its own team for per-tenant data
25    /// (e.g. `/costs`, `/agents/{id}/budget`). `None` on legacy tokens
26    /// issued before tenant claims existed, which therefore carry no team
27    /// scope (and remain admin-gated for cross-tenant data).
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub team_id: Option<String>,
30    /// AAASM-3139 — the org this token is scoped to. See [`Claims::team_id`].
31    #[serde(default, skip_serializing_if = "Option::is_none")]
32    pub org_id: Option<String>,
33}
34
35/// Signs JWTs using HMAC-SHA256.
36pub struct JwtSigner {
37    encoding_key: EncodingKey,
38}
39
40impl JwtSigner {
41    /// Create a new signer from a raw HMAC secret.
42    pub fn new(secret: &[u8]) -> Self {
43        Self {
44            encoding_key: EncodingKey::from_secret(secret),
45        }
46    }
47
48    /// Sign a new JWT for the given API key ID and scopes.
49    ///
50    /// The token expires after 24 hours.
51    pub fn sign(&self, key_id: &str, scopes: &[Scope]) -> Result<String, JwtError> {
52        self.sign_with_tenant(key_id, scopes, None, None)
53    }
54
55    /// Sign a JWT that additionally carries a tenant scope (AAASM-3139).
56    ///
57    /// `team_id` / `org_id` confine a non-admin caller to its own tenant for
58    /// per-tenant data endpoints. Pass `None` for either to leave it unscoped.
59    /// The token expires after 24 hours.
60    pub fn sign_with_tenant(
61        &self,
62        key_id: &str,
63        scopes: &[Scope],
64        team_id: Option<String>,
65        org_id: Option<String>,
66    ) -> Result<String, JwtError> {
67        let now = now_epoch_secs();
68        let claims = Claims {
69            sub: key_id.to_string(),
70            iat: now,
71            exp: now + TOKEN_EXPIRY_SECS,
72            scope: scopes.to_vec(),
73            team_id,
74            org_id,
75        };
76        encode(&Header::default(), &claims, &self.encoding_key).map_err(JwtError::Encode)
77    }
78
79    /// Sign a JWT with a custom expiry (for testing).
80    #[cfg(test)]
81    fn sign_with_expiry(&self, key_id: &str, scopes: &[Scope], exp: u64) -> Result<String, JwtError> {
82        let claims = Claims {
83            sub: key_id.to_string(),
84            iat: now_epoch_secs(),
85            exp,
86            scope: scopes.to_vec(),
87            team_id: None,
88            org_id: None,
89        };
90        encode(&Header::default(), &claims, &self.encoding_key).map_err(JwtError::Encode)
91    }
92}
93
94/// Verifies JWT tokens using HMAC-SHA256.
95pub struct JwtVerifier {
96    decoding_key: DecodingKey,
97    validation: Validation,
98}
99
100impl JwtVerifier {
101    /// Create a new verifier from a raw HMAC secret.
102    pub fn new(secret: &[u8]) -> Self {
103        let mut validation = Validation::default();
104        validation.set_required_spec_claims(&["sub", "iat", "exp"]);
105        Self {
106            decoding_key: DecodingKey::from_secret(secret),
107            validation,
108        }
109    }
110
111    /// Verify a JWT token and return its claims.
112    ///
113    /// Returns an error if the signature is invalid or the token has expired.
114    pub fn verify(&self, token: &str) -> Result<Claims, JwtError> {
115        let data = decode::<Claims>(token, &self.decoding_key, &self.validation).map_err(JwtError::Decode)?;
116        Ok(data.claims)
117    }
118}
119
120/// Return the current Unix epoch timestamp in seconds.
121fn now_epoch_secs() -> u64 {
122    std::time::SystemTime::now()
123        .duration_since(std::time::UNIX_EPOCH)
124        .expect("system clock before Unix epoch")
125        .as_secs()
126}
127
128/// Errors related to JWT operations.
129#[derive(Debug, Error)]
130pub enum JwtError {
131    #[error("failed to encode JWT: {0}")]
132    Encode(jsonwebtoken::errors::Error),
133    #[error("failed to decode JWT: {0}")]
134    Decode(jsonwebtoken::errors::Error),
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    const TEST_SECRET: &[u8] = b"test-secret-key-that-is-at-least-32-bytes-long!!";
142
143    #[test]
144    fn test_jwt_sign_verify_roundtrip() {
145        let signer = JwtSigner::new(TEST_SECRET);
146        let verifier = JwtVerifier::new(TEST_SECRET);
147
148        let token = signer
149            .sign("key-123", &[Scope::Read, Scope::Write])
150            .expect("signing should succeed");
151        let claims = verifier.verify(&token).expect("verification should succeed");
152
153        assert_eq!(claims.sub, "key-123");
154        assert_eq!(claims.scope, vec![Scope::Read, Scope::Write]);
155        assert!(claims.exp > claims.iat);
156    }
157
158    #[test]
159    fn test_jwt_expired_token_rejected() {
160        let signer = JwtSigner::new(TEST_SECRET);
161        let verifier = JwtVerifier::new(TEST_SECRET);
162
163        // Create a token that expired 1 hour ago.
164        let expired = now_epoch_secs() - 3600;
165        let token = signer
166            .sign_with_expiry("key-123", &[Scope::Read], expired)
167            .expect("signing should succeed");
168
169        let result = verifier.verify(&token);
170        assert!(result.is_err(), "expired token should be rejected");
171    }
172
173    #[test]
174    fn test_jwt_wrong_secret_rejected() {
175        let signer = JwtSigner::new(TEST_SECRET);
176        let verifier = JwtVerifier::new(b"different-secret-that-is-also-32-bytes-long!!");
177
178        let token = signer.sign("key-123", &[Scope::Read]).expect("signing should succeed");
179
180        let result = verifier.verify(&token);
181        assert!(result.is_err(), "token signed with different secret should be rejected");
182    }
183
184    #[test]
185    fn test_jwt_tenant_claim_roundtrip() {
186        // AAASM-3139: a tenant-scoped token must carry team_id back through
187        // verification; a plain token leaves the tenant claim None.
188        let signer = JwtSigner::new(TEST_SECRET);
189        let verifier = JwtVerifier::new(TEST_SECRET);
190
191        let scoped = signer
192            .sign_with_tenant("key-1", &[Scope::Read], Some("alpha".into()), Some("org-1".into()))
193            .unwrap();
194        let claims = verifier.verify(&scoped).unwrap();
195        assert_eq!(claims.team_id.as_deref(), Some("alpha"));
196        assert_eq!(claims.org_id.as_deref(), Some("org-1"));
197
198        let plain = signer.sign("key-2", &[Scope::Read]).unwrap();
199        let plain_claims = verifier.verify(&plain).unwrap();
200        assert_eq!(plain_claims.team_id, None);
201        assert_eq!(plain_claims.org_id, None);
202    }
203
204    #[test]
205    fn test_jwt_scopes_preserved() {
206        let signer = JwtSigner::new(TEST_SECRET);
207        let verifier = JwtVerifier::new(TEST_SECRET);
208
209        let scopes = vec![Scope::Read, Scope::Write, Scope::Admin];
210        let token = signer.sign("key-456", &scopes).expect("signing should succeed");
211        let claims = verifier.verify(&token).expect("verification should succeed");
212
213        assert_eq!(claims.scope, scopes);
214    }
215}
aa_auth/jwt.rs

aa_auth/
jwt.rs
