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

arete_server/websocket/
auth.rs
