dk_protocol/
auth.rs

1//! JWT and shared-secret authentication for the Agent Protocol.
2//!
3//! [`AuthConfig`] supports four modes:
4//! - **Jwt** -- HMAC-SHA256 JWT validation/issuance.
5//! - **SharedSecret** -- simple string comparison (legacy).
6//! - **Dual** -- try JWT first, fall back to shared secret.
7//! - **External** -- trust an outer layer (e.g. tonic interceptor) that already
8//!   validated the token; pass through the token value as agent id.
9
10use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
11use serde::{Deserialize, Serialize};
12use tonic::Status;
13
14// ── Claims ──────────────────────────────────────────────────────────
15
16/// JWT claims carried inside every agent token.
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct DkodClaims {
19    /// Subject -- the agent identity (e.g. "agent-42").
20    pub sub: String,
21    /// Issuer -- always "dkod".
22    pub iss: String,
23    /// Expiration (UTC epoch seconds).
24    pub exp: usize,
25    /// Issued-at (UTC epoch seconds).
26    pub iat: usize,
27    /// Permission scope (e.g. "read", "read+write", "admin").
28    pub scope: String,
29}
30
31// ── AuthConfig ──────────────────────────────────────────────────────
32
33/// Authentication configuration supporting JWT, shared-secret, or both.
34#[derive(Clone, Debug)]
35pub enum AuthConfig {
36    /// Pure JWT mode -- validate/issue using HMAC-SHA256.
37    Jwt { secret: String },
38    /// Legacy shared-secret mode -- simple string comparison.
39    SharedSecret { token: String },
40    /// Dual mode -- try JWT first, fall back to shared secret.
41    Dual {
42        jwt_secret: String,
43        shared_token: String,
44    },
45    /// External mode -- an outer authentication layer (e.g. a tonic
46    /// interceptor) has already validated the token. The raw token value
47    /// is passed through as the agent id. Useful for managed platforms
48    /// where a gateway handles JWT verification before the request
49    /// reaches the protocol server.
50    External,
51}
52
53impl AuthConfig {
54    /// Validate an incoming bearer token.
55    ///
56    /// Returns the agent id on success:
57    /// - JWT modes: the `sub` claim from the decoded token.
58    /// - SharedSecret mode: the literal `"anonymous"`.
59    ///
60    /// Empty tokens are always rejected regardless of auth mode.
61    pub fn validate(&self, token: &str) -> Result<String, Status> {
62        if token.is_empty() {
63            return Err(Status::unauthenticated("Auth token must not be empty"));
64        }
65
66        match self {
67            AuthConfig::Jwt { secret } => validate_jwt(token, secret),
68
69            AuthConfig::SharedSecret {
70                token: expected_token,
71            } => {
72                if token == expected_token {
73                    Ok("anonymous".to_string())
74                } else {
75                    Err(Status::unauthenticated("Invalid auth token"))
76                }
77            }
78
79            AuthConfig::Dual {
80                jwt_secret,
81                shared_token,
82            } => {
83                // Try JWT first; if that fails, try shared-secret.
84                match validate_jwt(token, jwt_secret) {
85                    Ok(agent_id) => Ok(agent_id),
86                    Err(_jwt_err) => {
87                        if token == shared_token {
88                            Ok("anonymous".to_string())
89                        } else {
90                            Err(Status::unauthenticated("Invalid auth token"))
91                        }
92                    }
93                }
94            }
95
96            AuthConfig::External => {
97                // The outer layer already validated the token; pass it
98                // through as the agent id (typically a user_id).
99                Ok(token.to_string())
100            }
101        }
102    }
103
104    /// Fully validate a bearer token and return the complete `DkodClaims`.
105    ///
106    /// Unlike [`validate`], which returns only the agent id, this method
107    /// returns the entire validated claims struct so callers can inspect
108    /// fields such as `scope` without an additional insecure decode pass.
109    ///
110    /// Returns `None` for non-JWT auth modes (SharedSecret / External) where
111    /// no claims struct exists.
112    pub fn validate_claims(&self, token: &str) -> Option<DkodClaims> {
113        match self {
114            AuthConfig::Jwt { secret } => validate_jwt_full(token, secret).ok(),
115            AuthConfig::Dual { jwt_secret, .. } => validate_jwt_full(token, jwt_secret).ok(),
116            AuthConfig::SharedSecret { .. } | AuthConfig::External => None,
117        }
118    }
119
120    /// Issue a new JWT for the given agent.
121    ///
122    /// Only available when a JWT secret is configured (Jwt or Dual mode).
123    /// Returns `Status::failed_precondition` if called in SharedSecret-only
124    /// mode.
125    pub fn issue_token(
126        &self,
127        agent_id: &str,
128        scope: &str,
129        ttl_secs: usize,
130    ) -> Result<String, Status> {
131        let secret = match self {
132            AuthConfig::Jwt { secret } => secret,
133            AuthConfig::Dual { jwt_secret, .. } => jwt_secret,
134            AuthConfig::SharedSecret { .. } | AuthConfig::External => {
135                return Err(Status::failed_precondition(
136                    "Cannot issue JWT tokens without a JWT secret",
137                ));
138            }
139        };
140
141        if secret.len() < 32 {
142            tracing::error!("JWT secret is too short (< 32 bytes); check server configuration");
143            return Err(Status::internal("server misconfiguration"));
144        }
145
146        let now = jsonwebtoken::get_current_timestamp() as usize;
147        let claims = DkodClaims {
148            sub: agent_id.to_string(),
149            iss: "dkod".to_string(),
150            exp: now + ttl_secs,
151            iat: now,
152            scope: scope.to_string(),
153        };
154
155        encode(
156            &Header::default(), // HS256
157            &claims,
158            &EncodingKey::from_secret(secret.as_bytes()),
159        )
160        .map_err(|e| Status::internal(format!("Failed to encode JWT: {e}")))
161    }
162}
163
164// ── Public helpers ──────────────────────────────────────────────────
165
166/// Returns `true` when `claims.scope` equals `"admin"` or contains the
167/// word `"admin"` (space-separated scopes).
168///
169/// Callers **must** pass claims that have already been fully verified
170/// (signature + expiry) by [`AuthConfig::validate`].  Never call this
171/// with claims obtained via an insecure decode path.
172pub fn claims_have_admin_scope(claims: &DkodClaims) -> bool {
173    let scope = &claims.scope;
174    scope == "admin" || scope.split_whitespace().any(|s| s == "admin")
175}
176
177// ── Private helpers ─────────────────────────────────────────────────
178
179/// Decode and validate a JWT using HMAC-SHA256.
180///
181/// Validates:
182/// - Algorithm: HS256
183/// - Issuer: "dkod"
184/// - Required claims: sub, exp, iss
185///
186/// Returns the full `DkodClaims` on success.
187fn validate_jwt_full(token: &str, secret: &str) -> Result<DkodClaims, Status> {
188    if secret.len() < 32 {
189        tracing::error!("JWT secret is too short (< 32 bytes); check server configuration");
190        return Err(Status::unauthenticated("JWT validation failed"));
191    }
192
193    let mut validation = Validation::new(jsonwebtoken::Algorithm::HS256);
194    validation.set_issuer(&["dkod"]);
195    validation.set_required_spec_claims(&["sub", "exp", "iss"]);
196
197    let token_data = decode::<DkodClaims>(
198        token,
199        &DecodingKey::from_secret(secret.as_bytes()),
200        &validation,
201    )
202    .map_err(|e| Status::unauthenticated(format!("JWT validation failed: {e}")))?;
203
204    Ok(token_data.claims)
205}
206
207/// Decode and validate a JWT, returning just the `sub` claim (agent id).
208fn validate_jwt(token: &str, secret: &str) -> Result<String, Status> {
209    validate_jwt_full(token, secret).map(|c| c.sub)
210}
211
212// ── Tests ───────────────────────────────────────────────────────────
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    const TEST_SECRET: &str = "test-secret-key-for-unit-tests!!";
219    const TEST_AGENT: &str = "agent-42";
220    const TEST_SCOPE: &str = "read+write";
221    const TTL: usize = 3600; // 1 hour
222
223    #[test]
224    fn jwt_roundtrip() {
225        let config = AuthConfig::Jwt {
226            secret: TEST_SECRET.to_string(),
227        };
228        let token = config
229            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
230            .expect("issue_token should succeed");
231        let agent_id = config.validate(&token).expect("validate should succeed");
232        assert_eq!(agent_id, TEST_AGENT);
233    }
234
235    #[test]
236    fn jwt_rejects_bad_token() {
237        let config = AuthConfig::Jwt {
238            secret: TEST_SECRET.to_string(),
239        };
240        let result = config.validate("not-a-jwt");
241        assert!(result.is_err(), "should reject garbage token");
242    }
243
244    #[test]
245    fn jwt_rejects_wrong_secret() {
246        let config1 = AuthConfig::Jwt {
247            secret: "secret-one-padding-for-32-bytes!".to_string(),
248        };
249        let config2 = AuthConfig::Jwt {
250            secret: "secret-two-padding-for-32-bytes!".to_string(),
251        };
252        let token = config1
253            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
254            .expect("issue_token should succeed");
255        let result = config2.validate(&token);
256        assert!(result.is_err(), "should reject token signed with different secret");
257    }
258
259    #[test]
260    fn shared_secret_accepts_correct_token() {
261        let config = AuthConfig::SharedSecret {
262            token: "my-shared-token".to_string(),
263        };
264        let agent_id = config
265            .validate("my-shared-token")
266            .expect("should accept correct token");
267        assert_eq!(agent_id, "anonymous");
268    }
269
270    #[test]
271    fn shared_secret_rejects_wrong_token() {
272        let config = AuthConfig::SharedSecret {
273            token: "correct-token".to_string(),
274        };
275        let result = config.validate("wrong-token");
276        assert!(result.is_err(), "should reject wrong token");
277    }
278
279    #[test]
280    fn dual_mode_accepts_jwt() {
281        let config = AuthConfig::Dual {
282            jwt_secret: TEST_SECRET.to_string(),
283            shared_token: "fallback-token".to_string(),
284        };
285        let token = config
286            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
287            .expect("issue_token should succeed");
288        let agent_id = config.validate(&token).expect("should accept valid JWT");
289        assert_eq!(agent_id, TEST_AGENT);
290    }
291
292    #[test]
293    fn dual_mode_falls_back_to_shared_secret() {
294        let config = AuthConfig::Dual {
295            jwt_secret: TEST_SECRET.to_string(),
296            shared_token: "fallback-token".to_string(),
297        };
298        let agent_id = config
299            .validate("fallback-token")
300            .expect("should fall back to shared secret");
301        assert_eq!(agent_id, "anonymous");
302    }
303
304    #[test]
305    fn dual_mode_rejects_invalid() {
306        let config = AuthConfig::Dual {
307            jwt_secret: TEST_SECRET.to_string(),
308            shared_token: "fallback-token".to_string(),
309        };
310        let result = config.validate("garbage-that-matches-nothing");
311        assert!(result.is_err(), "should reject invalid token in dual mode");
312    }
313
314    #[test]
315    fn empty_token_rejected_in_all_modes() {
316        let jwt = AuthConfig::Jwt {
317            secret: TEST_SECRET.to_string(),
318        };
319        assert!(jwt.validate("").is_err(), "JWT mode should reject empty token");
320
321        let shared = AuthConfig::SharedSecret {
322            token: "my-token".to_string(),
323        };
324        assert!(shared.validate("").is_err(), "SharedSecret should reject empty token");
325
326        let dual = AuthConfig::Dual {
327            jwt_secret: TEST_SECRET.to_string(),
328            shared_token: "fallback".to_string(),
329        };
330        assert!(dual.validate("").is_err(), "Dual mode should reject empty token");
331
332        let external = AuthConfig::External;
333        assert!(external.validate("").is_err(), "External mode should reject empty token");
334    }
335
336    #[test]
337    fn empty_shared_secret_never_matches() {
338        // Even if someone constructs SharedSecret with an empty token,
339        // empty incoming tokens are rejected before comparison.
340        let config = AuthConfig::SharedSecret {
341            token: "".to_string(),
342        };
343        assert!(config.validate("").is_err(), "empty token should be rejected even if shared secret is empty");
344    }
345
346    #[test]
347    fn external_passes_through_token_as_agent_id() {
348        let config = AuthConfig::External;
349        let agent_id = config
350            .validate("user_abc123")
351            .expect("External should accept any non-empty token");
352        assert_eq!(agent_id, "user_abc123");
353    }
354
355    #[test]
356    fn external_rejects_empty_token() {
357        let config = AuthConfig::External;
358        assert!(
359            config.validate("").is_err(),
360            "External should reject empty token"
361        );
362    }
363
364    #[test]
365    fn external_cannot_issue_tokens() {
366        let config = AuthConfig::External;
367        assert!(
368            config.issue_token("agent", "read", 3600).is_err(),
369            "External mode should not issue JWT tokens"
370        );
371    }
372
373    #[test]
374    fn rejects_short_jwt_secret() {
375        let config = AuthConfig::Jwt {
376            secret: "short".to_string(),
377        };
378        let result = config.issue_token("agent-1", "full", 3600);
379        assert!(result.is_err());
380    }
381}
dk_protocol/auth.rs

dk_protocol/
auth.rs
