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    /// Issue a new JWT for the given agent.
105    ///
106    /// Only available when a JWT secret is configured (Jwt or Dual mode).
107    /// Returns `Status::failed_precondition` if called in SharedSecret-only
108    /// mode.
109    pub fn issue_token(
110        &self,
111        agent_id: &str,
112        scope: &str,
113        ttl_secs: usize,
114    ) -> Result<String, Status> {
115        let secret = match self {
116            AuthConfig::Jwt { secret } => secret,
117            AuthConfig::Dual { jwt_secret, .. } => jwt_secret,
118            AuthConfig::SharedSecret { .. } | AuthConfig::External => {
119                return Err(Status::failed_precondition(
120                    "Cannot issue JWT tokens without a JWT secret",
121                ));
122            }
123        };
124
125        if secret.len() < 32 {
126            tracing::error!("JWT secret is too short (< 32 bytes); check server configuration");
127            return Err(Status::internal("server misconfiguration"));
128        }
129
130        let now = jsonwebtoken::get_current_timestamp() as usize;
131        let claims = DkodClaims {
132            sub: agent_id.to_string(),
133            iss: "dkod".to_string(),
134            exp: now + ttl_secs,
135            iat: now,
136            scope: scope.to_string(),
137        };
138
139        encode(
140            &Header::default(), // HS256
141            &claims,
142            &EncodingKey::from_secret(secret.as_bytes()),
143        )
144        .map_err(|e| Status::internal(format!("Failed to encode JWT: {e}")))
145    }
146}
147
148// ── Private helpers ─────────────────────────────────────────────────
149
150/// Decode and validate a JWT using HMAC-SHA256.
151///
152/// Validates:
153/// - Algorithm: HS256
154/// - Issuer: "dkod"
155/// - Required claims: sub, exp, iss
156///
157/// Returns the `sub` claim (agent id) on success.
158fn validate_jwt(token: &str, secret: &str) -> Result<String, Status> {
159    if secret.len() < 32 {
160        tracing::error!("JWT secret is too short (< 32 bytes); check server configuration");
161        return Err(Status::unauthenticated("JWT validation failed"));
162    }
163
164    let mut validation = Validation::new(jsonwebtoken::Algorithm::HS256);
165    validation.set_issuer(&["dkod"]);
166    validation.set_required_spec_claims(&["sub", "exp", "iss"]);
167
168    let token_data = decode::<DkodClaims>(
169        token,
170        &DecodingKey::from_secret(secret.as_bytes()),
171        &validation,
172    )
173    .map_err(|e| Status::unauthenticated(format!("JWT validation failed: {e}")))?;
174
175    Ok(token_data.claims.sub)
176}
177
178// ── Tests ───────────────────────────────────────────────────────────
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    const TEST_SECRET: &str = "test-secret-key-for-unit-tests!!";
185    const TEST_AGENT: &str = "agent-42";
186    const TEST_SCOPE: &str = "read+write";
187    const TTL: usize = 3600; // 1 hour
188
189    #[test]
190    fn jwt_roundtrip() {
191        let config = AuthConfig::Jwt {
192            secret: TEST_SECRET.to_string(),
193        };
194        let token = config
195            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
196            .expect("issue_token should succeed");
197        let agent_id = config.validate(&token).expect("validate should succeed");
198        assert_eq!(agent_id, TEST_AGENT);
199    }
200
201    #[test]
202    fn jwt_rejects_bad_token() {
203        let config = AuthConfig::Jwt {
204            secret: TEST_SECRET.to_string(),
205        };
206        let result = config.validate("not-a-jwt");
207        assert!(result.is_err(), "should reject garbage token");
208    }
209
210    #[test]
211    fn jwt_rejects_wrong_secret() {
212        let config1 = AuthConfig::Jwt {
213            secret: "secret-one-padding-for-32-bytes!".to_string(),
214        };
215        let config2 = AuthConfig::Jwt {
216            secret: "secret-two-padding-for-32-bytes!".to_string(),
217        };
218        let token = config1
219            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
220            .expect("issue_token should succeed");
221        let result = config2.validate(&token);
222        assert!(result.is_err(), "should reject token signed with different secret");
223    }
224
225    #[test]
226    fn shared_secret_accepts_correct_token() {
227        let config = AuthConfig::SharedSecret {
228            token: "my-shared-token".to_string(),
229        };
230        let agent_id = config
231            .validate("my-shared-token")
232            .expect("should accept correct token");
233        assert_eq!(agent_id, "anonymous");
234    }
235
236    #[test]
237    fn shared_secret_rejects_wrong_token() {
238        let config = AuthConfig::SharedSecret {
239            token: "correct-token".to_string(),
240        };
241        let result = config.validate("wrong-token");
242        assert!(result.is_err(), "should reject wrong token");
243    }
244
245    #[test]
246    fn dual_mode_accepts_jwt() {
247        let config = AuthConfig::Dual {
248            jwt_secret: TEST_SECRET.to_string(),
249            shared_token: "fallback-token".to_string(),
250        };
251        let token = config
252            .issue_token(TEST_AGENT, TEST_SCOPE, TTL)
253            .expect("issue_token should succeed");
254        let agent_id = config.validate(&token).expect("should accept valid JWT");
255        assert_eq!(agent_id, TEST_AGENT);
256    }
257
258    #[test]
259    fn dual_mode_falls_back_to_shared_secret() {
260        let config = AuthConfig::Dual {
261            jwt_secret: TEST_SECRET.to_string(),
262            shared_token: "fallback-token".to_string(),
263        };
264        let agent_id = config
265            .validate("fallback-token")
266            .expect("should fall back to shared secret");
267        assert_eq!(agent_id, "anonymous");
268    }
269
270    #[test]
271    fn dual_mode_rejects_invalid() {
272        let config = AuthConfig::Dual {
273            jwt_secret: TEST_SECRET.to_string(),
274            shared_token: "fallback-token".to_string(),
275        };
276        let result = config.validate("garbage-that-matches-nothing");
277        assert!(result.is_err(), "should reject invalid token in dual mode");
278    }
279
280    #[test]
281    fn empty_token_rejected_in_all_modes() {
282        let jwt = AuthConfig::Jwt {
283            secret: TEST_SECRET.to_string(),
284        };
285        assert!(jwt.validate("").is_err(), "JWT mode should reject empty token");
286
287        let shared = AuthConfig::SharedSecret {
288            token: "my-token".to_string(),
289        };
290        assert!(shared.validate("").is_err(), "SharedSecret should reject empty token");
291
292        let dual = AuthConfig::Dual {
293            jwt_secret: TEST_SECRET.to_string(),
294            shared_token: "fallback".to_string(),
295        };
296        assert!(dual.validate("").is_err(), "Dual mode should reject empty token");
297
298        let external = AuthConfig::External;
299        assert!(external.validate("").is_err(), "External mode should reject empty token");
300    }
301
302    #[test]
303    fn empty_shared_secret_never_matches() {
304        // Even if someone constructs SharedSecret with an empty token,
305        // empty incoming tokens are rejected before comparison.
306        let config = AuthConfig::SharedSecret {
307            token: "".to_string(),
308        };
309        assert!(config.validate("").is_err(), "empty token should be rejected even if shared secret is empty");
310    }
311
312    #[test]
313    fn external_passes_through_token_as_agent_id() {
314        let config = AuthConfig::External;
315        let agent_id = config
316            .validate("user_abc123")
317            .expect("External should accept any non-empty token");
318        assert_eq!(agent_id, "user_abc123");
319    }
320
321    #[test]
322    fn external_rejects_empty_token() {
323        let config = AuthConfig::External;
324        assert!(
325            config.validate("").is_err(),
326            "External should reject empty token"
327        );
328    }
329
330    #[test]
331    fn external_cannot_issue_tokens() {
332        let config = AuthConfig::External;
333        assert!(
334            config.issue_token("agent", "read", 3600).is_err(),
335            "External mode should not issue JWT tokens"
336        );
337    }
338
339    #[test]
340    fn rejects_short_jwt_secret() {
341        let config = AuthConfig::Jwt {
342            secret: "short".to_string(),
343        };
344        let result = config.issue_token("agent-1", "full", 3600);
345        assert!(result.is_err());
346    }
347}
dk_protocol/auth.rs

dk_protocol/
auth.rs
