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

arete_server/websocket/
auth.rs
