arete_server/websocket/
auth.rs

1use std::any::Any;
2use std::collections::{HashMap, HashSet};
3use std::net::SocketAddr;
4use std::sync::Arc;
5use std::time::Duration;
6
7use async_trait::async_trait;
8use tokio_tungstenite::tungstenite::http::Request;
9
10// Re-export AuthContext from arete-auth for convenience
11pub use arete_auth::AuthContext;
12// Re-export AuthErrorCode for convenience
13pub use arete_auth::AuthErrorCode;
14// Re-export RetryPolicy for convenience
15pub use arete_auth::RetryPolicy;
16// Re-export audit types
17pub use arete_auth::{
18    auth_failure_event, auth_success_event, rate_limit_event, AuditEvent, AuditSeverity,
19    ChannelAuditLogger, NoOpAuditLogger, SecurityAuditEvent, SecurityAuditLogger,
20};
21// Re-export metrics types
22pub use arete_auth::{AuthMetrics, AuthMetricsCollector, AuthMetricsSnapshot};
23// Re-export multi-key verifier types
24pub use arete_auth::{MultiKeyVerifier, MultiKeyVerifierBuilder, RotationKey};
25
26#[derive(Debug, Clone)]
27pub struct ConnectionAuthRequest {
28    pub remote_addr: SocketAddr,
29    pub path: String,
30    pub query: Option<String>,
31    pub headers: HashMap<String, String>,
32    /// Origin header from the request (for browser origin validation)
33    pub origin: Option<String>,
34}
35
36impl ConnectionAuthRequest {
37    pub fn from_http_request<B>(remote_addr: SocketAddr, request: &Request<B>) -> Self {
38        let mut headers = HashMap::new();
39        for (name, value) in request.headers() {
40            if let Ok(value_str) = value.to_str() {
41                headers.insert(name.as_str().to_ascii_lowercase(), value_str.to_string());
42            }
43        }
44
45        let origin = headers.get("origin").cloned();
46
47        Self {
48            remote_addr,
49            path: request.uri().path().to_string(),
50            query: request.uri().query().map(|q| q.to_string()),
51            headers,
52            origin,
53        }
54    }
55
56    pub fn header(&self, name: &str) -> Option<&str> {
57        self.headers
58            .get(&name.to_ascii_lowercase())
59            .map(String::as_str)
60    }
61
62    pub fn bearer_token(&self) -> Option<&str> {
63        let value = self.header("authorization")?;
64        let (scheme, token) = value.split_once(' ')?;
65        if scheme.eq_ignore_ascii_case("bearer") {
66            Some(token)
67        } else {
68            None
69        }
70    }
71
72    pub fn query_param(&self, key: &str) -> Option<&str> {
73        let query = self.query.as_deref()?;
74        query
75            .split('&')
76            .filter_map(|pair| pair.split_once('='))
77            .find_map(|(k, v)| if k == key { Some(v) } else { None })
78    }
79}
80
81/// Structured error details for machine-readable error handling
82#[derive(Debug, Clone, Default)]
83pub struct AuthErrorDetails {
84    /// The specific field or parameter that caused the error (if applicable)
85    pub field: Option<String>,
86    /// Additional context about the error
87    pub context: Option<String>,
88    /// Suggested action for the client to resolve the error
89    pub suggested_action: Option<String>,
90    /// Related documentation URL
91    pub docs_url: Option<String>,
92}
93
94/// Enhanced authentication denial with structured error information
95#[derive(Debug, Clone)]
96pub struct AuthDeny {
97    pub reason: String,
98    pub code: AuthErrorCode,
99    /// Structured error details for machine processing
100    pub details: AuthErrorDetails,
101    /// Retry policy hint
102    pub retry_policy: RetryPolicy,
103    /// HTTP status code equivalent for the error
104    pub http_status: u16,
105    /// When the error condition will reset (if applicable)
106    pub reset_at: Option<std::time::SystemTime>,
107}
108
109impl AuthDeny {
110    /// Create a new AuthDeny with the specified error code and reason
111    pub fn new(code: AuthErrorCode, reason: impl Into<String>) -> Self {
112        Self {
113            reason: reason.into(),
114            code,
115            details: AuthErrorDetails::default(),
116            retry_policy: code.default_retry_policy(),
117            http_status: code.http_status(),
118            reset_at: None,
119        }
120    }
121
122    /// Create an AuthDeny for missing token
123    pub fn token_missing() -> Self {
124        Self::new(
125            AuthErrorCode::TokenMissing,
126            "Missing session token (expected Authorization: Bearer <token> or query token)",
127        )
128        .with_suggested_action(
129            "Provide a valid session token in the Authorization header or as a query parameter",
130        )
131    }
132
133    /// Create an AuthDeny from a VerifyError
134    pub fn from_verify_error(err: arete_auth::VerifyError) -> Self {
135        let code = AuthErrorCode::from(&err);
136        Self::new(code, format!("Token verification failed: {}", err))
137    }
138
139    /// Add structured error details
140    pub fn with_details(mut self, details: AuthErrorDetails) -> Self {
141        self.details = details;
142        self
143    }
144
145    /// Add a specific field that caused the error
146    pub fn with_field(mut self, field: impl Into<String>) -> Self {
147        self.details.field = Some(field.into());
148        self
149    }
150
151    /// Add context to the error
152    pub fn with_context(mut self, context: impl Into<String>) -> Self {
153        self.details.context = Some(context.into());
154        self
155    }
156
157    /// Add a suggested action for the client
158    pub fn with_suggested_action(mut self, action: impl Into<String>) -> Self {
159        self.details.suggested_action = Some(action.into());
160        self
161    }
162
163    /// Add documentation URL
164    pub fn with_docs_url(mut self, url: impl Into<String>) -> Self {
165        self.details.docs_url = Some(url.into());
166        self
167    }
168
169    /// Set a custom retry policy
170    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
171        self.retry_policy = policy;
172        self
173    }
174
175    /// Set when the error condition will reset
176    pub fn with_reset_at(mut self, reset_at: std::time::SystemTime) -> Self {
177        self.reset_at = Some(reset_at);
178        self
179    }
180
181    /// Create an AuthDeny for rate limiting with retry information
182    pub fn rate_limited(retry_after: Duration, limit_type: &str) -> Self {
183        let reset_at = std::time::SystemTime::now() + retry_after;
184        Self::new(
185            AuthErrorCode::RateLimitExceeded,
186            format!(
187                "Rate limit exceeded for {}. Please retry after {:?}.",
188                limit_type, retry_after
189            ),
190        )
191        .with_retry_policy(RetryPolicy::RetryAfter(retry_after))
192        .with_reset_at(reset_at)
193        .with_suggested_action(format!(
194            "Wait {:?} before retrying the request",
195            retry_after
196        ))
197    }
198
199    /// Create an AuthDeny for connection limits
200    pub fn connection_limit_exceeded(limit_type: &str, current: usize, max: usize) -> Self {
201        Self::new(
202            AuthErrorCode::ConnectionLimitExceeded,
203            format!(
204                "Connection limit exceeded: {} has {} of {} allowed connections",
205                limit_type, current, max
206            ),
207        )
208        .with_suggested_action(
209            "Disconnect existing connections or wait for other connections to close",
210        )
211    }
212
213    /// The refusal a publishable-key session sees for a connection limit.
214    ///
215    /// Publishable keys ship to browsers, so the reason a visitor reads names
216    /// no identity and no counts. The detailed reason stays in
217    /// `details.context`, which is never sent to the client.
218    pub fn for_publishable_session(self) -> Self {
219        if self.code != AuthErrorCode::ConnectionLimitExceeded {
220            return self;
221        }
222        let detail = self.reason;
223        Self::new(
224            AuthErrorCode::ConnectionLimitExceeded,
225            "Connection limit reached",
226        )
227        .with_context(detail)
228        .with_suggested_action("Close other connections or retry later")
229    }
230
231    /// Convert to a JSON-serializable error response
232    pub fn to_error_response(&self) -> ErrorResponse {
233        ErrorResponse {
234            error: self.code.as_str().to_string(),
235            message: self.reason.clone(),
236            code: self.code.to_string(),
237            retryable: matches!(
238                self.retry_policy,
239                RetryPolicy::RetryImmediately
240                    | RetryPolicy::RetryAfter(_)
241                    | RetryPolicy::RetryWithBackoff { .. }
242                    | RetryPolicy::RetryWithFreshToken
243            ),
244            retry_after: match self.retry_policy {
245                RetryPolicy::RetryAfter(d) => Some(d.as_secs()),
246                _ => None,
247            },
248            suggested_action: self.details.suggested_action.clone(),
249            docs_url: self.details.docs_url.clone(),
250        }
251    }
252}
253
254/// JSON-serializable error response for clients
255#[derive(Debug, Clone, serde::Serialize)]
256pub struct ErrorResponse {
257    pub error: String,
258    pub message: String,
259    pub code: String,
260    pub retryable: bool,
261    #[serde(skip_serializing_if = "Option::is_none")]
262    pub retry_after: Option<u64>,
263    #[serde(skip_serializing_if = "Option::is_none")]
264    pub suggested_action: Option<String>,
265    #[serde(skip_serializing_if = "Option::is_none")]
266    pub docs_url: Option<String>,
267}
268
269/// Authentication decision with optional auth context
270#[derive(Debug, Clone)]
271#[allow(clippy::large_enum_variant)] // Keep the public plugin return type allocation-free.
272pub enum AuthDecision {
273    /// Connection is authorized with the given context
274    Allow(AuthContext),
275    /// Connection is denied
276    Deny(AuthDeny),
277}
278
279impl AuthDecision {
280    /// Check if the decision is Allow
281    pub fn is_allowed(&self) -> bool {
282        matches!(self, AuthDecision::Allow(_))
283    }
284
285    /// Get the auth context if allowed
286    pub fn auth_context(&self) -> Option<&AuthContext> {
287        match self {
288            AuthDecision::Allow(ctx) => Some(ctx),
289            AuthDecision::Deny(_) => None,
290        }
291    }
292}
293
294#[async_trait]
295pub trait WebSocketAuthPlugin: Send + Sync + Any {
296    async fn authorize(&self, request: &ConnectionAuthRequest) -> AuthDecision;
297
298    fn as_any(&self) -> &dyn Any;
299
300    /// Get the audit logger if configured
301    fn audit_logger(&self) -> Option<&dyn SecurityAuditLogger> {
302        None
303    }
304
305    /// Log a security audit event if audit logging is enabled
306    async fn log_audit(&self, event: SecurityAuditEvent) {
307        if let Some(logger) = self.audit_logger() {
308            logger.log(event).await;
309        }
310    }
311
312    /// Get auth metrics if configured
313    fn auth_metrics(&self) -> Option<&AuthMetrics> {
314        None
315    }
316}
317
318/// Development-only plugin that allows all connections
319///
320/// # Warning
321/// This should only be used for local development. Never use in production.
322pub struct AllowAllAuthPlugin;
323
324#[async_trait]
325impl WebSocketAuthPlugin for AllowAllAuthPlugin {
326    async fn authorize(&self, _request: &ConnectionAuthRequest) -> AuthDecision {
327        // Create a default auth context for development
328        let context = AuthContext {
329            subject: "anonymous".to_string(),
330            issuer: "allow-all".to_string(),
331            audience: "allow-all".to_string(),
332            key_class: arete_auth::KeyClass::Secret,
333            metering_key: "dev".to_string(),
334            deployment_id: None,
335            target_kind: None,
336            target_id: None,
337            program_id: None,
338            program_release_hash: None,
339            expires_at: u64::MAX, // Never expires
340            scope: "read write".to_string(),
341            limits: Default::default(),
342            plan: None,
343            origin: None,
344            client_ip: None,
345            jti: uuid::Uuid::new_v4().to_string(),
346            actor_key: None,
347            account_key: None,
348            consumer_key: None,
349            policy_version: None,
350            account_limits: Default::default(),
351            row_scopes: None,
352        };
353        AuthDecision::Allow(context)
354    }
355
356    fn as_any(&self) -> &dyn Any {
357        self
358    }
359}
360
361#[derive(Debug, Clone)]
362pub struct StaticTokenAuthPlugin {
363    tokens: HashSet<String>,
364    query_param_name: String,
365}
366
367impl StaticTokenAuthPlugin {
368    pub fn new(tokens: impl IntoIterator<Item = String>) -> Self {
369        Self {
370            tokens: tokens.into_iter().collect(),
371            query_param_name: "token".to_string(),
372        }
373    }
374
375    pub fn with_query_param_name(mut self, query_param_name: impl Into<String>) -> Self {
376        self.query_param_name = query_param_name.into();
377        self
378    }
379
380    fn extract_token<'a>(&self, request: &'a ConnectionAuthRequest) -> Option<&'a str> {
381        request
382            .bearer_token()
383            .or_else(|| request.query_param(&self.query_param_name))
384    }
385}
386
387#[async_trait]
388impl WebSocketAuthPlugin for StaticTokenAuthPlugin {
389    async fn authorize(&self, request: &ConnectionAuthRequest) -> AuthDecision {
390        let token = match self.extract_token(request) {
391            Some(token) => token,
392            None => {
393                return AuthDecision::Deny(AuthDeny::token_missing());
394            }
395        };
396
397        if self.tokens.contains(token) {
398            // Create auth context for static token
399            let context = AuthContext {
400                subject: format!("static:{}", &token[..token.len().min(8)]),
401                issuer: "static-token".to_string(),
402                audience: "static-token".to_string(),
403                key_class: arete_auth::KeyClass::Secret,
404                metering_key: token.to_string(),
405                deployment_id: None,
406                target_kind: None,
407                target_id: None,
408                program_id: None,
409                program_release_hash: None,
410                expires_at: u64::MAX, // Static tokens don't expire
411                scope: "read".to_string(),
412                limits: Default::default(),
413                plan: None,
414                origin: request.origin.clone(),
415                client_ip: None,
416                jti: uuid::Uuid::new_v4().to_string(),
417                actor_key: None,
418                account_key: None,
419                consumer_key: None,
420                policy_version: None,
421                account_limits: Default::default(),
422                row_scopes: None,
423            };
424            AuthDecision::Allow(context)
425        } else {
426            AuthDecision::Deny(AuthDeny::new(
427                AuthErrorCode::InvalidStaticToken,
428                "Invalid auth token",
429            ))
430        }
431    }
432
433    fn as_any(&self) -> &dyn Any {
434        self
435    }
436}
437
438/// Signed session token authentication plugin
439///
440/// This plugin verifies JWT session tokens using Ed25519 signatures.
441/// Tokens are expected to be passed either:
442/// - In the Authorization header: `Authorization: Bearer <token>`
443/// - As a query parameter: `?hs_token=<token>`
444enum SignedSessionVerifier {
445    Static(arete_auth::TokenVerifier),
446    CachedJwks(arete_auth::AsyncVerifier),
447    MultiKey(arete_auth::MultiKeyVerifier),
448}
449
450pub struct SignedSessionAuthPlugin {
451    verifier: SignedSessionVerifier,
452    query_param_name: String,
453    require_origin: bool,
454    audit_logger: Option<Arc<dyn SecurityAuditLogger>>,
455    metrics: Option<Arc<AuthMetrics>>,
456}
457
458impl SignedSessionAuthPlugin {
459    /// Create a new signed session auth plugin
460    pub fn new(verifier: arete_auth::TokenVerifier) -> Self {
461        Self {
462            verifier: SignedSessionVerifier::Static(verifier),
463            query_param_name: "hs_token".to_string(),
464            require_origin: false,
465            audit_logger: None,
466            metrics: None,
467        }
468    }
469
470    /// Create a signed session auth plugin backed by an async verifier, such as JWKS.
471    pub fn new_with_async_verifier(verifier: arete_auth::AsyncVerifier) -> Self {
472        Self {
473            verifier: SignedSessionVerifier::CachedJwks(verifier),
474            query_param_name: "hs_token".to_string(),
475            require_origin: false,
476            audit_logger: None,
477            metrics: None,
478        }
479    }
480
481    /// Create a signed session auth plugin backed by a multi-key verifier for key rotation.
482    pub fn new_with_multi_key_verifier(verifier: arete_auth::MultiKeyVerifier) -> Self {
483        Self {
484            verifier: SignedSessionVerifier::MultiKey(verifier),
485            query_param_name: "hs_token".to_string(),
486            require_origin: false,
487            audit_logger: None,
488            metrics: None,
489        }
490    }
491
492    /// Set a custom query parameter name for the token
493    pub fn with_query_param_name(mut self, name: impl Into<String>) -> Self {
494        self.query_param_name = name.into();
495        self
496    }
497
498    /// Require origin validation (defense-in-depth for browser clients)
499    pub fn with_origin_validation(mut self) -> Self {
500        self.require_origin = true;
501        self
502    }
503
504    /// Set an audit logger for security events
505    pub fn with_audit_logger(mut self, logger: Arc<dyn SecurityAuditLogger>) -> Self {
506        self.audit_logger = Some(logger);
507        self
508    }
509
510    /// Set metrics collector for auth operations
511    pub fn with_metrics(mut self, metrics: Arc<AuthMetrics>) -> Self {
512        self.metrics = Some(metrics);
513        self
514    }
515
516    /// Get metrics snapshot if metrics are enabled
517    pub fn metrics_snapshot(&self) -> Option<AuthMetricsSnapshot> {
518        self.metrics.as_ref().map(|m| m.snapshot())
519    }
520
521    fn extract_token<'a>(&self, request: &'a ConnectionAuthRequest) -> Option<&'a str> {
522        request
523            .bearer_token()
524            .or_else(|| request.query_param(&self.query_param_name))
525    }
526
527    /// Verify a token for in-band refresh and return the auth context
528    ///
529    /// This is used when a client wants to refresh their auth without reconnecting.
530    /// The origin is NOT validated here - we assume the client has already proven
531    /// origin at connection time, and we're just refreshing the session token.
532    #[allow(clippy::result_large_err)]
533    pub async fn verify_refresh_token(&self, token: &str) -> Result<AuthContext, AuthDeny> {
534        let result = match &self.verifier {
535            SignedSessionVerifier::Static(verifier) => verifier.verify(token, None, None),
536            SignedSessionVerifier::CachedJwks(verifier) => {
537                verifier.verify_with_cache(token, None, None).await
538            }
539            SignedSessionVerifier::MultiKey(verifier) => verifier.verify(token, None, None).await,
540        };
541
542        match result {
543            Ok(context) => Ok(context),
544            Err(e) => Err(AuthDeny::from_verify_error(e)),
545        }
546    }
547}
548
549#[async_trait]
550impl WebSocketAuthPlugin for SignedSessionAuthPlugin {
551    async fn authorize(&self, request: &ConnectionAuthRequest) -> AuthDecision {
552        let token = match self.extract_token(request) {
553            Some(token) => token,
554            None => {
555                return AuthDecision::Deny(AuthDeny::token_missing());
556            }
557        };
558
559        let expected_origin = request.origin.as_deref();
560
561        let expected_client_ip = None; // IP validation can be added here if needed
562
563        let result = match &self.verifier {
564            SignedSessionVerifier::Static(verifier) => {
565                verifier.verify(token, expected_origin, expected_client_ip)
566            }
567            SignedSessionVerifier::CachedJwks(verifier) => {
568                verifier
569                    .verify_with_cache(token, expected_origin, expected_client_ip)
570                    .await
571            }
572            SignedSessionVerifier::MultiKey(verifier) => {
573                verifier
574                    .verify(token, expected_origin, expected_client_ip)
575                    .await
576            }
577        };
578
579        match result {
580            Ok(context) => {
581                // Log successful authentication
582                let event = auth_success_event(&context.subject)
583                    .with_client_ip(request.remote_addr)
584                    .with_path(&request.path);
585                if let Some(origin) = &request.origin {
586                    let event = event.with_origin(origin.clone());
587                    self.log_audit(event).await;
588                } else {
589                    self.log_audit(event).await;
590                }
591                AuthDecision::Allow(context)
592            }
593            Err(e) => {
594                let deny = AuthDeny::from_verify_error(e);
595                // Log failed authentication
596                let event = auth_failure_event(&deny.code, &deny.reason)
597                    .with_client_ip(request.remote_addr)
598                    .with_path(&request.path);
599                let event = if let Some(origin) = &request.origin {
600                    event.with_origin(origin.clone())
601                } else {
602                    event
603                };
604                self.log_audit(event).await;
605                AuthDecision::Deny(deny)
606            }
607        }
608    }
609
610    fn as_any(&self) -> &dyn Any {
611        self
612    }
613
614    fn audit_logger(&self) -> Option<&dyn SecurityAuditLogger> {
615        self.audit_logger.as_ref().map(|l| l.as_ref())
616    }
617
618    fn auth_metrics(&self) -> Option<&AuthMetrics> {
619        self.metrics.as_ref().map(|m| m.as_ref())
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626
627    #[test]
628    fn extracts_bearer_and_query_tokens() {
629        let request = Request::builder()
630            .uri("/ws?token=query-token")
631            .header("Authorization", "Bearer header-token")
632            .body(())
633            .expect("request should build");
634
635        let auth_request = ConnectionAuthRequest::from_http_request(
636            "127.0.0.1:8877".parse().expect("socket addr should parse"),
637            &request,
638        );
639
640        assert_eq!(auth_request.bearer_token(), Some("header-token"));
641        assert_eq!(auth_request.query_param("token"), Some("query-token"));
642    }
643
644    #[tokio::test]
645    async fn static_token_plugin_allows_matching_token() {
646        let plugin = StaticTokenAuthPlugin::new(["secret".to_string()]);
647        let request = Request::builder()
648            .uri("/ws?token=secret")
649            .body(())
650            .expect("request should build");
651        let auth_request = ConnectionAuthRequest::from_http_request(
652            "127.0.0.1:8877".parse().expect("socket addr should parse"),
653            &request,
654        );
655
656        let decision = plugin.authorize(&auth_request).await;
657        assert!(decision.is_allowed());
658        assert!(decision.auth_context().is_some());
659    }
660
661    #[tokio::test]
662    async fn static_token_plugin_denies_missing_token() {
663        let plugin = StaticTokenAuthPlugin::new(["secret".to_string()]);
664        let request = Request::builder()
665            .uri("/ws")
666            .body(())
667            .expect("request should build");
668        let auth_request = ConnectionAuthRequest::from_http_request(
669            "127.0.0.1:8877".parse().expect("socket addr should parse"),
670            &request,
671        );
672
673        let decision = plugin.authorize(&auth_request).await;
674        assert!(!decision.is_allowed());
675    }
676
677    #[tokio::test]
678    async fn allow_all_plugin_allows_with_context() {
679        let plugin = AllowAllAuthPlugin;
680        let request = Request::builder()
681            .uri("/ws")
682            .body(())
683            .expect("request should build");
684        let auth_request = ConnectionAuthRequest::from_http_request(
685            "127.0.0.1:8877".parse().expect("socket addr should parse"),
686            &request,
687        );
688
689        let decision = plugin.authorize(&auth_request).await;
690        assert!(decision.is_allowed());
691        let ctx = decision.auth_context().unwrap();
692        assert_eq!(ctx.subject, "anonymous");
693    }
694
695    // Integration tests for handshake auth failures
696
697    #[tokio::test]
698    async fn signed_session_plugin_denies_missing_token() {
699        let signing_key = arete_auth::SigningKey::generate();
700        let verifying_key = signing_key.verifying_key();
701        let verifier =
702            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
703        let plugin = SignedSessionAuthPlugin::new(verifier);
704
705        let request = Request::builder()
706            .uri("/ws")
707            .body(())
708            .expect("request should build");
709        let auth_request = ConnectionAuthRequest::from_http_request(
710            "127.0.0.1:8877".parse().expect("socket addr should parse"),
711            &request,
712        );
713
714        let decision = plugin.authorize(&auth_request).await;
715        assert!(!decision.is_allowed());
716
717        if let AuthDecision::Deny(deny) = decision {
718            assert_eq!(deny.code, AuthErrorCode::TokenMissing);
719        } else {
720            panic!("Expected Deny decision");
721        }
722    }
723
724    #[tokio::test]
725    async fn signed_session_plugin_denies_expired_token() {
726        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
727        use std::time::{SystemTime, UNIX_EPOCH};
728
729        let signing_key = arete_auth::SigningKey::generate();
730        let verifying_key = signing_key.verifying_key();
731        let signer = TokenSigner::new(signing_key, "test-issuer");
732        let verifier =
733            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
734        let plugin = SignedSessionAuthPlugin::new(verifier);
735
736        // Create a token that expired 1 hour ago
737        let now = SystemTime::now()
738            .duration_since(UNIX_EPOCH)
739            .unwrap()
740            .as_secs();
741        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
742            .with_scope("read")
743            .with_key_class(KeyClass::Secret)
744            .build();
745
746        // Manually create expired claims
747        let mut expired_claims = claims;
748        expired_claims.exp = now - 3600; // Expired 1 hour ago
749        expired_claims.iat = now - 7200; // Issued 2 hours ago
750        expired_claims.nbf = now - 7200;
751
752        let token = signer.sign(expired_claims).unwrap();
753
754        let request = Request::builder()
755            .uri(format!("/ws?hs_token={}", token))
756            .body(())
757            .expect("request should build");
758        let auth_request = ConnectionAuthRequest::from_http_request(
759            "127.0.0.1:8877".parse().expect("socket addr should parse"),
760            &request,
761        );
762
763        let decision = plugin.authorize(&auth_request).await;
764        assert!(!decision.is_allowed());
765
766        if let AuthDecision::Deny(deny) = decision {
767            assert_eq!(deny.code, AuthErrorCode::TokenExpired);
768        } else {
769            panic!("Expected Deny decision for expired token");
770        }
771    }
772
773    #[tokio::test]
774    async fn signed_session_plugin_denies_invalid_signature() {
775        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
776
777        // Create two different key pairs
778        let signing_key = arete_auth::SigningKey::generate();
779        let wrong_key = arete_auth::SigningKey::generate();
780
781        // Sign with one key, verify with another
782        let signer = TokenSigner::new(signing_key, "test-issuer");
783        let wrong_verifying_key = wrong_key.verifying_key();
784        let verifier =
785            arete_auth::TokenVerifier::new(wrong_verifying_key, "test-issuer", "test-audience");
786        let plugin = SignedSessionAuthPlugin::new(verifier);
787
788        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
789            .with_scope("read")
790            .with_key_class(KeyClass::Secret)
791            .build();
792
793        let token = signer.sign(claims).unwrap();
794
795        let request = Request::builder()
796            .uri(format!("/ws?hs_token={}", token))
797            .body(())
798            .expect("request should build");
799        let auth_request = ConnectionAuthRequest::from_http_request(
800            "127.0.0.1:8877".parse().expect("socket addr should parse"),
801            &request,
802        );
803
804        let decision = plugin.authorize(&auth_request).await;
805        assert!(!decision.is_allowed());
806
807        if let AuthDecision::Deny(deny) = decision {
808            assert_eq!(deny.code, AuthErrorCode::TokenInvalidSignature);
809        } else {
810            panic!("Expected Deny decision for invalid signature");
811        }
812    }
813
814    #[tokio::test]
815    async fn signed_session_plugin_denies_wrong_audience() {
816        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
817
818        let signing_key = arete_auth::SigningKey::generate();
819        let verifying_key = signing_key.verifying_key();
820        let signer = TokenSigner::new(signing_key, "test-issuer");
821
822        // Verifier expects "test-audience", token is for "wrong-audience"
823        let verifier =
824            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
825        let plugin = SignedSessionAuthPlugin::new(verifier);
826
827        let claims = SessionClaims::builder("test-issuer", "test-subject", "wrong-audience")
828            .with_scope("read")
829            .with_key_class(KeyClass::Secret)
830            .build();
831
832        let token = signer.sign(claims).unwrap();
833
834        let request = Request::builder()
835            .uri(format!("/ws?hs_token={}", token))
836            .body(())
837            .expect("request should build");
838        let auth_request = ConnectionAuthRequest::from_http_request(
839            "127.0.0.1:8877".parse().expect("socket addr should parse"),
840            &request,
841        );
842
843        let decision = plugin.authorize(&auth_request).await;
844        assert!(!decision.is_allowed());
845
846        if let AuthDecision::Deny(deny) = decision {
847            assert_eq!(deny.code, AuthErrorCode::TokenInvalidAudience);
848        } else {
849            panic!("Expected Deny decision for wrong audience");
850        }
851    }
852
853    #[tokio::test]
854    async fn signed_session_plugin_denies_origin_mismatch() {
855        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
856
857        let signing_key = arete_auth::SigningKey::generate();
858        let verifying_key = signing_key.verifying_key();
859        let signer = TokenSigner::new(signing_key, "test-issuer");
860
861        // Verifier requires origin validation
862        let verifier =
863            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
864                .with_origin_validation();
865        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
866
867        // Token bound to specific origin
868        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
869            .with_scope("read")
870            .with_key_class(KeyClass::Secret)
871            .with_origin("https://allowed.example.com")
872            .build();
873
874        let token = signer.sign(claims).unwrap();
875
876        // Request from different origin
877        let request = Request::builder()
878            .uri(format!("/ws?hs_token={}", token))
879            .header("Origin", "https://evil.example.com")
880            .body(())
881            .expect("request should build");
882        let auth_request = ConnectionAuthRequest::from_http_request(
883            "127.0.0.1:8877".parse().expect("socket addr should parse"),
884            &request,
885        );
886
887        let decision = plugin.authorize(&auth_request).await;
888        assert!(!decision.is_allowed());
889
890        if let AuthDecision::Deny(deny) = decision {
891            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
892        } else {
893            panic!("Expected Deny decision for origin mismatch");
894        }
895    }
896
897    #[tokio::test]
898    async fn signed_session_plugin_allows_valid_token() {
899        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
900
901        let signing_key = arete_auth::SigningKey::generate();
902        let verifying_key = signing_key.verifying_key();
903        let signer = TokenSigner::new(signing_key, "test-issuer");
904        let verifier =
905            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
906        let plugin = SignedSessionAuthPlugin::new(verifier);
907
908        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
909            .with_scope("read")
910            .with_key_class(KeyClass::Secret)
911            .with_metering_key("meter-123")
912            .build();
913
914        let token = signer.sign(claims).unwrap();
915
916        let request = Request::builder()
917            .uri(format!("/ws?hs_token={}", token))
918            .body(())
919            .expect("request should build");
920        let auth_request = ConnectionAuthRequest::from_http_request(
921            "127.0.0.1:8877".parse().expect("socket addr should parse"),
922            &request,
923        );
924
925        let decision = plugin.authorize(&auth_request).await;
926        assert!(decision.is_allowed());
927
928        if let AuthDecision::Allow(ctx) = decision {
929            assert_eq!(ctx.subject, "test-subject");
930            assert_eq!(ctx.metering_key, "meter-123");
931            assert_eq!(ctx.key_class, KeyClass::Secret);
932        } else {
933            panic!("Expected Allow decision");
934        }
935    }
936
937    #[tokio::test]
938    async fn signed_session_plugin_allows_with_matching_origin() {
939        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
940
941        let signing_key = arete_auth::SigningKey::generate();
942        let verifying_key = signing_key.verifying_key();
943        let signer = TokenSigner::new(signing_key, "test-issuer");
944
945        let verifier =
946            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
947                .with_origin_validation();
948        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
949
950        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
951            .with_scope("read")
952            .with_key_class(KeyClass::Secret)
953            .with_origin("https://trusted.example.com")
954            .build();
955
956        let token = signer.sign(claims).unwrap();
957
958        let request = Request::builder()
959            .uri(format!("/ws?hs_token={}", token))
960            .header("Origin", "https://trusted.example.com")
961            .body(())
962            .expect("request should build");
963        let auth_request = ConnectionAuthRequest::from_http_request(
964            "127.0.0.1:8877".parse().expect("socket addr should parse"),
965            &request,
966        );
967
968        let decision = plugin.authorize(&auth_request).await;
969        assert!(decision.is_allowed());
970
971        if let AuthDecision::Allow(ctx) = decision {
972            assert_eq!(ctx.origin, Some("https://trusted.example.com".to_string()));
973        } else {
974            panic!("Expected Allow decision");
975        }
976    }
977
978    #[tokio::test]
979    async fn signed_session_plugin_allows_token_with_origin_when_no_origin_provided_and_not_required(
980    ) {
981        // This tests the non-browser client scenario (Rust, Python, etc.)
982        // where the client doesn't send an Origin header.
983        // The token has an origin claim from when it was minted via browser/API,
984        // but when the plugin doesn't require origin, the connection should still be allowed.
985        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
986
987        let signing_key = arete_auth::SigningKey::generate();
988        let verifying_key = signing_key.verifying_key();
989        let signer = TokenSigner::new(signing_key, "test-issuer");
990
991        // Plugin WITHOUT origin validation (default for public stacks)
992        let verifier =
993            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
994        let plugin = SignedSessionAuthPlugin::new(verifier);
995
996        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
997            .with_scope("read")
998            .with_key_class(KeyClass::Publishable)
999            .with_origin("https://example.com") // Token has origin claim
1000            .build();
1001
1002        let token = signer.sign(claims).unwrap();
1003
1004        // No Origin header provided (simulating non-browser client)
1005        let request = Request::builder()
1006            .uri(format!("/ws?hs_token={}", token))
1007            .body(())
1008            .expect("request should build");
1009        let auth_request = ConnectionAuthRequest::from_http_request(
1010            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1011            &request,
1012        );
1013
1014        // Should succeed even without Origin header
1015        let decision = plugin.authorize(&auth_request).await;
1016        assert!(
1017            decision.is_allowed(),
1018            "Expected Allow decision for non-browser client without Origin"
1019        );
1020
1021        if let AuthDecision::Allow(ctx) = decision {
1022            assert_eq!(ctx.origin, Some("https://example.com".to_string()));
1023        } else {
1024            panic!("Expected Allow decision");
1025        }
1026    }
1027
1028    #[tokio::test]
1029    async fn signed_session_plugin_validates_origin_when_provided_even_when_not_required() {
1030        // When origin IS provided, it should still be validated against the token
1031        // even when require_origin is false (defense-in-depth)
1032        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1033
1034        let signing_key = arete_auth::SigningKey::generate();
1035        let verifying_key = signing_key.verifying_key();
1036        let signer = TokenSigner::new(signing_key, "test-issuer");
1037
1038        // Plugin WITHOUT origin validation (default)
1039        let verifier =
1040            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
1041        let plugin = SignedSessionAuthPlugin::new(verifier);
1042
1043        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1044            .with_scope("read")
1045            .with_key_class(KeyClass::Publishable)
1046            .with_origin("https://allowed.example.com")
1047            .build();
1048
1049        let token = signer.sign(claims).unwrap();
1050
1051        // Origin provided and matches - should succeed
1052        let request = Request::builder()
1053            .uri(format!("/ws?hs_token={}", token))
1054            .header("Origin", "https://allowed.example.com")
1055            .body(())
1056            .expect("request should build");
1057        let auth_request = ConnectionAuthRequest::from_http_request(
1058            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1059            &request,
1060        );
1061
1062        let decision = plugin.authorize(&auth_request).await;
1063        assert!(decision.is_allowed());
1064
1065        // Origin provided but doesn't match - should fail
1066        let request = Request::builder()
1067            .uri(format!("/ws?hs_token={}", token))
1068            .header("Origin", "https://evil.example.com")
1069            .body(())
1070            .expect("request should build");
1071        let auth_request = ConnectionAuthRequest::from_http_request(
1072            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1073            &request,
1074        );
1075
1076        let decision = plugin.authorize(&auth_request).await;
1077        assert!(!decision.is_allowed());
1078
1079        if let AuthDecision::Deny(deny) = decision {
1080            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
1081        } else {
1082            panic!("Expected Deny decision for origin mismatch");
1083        }
1084    }
1085
1086    // Tests for AuthErrorCode utility methods
1087    #[test]
1088    fn auth_error_code_should_retry_logic() {
1089        assert!(AuthErrorCode::RateLimitExceeded.should_retry());
1090        assert!(AuthErrorCode::InternalError.should_retry());
1091        assert!(!AuthErrorCode::TokenExpired.should_retry());
1092        assert!(!AuthErrorCode::TokenInvalidSignature.should_retry());
1093        assert!(!AuthErrorCode::TokenMissing.should_retry());
1094    }
1095
1096    #[test]
1097    fn auth_error_code_should_refresh_token_logic() {
1098        assert!(AuthErrorCode::TokenExpired.should_refresh_token());
1099        assert!(AuthErrorCode::TokenInvalidSignature.should_refresh_token());
1100        assert!(AuthErrorCode::TokenInvalidFormat.should_refresh_token());
1101        assert!(AuthErrorCode::TokenInvalidIssuer.should_refresh_token());
1102        assert!(AuthErrorCode::TokenInvalidAudience.should_refresh_token());
1103        assert!(AuthErrorCode::TokenKeyNotFound.should_refresh_token());
1104        assert!(!AuthErrorCode::TokenMissing.should_refresh_token());
1105        assert!(!AuthErrorCode::RateLimitExceeded.should_refresh_token());
1106        assert!(!AuthErrorCode::ConnectionLimitExceeded.should_refresh_token());
1107    }
1108
1109    #[test]
1110    fn auth_error_code_string_representation() {
1111        assert_eq!(AuthErrorCode::TokenMissing.as_str(), "token-missing");
1112        assert_eq!(AuthErrorCode::TokenExpired.as_str(), "token-expired");
1113        assert_eq!(
1114            AuthErrorCode::TokenInvalidSignature.as_str(),
1115            "token-invalid-signature"
1116        );
1117        assert_eq!(
1118            AuthErrorCode::RateLimitExceeded.as_str(),
1119            "rate-limit-exceeded"
1120        );
1121        assert_eq!(
1122            AuthErrorCode::ConnectionLimitExceeded.as_str(),
1123            "connection-limit-exceeded"
1124        );
1125    }
1126
1127    // Tests for AuthDeny construction
1128    #[test]
1129    fn auth_deny_token_missing_factory() {
1130        let deny = AuthDeny::token_missing();
1131        assert_eq!(deny.code, AuthErrorCode::TokenMissing);
1132        assert!(deny.reason.contains("Missing session token"));
1133    }
1134
1135    #[test]
1136    fn auth_deny_from_verify_error_mapping() {
1137        use arete_auth::VerifyError;
1138
1139        let test_cases = vec![
1140            (VerifyError::Expired, AuthErrorCode::TokenExpired),
1141            (
1142                VerifyError::InvalidSignature,
1143                AuthErrorCode::TokenInvalidSignature,
1144            ),
1145            (
1146                VerifyError::InvalidIssuer,
1147                AuthErrorCode::TokenInvalidIssuer,
1148            ),
1149            (
1150                VerifyError::InvalidAudience,
1151                AuthErrorCode::TokenInvalidAudience,
1152            ),
1153            (
1154                VerifyError::KeyNotFound("kid123".to_string()),
1155                AuthErrorCode::TokenKeyNotFound,
1156            ),
1157            (
1158                VerifyError::OriginMismatch {
1159                    expected: "a".to_string(),
1160                    actual: "b".to_string(),
1161                },
1162                AuthErrorCode::OriginMismatch,
1163            ),
1164        ];
1165
1166        for (err, expected_code) in test_cases {
1167            let deny = AuthDeny::from_verify_error(err);
1168            assert_eq!(deny.code, expected_code);
1169        }
1170    }
1171
1172    // Tests for multiple auth failure scenarios in sequence
1173    #[tokio::test]
1174    async fn signed_session_plugin_handles_multiple_failure_reasons() {
1175        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1176
1177        let signing_key = arete_auth::SigningKey::generate();
1178        let verifying_key = signing_key.verifying_key();
1179        let signer = TokenSigner::new(signing_key, "test-issuer");
1180        let verifier =
1181            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
1182                .with_origin_validation();
1183        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
1184
1185        // Test 1: Missing token
1186        let request = Request::builder()
1187            .uri("/ws")
1188            .body(())
1189            .expect("request should build");
1190        let auth_request = ConnectionAuthRequest::from_http_request(
1191            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1192            &request,
1193        );
1194        let decision = plugin.authorize(&auth_request).await;
1195        assert!(!decision.is_allowed());
1196        match decision {
1197            AuthDecision::Deny(deny) => assert_eq!(deny.code, AuthErrorCode::TokenMissing),
1198            _ => panic!("Expected Deny decision"),
1199        }
1200
1201        // Test 2: Valid token with wrong origin
1202        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1203            .with_scope("read")
1204            .with_key_class(KeyClass::Secret)
1205            .with_origin("https://allowed.example.com")
1206            .build();
1207        let token = signer.sign(claims).unwrap();
1208
1209        let request = Request::builder()
1210            .uri(format!("/ws?hs_token={}", token))
1211            .header("Origin", "https://evil.example.com")
1212            .body(())
1213            .expect("request should build");
1214        let auth_request = ConnectionAuthRequest::from_http_request(
1215            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1216            &request,
1217        );
1218        let decision = plugin.authorize(&auth_request).await;
1219        assert!(!decision.is_allowed());
1220        match decision {
1221            AuthDecision::Deny(deny) => assert_eq!(deny.code, AuthErrorCode::OriginMismatch),
1222            _ => panic!("Expected Deny decision for origin mismatch"),
1223        }
1224
1225        // Test 3: Valid token with correct origin
1226        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1227            .with_scope("read")
1228            .with_key_class(KeyClass::Secret)
1229            .with_origin("https://allowed.example.com")
1230            .build();
1231        let token = signer.sign(claims).unwrap();
1232
1233        let request = Request::builder()
1234            .uri(format!("/ws?hs_token={}", token))
1235            .header("Origin", "https://allowed.example.com")
1236            .body(())
1237            .expect("request should build");
1238        let auth_request = ConnectionAuthRequest::from_http_request(
1239            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1240            &request,
1241        );
1242        let decision = plugin.authorize(&auth_request).await;
1243        assert!(decision.is_allowed());
1244    }
1245
1246    // Test for rate limit error code
1247    #[tokio::test]
1248    async fn auth_deney_with_rate_limit_code() {
1249        let deny = AuthDeny::new(
1250            AuthErrorCode::RateLimitExceeded,
1251            "Too many requests from this IP",
1252        );
1253        assert_eq!(deny.code, AuthErrorCode::RateLimitExceeded);
1254        assert!(deny.code.should_retry());
1255        assert!(!deny.code.should_refresh_token());
1256    }
1257
1258    // Test for connection limit error code
1259    #[tokio::test]
1260    async fn auth_deny_with_connection_limit_code() {
1261        let deny = AuthDeny::new(
1262            AuthErrorCode::ConnectionLimitExceeded,
1263            "Maximum connections exceeded for subject user-123",
1264        );
1265        assert_eq!(deny.code, AuthErrorCode::ConnectionLimitExceeded);
1266        assert!(!deny.code.should_retry());
1267        assert!(!deny.code.should_refresh_token());
1268    }
1269
1270    // Integration-style test: Token extraction from various sources
1271    #[test]
1272    fn token_extraction_priority() {
1273        // Header takes priority over query param
1274        let request = Request::builder()
1275            .uri("/ws?hs_token=query-value")
1276            .header("Authorization", "Bearer header-value")
1277            .body(())
1278            .expect("request should build");
1279        let auth_request = ConnectionAuthRequest::from_http_request(
1280            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1281            &request,
1282        );
1283
1284        // bearer_token should return header value
1285        assert_eq!(auth_request.bearer_token(), Some("header-value"));
1286        // query_param should return query value
1287        assert_eq!(auth_request.query_param("hs_token"), Some("query-value"));
1288    }
1289
1290    // Test malformed authorization header handling
1291    #[test]
1292    fn malformed_authorization_header() {
1293        let test_cases = vec![
1294            ("Basic dXNlcjpwYXNz", None),                // Wrong scheme
1295            ("Bearer", None),                            // Missing token (no space after Bearer)
1296            ("", None),                                  // Empty
1297            ("Bearer token extra", Some("token extra")), // Extra parts (token includes everything after scheme)
1298        ];
1299
1300        for (header_value, expected) in test_cases {
1301            let request = Request::builder()
1302                .uri("/ws")
1303                .header("Authorization", header_value)
1304                .body(())
1305                .expect("request should build");
1306            let auth_request = ConnectionAuthRequest::from_http_request(
1307                "127.0.0.1:8877".parse().expect("socket addr should parse"),
1308                &request,
1309            );
1310            assert_eq!(
1311                auth_request.bearer_token(),
1312                expected,
1313                "Failed for header: {}",
1314                header_value
1315            );
1316        }
1317    }
1318
1319    // ============================================
1320    // WEBSOCKET HANDSHAKE AUTH FAILURE TESTS
1321    // ============================================
1322    // These tests simulate real-world handshake failure scenarios
1323
1324    #[test]
1325    fn auth_deny_error_response_structure() {
1326        let deny = AuthDeny::new(AuthErrorCode::TokenExpired, "Token has expired")
1327            .with_field("exp")
1328            .with_context("Token expired 5 minutes ago")
1329            .with_suggested_action("Refresh your authentication token")
1330            .with_docs_url("https://docs.arete.run/auth/errors#token-expired");
1331
1332        let response = deny.to_error_response();
1333
1334        assert_eq!(response.code, "token-expired");
1335        assert_eq!(response.message, "Token has expired");
1336        assert_eq!(response.error, "token-expired");
1337        assert!(response.retryable);
1338        assert_eq!(
1339            response.suggested_action,
1340            Some("Refresh your authentication token".to_string())
1341        );
1342        assert_eq!(
1343            response.docs_url,
1344            Some("https://docs.arete.run/auth/errors#token-expired".to_string())
1345        );
1346    }
1347
1348    #[test]
1349    fn auth_deny_rate_limited_response() {
1350        use std::time::Duration;
1351
1352        let deny = AuthDeny::rate_limited(Duration::from_secs(30), "websocket connections");
1353        let response = deny.to_error_response();
1354
1355        assert_eq!(response.code, "rate-limit-exceeded");
1356        assert!(response.message.contains("30s"));
1357        assert!(response.retryable);
1358        assert_eq!(response.retry_after, Some(30));
1359    }
1360
1361    #[test]
1362    fn auth_deny_connection_limit_response() {
1363        let deny = AuthDeny::connection_limit_exceeded("user-123", 5, 5);
1364        let response = deny.to_error_response();
1365
1366        assert_eq!(response.code, "connection-limit-exceeded");
1367        assert!(response.message.contains("user-123"));
1368        assert!(response.message.contains("5 of 5"));
1369        assert!(response.retryable); // Connection limits are retryable (may become available)
1370    }
1371
1372    #[test]
1373    fn retry_policy_immediate() {
1374        let deny = AuthDeny::new(AuthErrorCode::InternalError, "Transient error")
1375            .with_retry_policy(RetryPolicy::RetryImmediately);
1376
1377        assert_eq!(deny.retry_policy, RetryPolicy::RetryImmediately);
1378    }
1379
1380    #[test]
1381    fn retry_policy_with_backoff() {
1382        use std::time::Duration;
1383
1384        let deny = AuthDeny::new(AuthErrorCode::RateLimitExceeded, "Too many requests")
1385            .with_retry_policy(RetryPolicy::RetryWithBackoff {
1386                initial: Duration::from_secs(1),
1387                max: Duration::from_secs(60),
1388            });
1389
1390        match deny.retry_policy {
1391            RetryPolicy::RetryWithBackoff { initial, max } => {
1392                assert_eq!(initial, Duration::from_secs(1));
1393                assert_eq!(max, Duration::from_secs(60));
1394            }
1395            _ => panic!("Expected RetryWithBackoff"),
1396        }
1397    }
1398
1399    #[test]
1400    fn auth_error_code_http_status_mapping() {
1401        assert_eq!(AuthErrorCode::TokenMissing.http_status(), 401);
1402        assert_eq!(AuthErrorCode::TokenExpired.http_status(), 401);
1403        assert_eq!(AuthErrorCode::TokenInvalidSignature.http_status(), 401);
1404        assert_eq!(AuthErrorCode::OriginMismatch.http_status(), 403);
1405        assert_eq!(AuthErrorCode::RateLimitExceeded.http_status(), 429);
1406        assert_eq!(AuthErrorCode::ConnectionLimitExceeded.http_status(), 429);
1407        assert_eq!(AuthErrorCode::InternalError.http_status(), 500);
1408    }
1409
1410    #[test]
1411    fn auth_error_code_default_retry_policies() {
1412        // Should refresh token
1413        assert!(matches!(
1414            AuthErrorCode::TokenExpired.default_retry_policy(),
1415            RetryPolicy::RetryWithFreshToken
1416        ));
1417        assert!(matches!(
1418            AuthErrorCode::TokenInvalidSignature.default_retry_policy(),
1419            RetryPolicy::RetryWithFreshToken
1420        ));
1421
1422        // Should retry with backoff
1423        assert!(matches!(
1424            AuthErrorCode::RateLimitExceeded.default_retry_policy(),
1425            RetryPolicy::RetryWithBackoff { .. }
1426        ));
1427        assert!(matches!(
1428            AuthErrorCode::InternalError.default_retry_policy(),
1429            RetryPolicy::RetryWithBackoff { .. }
1430        ));
1431
1432        // Should not retry
1433        assert!(matches!(
1434            AuthErrorCode::TokenMissing.default_retry_policy(),
1435            RetryPolicy::NoRetry
1436        ));
1437        assert!(matches!(
1438            AuthErrorCode::OriginMismatch.default_retry_policy(),
1439            RetryPolicy::NoRetry
1440        ));
1441    }
1442
1443    // Simulated handshake scenarios
1444
1445    #[tokio::test]
1446    async fn handshake_rejects_missing_token_with_proper_error() {
1447        // Create a request without a token
1448        let request = Request::builder()
1449            .uri("/ws")
1450            .body(())
1451            .expect("request should build");
1452
1453        let auth_request = ConnectionAuthRequest::from_http_request(
1454            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1455            &request,
1456        );
1457
1458        // For this test, we'll use a plugin that requires tokens
1459        // Actually AllowAllAuthPlugin doesn't require tokens, so let's create a static token plugin
1460        let static_plugin = StaticTokenAuthPlugin::new(["valid-token".to_string()]);
1461        let decision = static_plugin.authorize(&auth_request).await;
1462
1463        assert!(!decision.is_allowed());
1464
1465        if let AuthDecision::Deny(deny) = decision {
1466            assert_eq!(deny.code, AuthErrorCode::TokenMissing);
1467            assert_eq!(deny.http_status, 401);
1468            assert!(deny.reason.contains("Missing"));
1469        } else {
1470            panic!("Expected Deny decision");
1471        }
1472    }
1473
1474    #[tokio::test]
1475    async fn handshake_rejects_expired_token_with_retry_hint() {
1476        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1477        use std::time::{SystemTime, UNIX_EPOCH};
1478
1479        let signing_key = arete_auth::SigningKey::generate();
1480        let verifying_key = signing_key.verifying_key();
1481        let signer = TokenSigner::new(signing_key, "test-issuer");
1482
1483        // Create an expired token
1484        let now = SystemTime::now()
1485            .duration_since(UNIX_EPOCH)
1486            .unwrap()
1487            .as_secs();
1488        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1489            .with_scope("read")
1490            .with_key_class(KeyClass::Secret)
1491            .build();
1492
1493        let mut expired_claims = claims;
1494        expired_claims.exp = now - 3600;
1495        expired_claims.iat = now - 7200;
1496        expired_claims.nbf = now - 7200;
1497
1498        let token = signer.sign(expired_claims).unwrap();
1499
1500        // Create verifier and plugin
1501        let verifier =
1502            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
1503        let plugin = SignedSessionAuthPlugin::new(verifier);
1504
1505        let request = Request::builder()
1506            .uri(format!("/ws?hs_token={}", token))
1507            .body(())
1508            .expect("request should build");
1509
1510        let auth_request = ConnectionAuthRequest::from_http_request(
1511            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1512            &request,
1513        );
1514
1515        let decision = plugin.authorize(&auth_request).await;
1516
1517        assert!(!decision.is_allowed());
1518
1519        if let AuthDecision::Deny(deny) = decision {
1520            assert_eq!(deny.code, AuthErrorCode::TokenExpired);
1521            assert_eq!(deny.http_status, 401);
1522            // Should suggest refreshing the token
1523            assert!(matches!(
1524                deny.retry_policy,
1525                RetryPolicy::RetryWithFreshToken
1526            ));
1527        } else {
1528            panic!("Expected Deny decision");
1529        }
1530    }
1531
1532    #[tokio::test]
1533    async fn handshake_rejects_invalid_signature_with_retry_hint() {
1534        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1535
1536        // Create two different key pairs
1537        let signing_key = arete_auth::SigningKey::generate();
1538        let wrong_key = arete_auth::SigningKey::generate();
1539
1540        // Sign with one key, verify with another
1541        let signer = TokenSigner::new(signing_key, "test-issuer");
1542        let wrong_verifying_key = wrong_key.verifying_key();
1543        let verifier =
1544            arete_auth::TokenVerifier::new(wrong_verifying_key, "test-issuer", "test-audience");
1545        let plugin = SignedSessionAuthPlugin::new(verifier);
1546
1547        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1548            .with_scope("read")
1549            .with_key_class(KeyClass::Secret)
1550            .build();
1551
1552        let token = signer.sign(claims).unwrap();
1553
1554        let request = Request::builder()
1555            .uri(format!("/ws?hs_token={}", token))
1556            .body(())
1557            .expect("request should build");
1558
1559        let auth_request = ConnectionAuthRequest::from_http_request(
1560            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1561            &request,
1562        );
1563
1564        let decision = plugin.authorize(&auth_request).await;
1565
1566        assert!(!decision.is_allowed());
1567
1568        if let AuthDecision::Deny(deny) = decision {
1569            assert_eq!(deny.code, AuthErrorCode::TokenInvalidSignature);
1570            assert_eq!(deny.http_status, 401);
1571            // Should suggest refreshing the token
1572            assert!(matches!(
1573                deny.retry_policy,
1574                RetryPolicy::RetryWithFreshToken
1575            ));
1576        } else {
1577            panic!("Expected Deny decision");
1578        }
1579    }
1580
1581    #[tokio::test]
1582    async fn handshake_rejects_origin_mismatch_without_retry() {
1583        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1584
1585        let signing_key = arete_auth::SigningKey::generate();
1586        let verifying_key = signing_key.verifying_key();
1587        let signer = TokenSigner::new(signing_key, "test-issuer");
1588
1589        let verifier =
1590            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
1591                .with_origin_validation();
1592        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
1593
1594        // Token bound to specific origin
1595        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1596            .with_scope("read")
1597            .with_key_class(KeyClass::Secret)
1598            .with_origin("https://allowed.example.com")
1599            .build();
1600
1601        let token = signer.sign(claims).unwrap();
1602
1603        // Request from different origin
1604        let request = Request::builder()
1605            .uri(format!("/ws?hs_token={}", token))
1606            .header("Origin", "https://evil.example.com")
1607            .body(())
1608            .expect("request should build");
1609
1610        let auth_request = ConnectionAuthRequest::from_http_request(
1611            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1612            &request,
1613        );
1614
1615        let decision = plugin.authorize(&auth_request).await;
1616
1617        assert!(!decision.is_allowed());
1618
1619        if let AuthDecision::Deny(deny) = decision {
1620            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
1621            assert_eq!(deny.http_status, 403);
1622            // Should NOT suggest retrying - this is a security issue
1623            assert!(matches!(deny.retry_policy, RetryPolicy::NoRetry));
1624        } else {
1625            panic!("Expected Deny decision");
1626        }
1627    }
1628
1629    // Test that AuthDeny can be converted to HTTP error response
1630    #[test]
1631    fn auth_deny_to_http_response() {
1632        let deny = AuthDeny::new(AuthErrorCode::RateLimitExceeded, "Too many requests")
1633            .with_suggested_action("Wait before retrying")
1634            .with_retry_policy(RetryPolicy::RetryAfter(Duration::from_secs(30)));
1635
1636        let response = deny.to_error_response();
1637
1638        // Verify the response is serializable
1639        let json = serde_json::to_string(&response).expect("Should serialize");
1640        assert!(json.contains("rate-limit-exceeded"));
1641        assert!(json.contains("Too many requests"));
1642        assert!(json.contains("Wait before retrying"));
1643        assert!(json.contains("\"retryable\":true"));
1644        assert!(json.contains("\"retry_after\":30"));
1645    }
1646
1647    // Test comprehensive error scenarios
1648    #[tokio::test]
1649    async fn comprehensive_auth_error_scenarios() {
1650        let signing_key = arete_auth::SigningKey::generate();
1651        let verifying_key = signing_key.verifying_key();
1652        let verifier =
1653            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
1654        let plugin = SignedSessionAuthPlugin::new(verifier);
1655
1656        let test_cases = vec![
1657            ("missing_token", None, AuthErrorCode::TokenMissing),
1658            (
1659                "invalid_format",
1660                Some("not-a-valid-token"),
1661                AuthErrorCode::TokenInvalidFormat,
1662            ),
1663        ];
1664
1665        for (name, token, expected_code) in test_cases {
1666            let uri = token.map_or_else(|| "/ws".to_string(), |t| format!("/ws?hs_token={}", t));
1667
1668            let request = Request::builder()
1669                .uri(&uri)
1670                .body(())
1671                .expect("request should build");
1672
1673            let auth_request = ConnectionAuthRequest::from_http_request(
1674                "127.0.0.1:8877".parse().expect("socket addr should parse"),
1675                &request,
1676            );
1677
1678            let decision = plugin.authorize(&auth_request).await;
1679
1680            assert!(!decision.is_allowed(), "{}: should deny", name);
1681
1682            if let AuthDecision::Deny(deny) = decision {
1683                assert_eq!(deny.code, expected_code, "{}: wrong error code", name);
1684            } else {
1685                panic!("{}: Expected Deny decision", name);
1686            }
1687        }
1688    }
1689}
arete_server/websocket/auth.rs

arete_server/websocket/
auth.rs
