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        use arete_auth::TokenSigner;
658
659        let signing_key = arete_auth::SigningKey::generate();
660        let verifying_key = signing_key.verifying_key();
661        let verifier =
662            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
663        let plugin = SignedSessionAuthPlugin::new(verifier);
664
665        let request = Request::builder()
666            .uri("/ws")
667            .body(())
668            .expect("request should build");
669        let auth_request = ConnectionAuthRequest::from_http_request(
670            "127.0.0.1:8877".parse().expect("socket addr should parse"),
671            &request,
672        );
673
674        let decision = plugin.authorize(&auth_request).await;
675        assert!(!decision.is_allowed());
676
677        if let AuthDecision::Deny(deny) = decision {
678            assert_eq!(deny.code, AuthErrorCode::TokenMissing);
679        } else {
680            panic!("Expected Deny decision");
681        }
682    }
683
684    #[tokio::test]
685    async fn signed_session_plugin_denies_expired_token() {
686        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
687        use std::time::{SystemTime, UNIX_EPOCH};
688
689        let signing_key = arete_auth::SigningKey::generate();
690        let verifying_key = signing_key.verifying_key();
691        let signer = TokenSigner::new(signing_key, "test-issuer");
692        let verifier =
693            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
694        let plugin = SignedSessionAuthPlugin::new(verifier);
695
696        // Create a token that expired 1 hour ago
697        let now = SystemTime::now()
698            .duration_since(UNIX_EPOCH)
699            .unwrap()
700            .as_secs();
701        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
702            .with_scope("read")
703            .with_key_class(KeyClass::Secret)
704            .build();
705
706        // Manually create expired claims
707        let mut expired_claims = claims;
708        expired_claims.exp = now - 3600; // Expired 1 hour ago
709        expired_claims.iat = now - 7200; // Issued 2 hours ago
710        expired_claims.nbf = now - 7200;
711
712        let token = signer.sign(expired_claims).unwrap();
713
714        let request = Request::builder()
715            .uri(format!("/ws?hs_token={}", token))
716            .body(())
717            .expect("request should build");
718        let auth_request = ConnectionAuthRequest::from_http_request(
719            "127.0.0.1:8877".parse().expect("socket addr should parse"),
720            &request,
721        );
722
723        let decision = plugin.authorize(&auth_request).await;
724        assert!(!decision.is_allowed());
725
726        if let AuthDecision::Deny(deny) = decision {
727            assert_eq!(deny.code, AuthErrorCode::TokenExpired);
728        } else {
729            panic!("Expected Deny decision for expired token");
730        }
731    }
732
733    #[tokio::test]
734    async fn signed_session_plugin_denies_invalid_signature() {
735        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
736
737        // Create two different key pairs
738        let signing_key = arete_auth::SigningKey::generate();
739        let wrong_key = arete_auth::SigningKey::generate();
740
741        // Sign with one key, verify with another
742        let signer = TokenSigner::new(signing_key, "test-issuer");
743        let wrong_verifying_key = wrong_key.verifying_key();
744        let verifier =
745            arete_auth::TokenVerifier::new(wrong_verifying_key, "test-issuer", "test-audience");
746        let plugin = SignedSessionAuthPlugin::new(verifier);
747
748        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
749            .with_scope("read")
750            .with_key_class(KeyClass::Secret)
751            .build();
752
753        let token = signer.sign(claims).unwrap();
754
755        let request = Request::builder()
756            .uri(format!("/ws?hs_token={}", token))
757            .body(())
758            .expect("request should build");
759        let auth_request = ConnectionAuthRequest::from_http_request(
760            "127.0.0.1:8877".parse().expect("socket addr should parse"),
761            &request,
762        );
763
764        let decision = plugin.authorize(&auth_request).await;
765        assert!(!decision.is_allowed());
766
767        if let AuthDecision::Deny(deny) = decision {
768            assert_eq!(deny.code, AuthErrorCode::TokenInvalidSignature);
769        } else {
770            panic!("Expected Deny decision for invalid signature");
771        }
772    }
773
774    #[tokio::test]
775    async fn signed_session_plugin_denies_wrong_audience() {
776        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
777
778        let signing_key = arete_auth::SigningKey::generate();
779        let verifying_key = signing_key.verifying_key();
780        let signer = TokenSigner::new(signing_key, "test-issuer");
781
782        // Verifier expects "test-audience", token is for "wrong-audience"
783        let verifier =
784            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
785        let plugin = SignedSessionAuthPlugin::new(verifier);
786
787        let claims = SessionClaims::builder("test-issuer", "test-subject", "wrong-audience")
788            .with_scope("read")
789            .with_key_class(KeyClass::Secret)
790            .build();
791
792        let token = signer.sign(claims).unwrap();
793
794        let request = Request::builder()
795            .uri(format!("/ws?hs_token={}", token))
796            .body(())
797            .expect("request should build");
798        let auth_request = ConnectionAuthRequest::from_http_request(
799            "127.0.0.1:8877".parse().expect("socket addr should parse"),
800            &request,
801        );
802
803        let decision = plugin.authorize(&auth_request).await;
804        assert!(!decision.is_allowed());
805
806        if let AuthDecision::Deny(deny) = decision {
807            assert_eq!(deny.code, AuthErrorCode::TokenInvalidAudience);
808        } else {
809            panic!("Expected Deny decision for wrong audience");
810        }
811    }
812
813    #[tokio::test]
814    async fn signed_session_plugin_denies_origin_mismatch() {
815        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
816
817        let signing_key = arete_auth::SigningKey::generate();
818        let verifying_key = signing_key.verifying_key();
819        let signer = TokenSigner::new(signing_key, "test-issuer");
820
821        // Verifier requires origin validation
822        let verifier =
823            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
824                .with_origin_validation();
825        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
826
827        // Token bound to specific origin
828        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
829            .with_scope("read")
830            .with_key_class(KeyClass::Secret)
831            .with_origin("https://allowed.example.com")
832            .build();
833
834        let token = signer.sign(claims).unwrap();
835
836        // Request from different origin
837        let request = Request::builder()
838            .uri(format!("/ws?hs_token={}", token))
839            .header("Origin", "https://evil.example.com")
840            .body(())
841            .expect("request should build");
842        let auth_request = ConnectionAuthRequest::from_http_request(
843            "127.0.0.1:8877".parse().expect("socket addr should parse"),
844            &request,
845        );
846
847        let decision = plugin.authorize(&auth_request).await;
848        assert!(!decision.is_allowed());
849
850        if let AuthDecision::Deny(deny) = decision {
851            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
852        } else {
853            panic!("Expected Deny decision for origin mismatch");
854        }
855    }
856
857    #[tokio::test]
858    async fn signed_session_plugin_allows_valid_token() {
859        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
860
861        let signing_key = arete_auth::SigningKey::generate();
862        let verifying_key = signing_key.verifying_key();
863        let signer = TokenSigner::new(signing_key, "test-issuer");
864        let verifier =
865            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
866        let plugin = SignedSessionAuthPlugin::new(verifier);
867
868        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
869            .with_scope("read")
870            .with_key_class(KeyClass::Secret)
871            .with_metering_key("meter-123")
872            .build();
873
874        let token = signer.sign(claims).unwrap();
875
876        let request = Request::builder()
877            .uri(format!("/ws?hs_token={}", token))
878            .body(())
879            .expect("request should build");
880        let auth_request = ConnectionAuthRequest::from_http_request(
881            "127.0.0.1:8877".parse().expect("socket addr should parse"),
882            &request,
883        );
884
885        let decision = plugin.authorize(&auth_request).await;
886        assert!(decision.is_allowed());
887
888        if let AuthDecision::Allow(ctx) = decision {
889            assert_eq!(ctx.subject, "test-subject");
890            assert_eq!(ctx.metering_key, "meter-123");
891            assert_eq!(ctx.key_class, KeyClass::Secret);
892        } else {
893            panic!("Expected Allow decision");
894        }
895    }
896
897    #[tokio::test]
898    async fn signed_session_plugin_allows_with_matching_origin() {
899        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
900
901        let signing_key = arete_auth::SigningKey::generate();
902        let verifying_key = signing_key.verifying_key();
903        let signer = TokenSigner::new(signing_key, "test-issuer");
904
905        let verifier =
906            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
907                .with_origin_validation();
908        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
909
910        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
911            .with_scope("read")
912            .with_key_class(KeyClass::Secret)
913            .with_origin("https://trusted.example.com")
914            .build();
915
916        let token = signer.sign(claims).unwrap();
917
918        let request = Request::builder()
919            .uri(format!("/ws?hs_token={}", token))
920            .header("Origin", "https://trusted.example.com")
921            .body(())
922            .expect("request should build");
923        let auth_request = ConnectionAuthRequest::from_http_request(
924            "127.0.0.1:8877".parse().expect("socket addr should parse"),
925            &request,
926        );
927
928        let decision = plugin.authorize(&auth_request).await;
929        assert!(decision.is_allowed());
930
931        if let AuthDecision::Allow(ctx) = decision {
932            assert_eq!(ctx.origin, Some("https://trusted.example.com".to_string()));
933        } else {
934            panic!("Expected Allow decision");
935        }
936    }
937
938    #[tokio::test]
939    async fn signed_session_plugin_allows_token_with_origin_when_no_origin_provided_and_not_required(
940    ) {
941        // This tests the non-browser client scenario (Rust, Python, etc.)
942        // where the client doesn't send an Origin header.
943        // The token has an origin claim from when it was minted via browser/API,
944        // but when the plugin doesn't require origin, the connection should still be allowed.
945        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
946
947        let signing_key = arete_auth::SigningKey::generate();
948        let verifying_key = signing_key.verifying_key();
949        let signer = TokenSigner::new(signing_key, "test-issuer");
950
951        // Plugin WITHOUT origin validation (default for public stacks)
952        let verifier =
953            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
954        let plugin = SignedSessionAuthPlugin::new(verifier);
955
956        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
957            .with_scope("read")
958            .with_key_class(KeyClass::Publishable)
959            .with_origin("https://example.com") // Token has origin claim
960            .build();
961
962        let token = signer.sign(claims).unwrap();
963
964        // No Origin header provided (simulating non-browser client)
965        let request = Request::builder()
966            .uri(format!("/ws?hs_token={}", token))
967            .body(())
968            .expect("request should build");
969        let auth_request = ConnectionAuthRequest::from_http_request(
970            "127.0.0.1:8877".parse().expect("socket addr should parse"),
971            &request,
972        );
973
974        // Should succeed even without Origin header
975        let decision = plugin.authorize(&auth_request).await;
976        assert!(
977            decision.is_allowed(),
978            "Expected Allow decision for non-browser client without Origin"
979        );
980
981        if let AuthDecision::Allow(ctx) = decision {
982            assert_eq!(ctx.origin, Some("https://example.com".to_string()));
983        } else {
984            panic!("Expected Allow decision");
985        }
986    }
987
988    #[tokio::test]
989    async fn signed_session_plugin_validates_origin_when_provided_even_when_not_required() {
990        // When origin IS provided, it should still be validated against the token
991        // even when require_origin is false (defense-in-depth)
992        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
993
994        let signing_key = arete_auth::SigningKey::generate();
995        let verifying_key = signing_key.verifying_key();
996        let signer = TokenSigner::new(signing_key, "test-issuer");
997
998        // Plugin WITHOUT origin validation (default)
999        let verifier =
1000            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
1001        let plugin = SignedSessionAuthPlugin::new(verifier);
1002
1003        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1004            .with_scope("read")
1005            .with_key_class(KeyClass::Publishable)
1006            .with_origin("https://allowed.example.com")
1007            .build();
1008
1009        let token = signer.sign(claims).unwrap();
1010
1011        // Origin provided and matches - should succeed
1012        let request = Request::builder()
1013            .uri(format!("/ws?hs_token={}", token))
1014            .header("Origin", "https://allowed.example.com")
1015            .body(())
1016            .expect("request should build");
1017        let auth_request = ConnectionAuthRequest::from_http_request(
1018            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1019            &request,
1020        );
1021
1022        let decision = plugin.authorize(&auth_request).await;
1023        assert!(decision.is_allowed());
1024
1025        // Origin provided but doesn't match - should fail
1026        let request = Request::builder()
1027            .uri(format!("/ws?hs_token={}", token))
1028            .header("Origin", "https://evil.example.com")
1029            .body(())
1030            .expect("request should build");
1031        let auth_request = ConnectionAuthRequest::from_http_request(
1032            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1033            &request,
1034        );
1035
1036        let decision = plugin.authorize(&auth_request).await;
1037        assert!(!decision.is_allowed());
1038
1039        if let AuthDecision::Deny(deny) = decision {
1040            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
1041        } else {
1042            panic!("Expected Deny decision for origin mismatch");
1043        }
1044    }
1045
1046    // Tests for AuthErrorCode utility methods
1047    #[test]
1048    fn auth_error_code_should_retry_logic() {
1049        assert!(AuthErrorCode::RateLimitExceeded.should_retry());
1050        assert!(AuthErrorCode::InternalError.should_retry());
1051        assert!(!AuthErrorCode::TokenExpired.should_retry());
1052        assert!(!AuthErrorCode::TokenInvalidSignature.should_retry());
1053        assert!(!AuthErrorCode::TokenMissing.should_retry());
1054    }
1055
1056    #[test]
1057    fn auth_error_code_should_refresh_token_logic() {
1058        assert!(AuthErrorCode::TokenExpired.should_refresh_token());
1059        assert!(AuthErrorCode::TokenInvalidSignature.should_refresh_token());
1060        assert!(AuthErrorCode::TokenInvalidFormat.should_refresh_token());
1061        assert!(AuthErrorCode::TokenInvalidIssuer.should_refresh_token());
1062        assert!(AuthErrorCode::TokenInvalidAudience.should_refresh_token());
1063        assert!(AuthErrorCode::TokenKeyNotFound.should_refresh_token());
1064        assert!(!AuthErrorCode::TokenMissing.should_refresh_token());
1065        assert!(!AuthErrorCode::RateLimitExceeded.should_refresh_token());
1066        assert!(!AuthErrorCode::ConnectionLimitExceeded.should_refresh_token());
1067    }
1068
1069    #[test]
1070    fn auth_error_code_string_representation() {
1071        assert_eq!(AuthErrorCode::TokenMissing.as_str(), "token-missing");
1072        assert_eq!(AuthErrorCode::TokenExpired.as_str(), "token-expired");
1073        assert_eq!(
1074            AuthErrorCode::TokenInvalidSignature.as_str(),
1075            "token-invalid-signature"
1076        );
1077        assert_eq!(
1078            AuthErrorCode::RateLimitExceeded.as_str(),
1079            "rate-limit-exceeded"
1080        );
1081        assert_eq!(
1082            AuthErrorCode::ConnectionLimitExceeded.as_str(),
1083            "connection-limit-exceeded"
1084        );
1085    }
1086
1087    // Tests for AuthDeny construction
1088    #[test]
1089    fn auth_deny_token_missing_factory() {
1090        let deny = AuthDeny::token_missing();
1091        assert_eq!(deny.code, AuthErrorCode::TokenMissing);
1092        assert!(deny.reason.contains("Missing session token"));
1093    }
1094
1095    #[test]
1096    fn auth_deny_from_verify_error_mapping() {
1097        use arete_auth::VerifyError;
1098
1099        let test_cases = vec![
1100            (VerifyError::Expired, AuthErrorCode::TokenExpired),
1101            (
1102                VerifyError::InvalidSignature,
1103                AuthErrorCode::TokenInvalidSignature,
1104            ),
1105            (
1106                VerifyError::InvalidIssuer,
1107                AuthErrorCode::TokenInvalidIssuer,
1108            ),
1109            (
1110                VerifyError::InvalidAudience,
1111                AuthErrorCode::TokenInvalidAudience,
1112            ),
1113            (
1114                VerifyError::KeyNotFound("kid123".to_string()),
1115                AuthErrorCode::TokenKeyNotFound,
1116            ),
1117            (
1118                VerifyError::OriginMismatch {
1119                    expected: "a".to_string(),
1120                    actual: "b".to_string(),
1121                },
1122                AuthErrorCode::OriginMismatch,
1123            ),
1124        ];
1125
1126        for (err, expected_code) in test_cases {
1127            let deny = AuthDeny::from_verify_error(err);
1128            assert_eq!(deny.code, expected_code);
1129        }
1130    }
1131
1132    // Tests for multiple auth failure scenarios in sequence
1133    #[tokio::test]
1134    async fn signed_session_plugin_handles_multiple_failure_reasons() {
1135        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1136
1137        let signing_key = arete_auth::SigningKey::generate();
1138        let verifying_key = signing_key.verifying_key();
1139        let signer = TokenSigner::new(signing_key, "test-issuer");
1140        let verifier =
1141            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
1142                .with_origin_validation();
1143        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
1144
1145        // Test 1: Missing token
1146        let request = Request::builder()
1147            .uri("/ws")
1148            .body(())
1149            .expect("request should build");
1150        let auth_request = ConnectionAuthRequest::from_http_request(
1151            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1152            &request,
1153        );
1154        let decision = plugin.authorize(&auth_request).await;
1155        assert!(!decision.is_allowed());
1156        match decision {
1157            AuthDecision::Deny(deny) => assert_eq!(deny.code, AuthErrorCode::TokenMissing),
1158            _ => panic!("Expected Deny decision"),
1159        }
1160
1161        // Test 2: Valid token with wrong origin
1162        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1163            .with_scope("read")
1164            .with_key_class(KeyClass::Secret)
1165            .with_origin("https://allowed.example.com")
1166            .build();
1167        let token = signer.sign(claims).unwrap();
1168
1169        let request = Request::builder()
1170            .uri(format!("/ws?hs_token={}", token))
1171            .header("Origin", "https://evil.example.com")
1172            .body(())
1173            .expect("request should build");
1174        let auth_request = ConnectionAuthRequest::from_http_request(
1175            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1176            &request,
1177        );
1178        let decision = plugin.authorize(&auth_request).await;
1179        assert!(!decision.is_allowed());
1180        match decision {
1181            AuthDecision::Deny(deny) => assert_eq!(deny.code, AuthErrorCode::OriginMismatch),
1182            _ => panic!("Expected Deny decision for origin mismatch"),
1183        }
1184
1185        // Test 3: Valid token with correct origin
1186        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1187            .with_scope("read")
1188            .with_key_class(KeyClass::Secret)
1189            .with_origin("https://allowed.example.com")
1190            .build();
1191        let token = signer.sign(claims).unwrap();
1192
1193        let request = Request::builder()
1194            .uri(format!("/ws?hs_token={}", token))
1195            .header("Origin", "https://allowed.example.com")
1196            .body(())
1197            .expect("request should build");
1198        let auth_request = ConnectionAuthRequest::from_http_request(
1199            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1200            &request,
1201        );
1202        let decision = plugin.authorize(&auth_request).await;
1203        assert!(decision.is_allowed());
1204    }
1205
1206    // Test for rate limit error code
1207    #[tokio::test]
1208    async fn auth_deney_with_rate_limit_code() {
1209        let deny = AuthDeny::new(
1210            AuthErrorCode::RateLimitExceeded,
1211            "Too many requests from this IP",
1212        );
1213        assert_eq!(deny.code, AuthErrorCode::RateLimitExceeded);
1214        assert!(deny.code.should_retry());
1215        assert!(!deny.code.should_refresh_token());
1216    }
1217
1218    // Test for connection limit error code
1219    #[tokio::test]
1220    async fn auth_deny_with_connection_limit_code() {
1221        let deny = AuthDeny::new(
1222            AuthErrorCode::ConnectionLimitExceeded,
1223            "Maximum connections exceeded for subject user-123",
1224        );
1225        assert_eq!(deny.code, AuthErrorCode::ConnectionLimitExceeded);
1226        assert!(!deny.code.should_retry());
1227        assert!(!deny.code.should_refresh_token());
1228    }
1229
1230    // Integration-style test: Token extraction from various sources
1231    #[test]
1232    fn token_extraction_priority() {
1233        // Header takes priority over query param
1234        let request = Request::builder()
1235            .uri("/ws?hs_token=query-value")
1236            .header("Authorization", "Bearer header-value")
1237            .body(())
1238            .expect("request should build");
1239        let auth_request = ConnectionAuthRequest::from_http_request(
1240            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1241            &request,
1242        );
1243
1244        // bearer_token should return header value
1245        assert_eq!(auth_request.bearer_token(), Some("header-value"));
1246        // query_param should return query value
1247        assert_eq!(auth_request.query_param("hs_token"), Some("query-value"));
1248    }
1249
1250    // Test malformed authorization header handling
1251    #[test]
1252    fn malformed_authorization_header() {
1253        let test_cases = vec![
1254            ("Basic dXNlcjpwYXNz", None),                // Wrong scheme
1255            ("Bearer", None),                            // Missing token (no space after Bearer)
1256            ("", None),                                  // Empty
1257            ("Bearer token extra", Some("token extra")), // Extra parts (token includes everything after scheme)
1258        ];
1259
1260        for (header_value, expected) in test_cases {
1261            let request = Request::builder()
1262                .uri("/ws")
1263                .header("Authorization", header_value)
1264                .body(())
1265                .expect("request should build");
1266            let auth_request = ConnectionAuthRequest::from_http_request(
1267                "127.0.0.1:8877".parse().expect("socket addr should parse"),
1268                &request,
1269            );
1270            assert_eq!(
1271                auth_request.bearer_token(),
1272                expected,
1273                "Failed for header: {}",
1274                header_value
1275            );
1276        }
1277    }
1278
1279    // ============================================
1280    // WEBSOCKET HANDSHAKE AUTH FAILURE TESTS
1281    // ============================================
1282    // These tests simulate real-world handshake failure scenarios
1283
1284    #[test]
1285    fn auth_deny_error_response_structure() {
1286        let deny = AuthDeny::new(AuthErrorCode::TokenExpired, "Token has expired")
1287            .with_field("exp")
1288            .with_context("Token expired 5 minutes ago")
1289            .with_suggested_action("Refresh your authentication token")
1290            .with_docs_url("https://docs.arete.run/auth/errors#token-expired");
1291
1292        let response = deny.to_error_response();
1293
1294        assert_eq!(response.code, "token-expired");
1295        assert_eq!(response.message, "Token has expired");
1296        assert_eq!(response.error, "token-expired");
1297        assert!(response.retryable);
1298        assert_eq!(
1299            response.suggested_action,
1300            Some("Refresh your authentication token".to_string())
1301        );
1302        assert_eq!(
1303            response.docs_url,
1304            Some("https://docs.arete.run/auth/errors#token-expired".to_string())
1305        );
1306    }
1307
1308    #[test]
1309    fn auth_deny_rate_limited_response() {
1310        use std::time::Duration;
1311
1312        let deny = AuthDeny::rate_limited(Duration::from_secs(30), "websocket connections");
1313        let response = deny.to_error_response();
1314
1315        assert_eq!(response.code, "rate-limit-exceeded");
1316        assert!(response.message.contains("30s"));
1317        assert!(response.retryable);
1318        assert_eq!(response.retry_after, Some(30));
1319    }
1320
1321    #[test]
1322    fn auth_deny_connection_limit_response() {
1323        let deny = AuthDeny::connection_limit_exceeded("user-123", 5, 5);
1324        let response = deny.to_error_response();
1325
1326        assert_eq!(response.code, "connection-limit-exceeded");
1327        assert!(response.message.contains("user-123"));
1328        assert!(response.message.contains("5 of 5"));
1329        assert!(response.retryable); // Connection limits are retryable (may become available)
1330    }
1331
1332    #[test]
1333    fn retry_policy_immediate() {
1334        let deny = AuthDeny::new(AuthErrorCode::InternalError, "Transient error")
1335            .with_retry_policy(RetryPolicy::RetryImmediately);
1336
1337        assert_eq!(deny.retry_policy, RetryPolicy::RetryImmediately);
1338    }
1339
1340    #[test]
1341    fn retry_policy_with_backoff() {
1342        use std::time::Duration;
1343
1344        let deny = AuthDeny::new(AuthErrorCode::RateLimitExceeded, "Too many requests")
1345            .with_retry_policy(RetryPolicy::RetryWithBackoff {
1346                initial: Duration::from_secs(1),
1347                max: Duration::from_secs(60),
1348            });
1349
1350        match deny.retry_policy {
1351            RetryPolicy::RetryWithBackoff { initial, max } => {
1352                assert_eq!(initial, Duration::from_secs(1));
1353                assert_eq!(max, Duration::from_secs(60));
1354            }
1355            _ => panic!("Expected RetryWithBackoff"),
1356        }
1357    }
1358
1359    #[test]
1360    fn auth_error_code_http_status_mapping() {
1361        assert_eq!(AuthErrorCode::TokenMissing.http_status(), 401);
1362        assert_eq!(AuthErrorCode::TokenExpired.http_status(), 401);
1363        assert_eq!(AuthErrorCode::TokenInvalidSignature.http_status(), 401);
1364        assert_eq!(AuthErrorCode::OriginMismatch.http_status(), 403);
1365        assert_eq!(AuthErrorCode::RateLimitExceeded.http_status(), 429);
1366        assert_eq!(AuthErrorCode::ConnectionLimitExceeded.http_status(), 429);
1367        assert_eq!(AuthErrorCode::InternalError.http_status(), 500);
1368    }
1369
1370    #[test]
1371    fn auth_error_code_default_retry_policies() {
1372        use std::time::Duration;
1373
1374        // Should refresh token
1375        assert!(matches!(
1376            AuthErrorCode::TokenExpired.default_retry_policy(),
1377            RetryPolicy::RetryWithFreshToken
1378        ));
1379        assert!(matches!(
1380            AuthErrorCode::TokenInvalidSignature.default_retry_policy(),
1381            RetryPolicy::RetryWithFreshToken
1382        ));
1383
1384        // Should retry with backoff
1385        assert!(matches!(
1386            AuthErrorCode::RateLimitExceeded.default_retry_policy(),
1387            RetryPolicy::RetryWithBackoff { .. }
1388        ));
1389        assert!(matches!(
1390            AuthErrorCode::InternalError.default_retry_policy(),
1391            RetryPolicy::RetryWithBackoff { .. }
1392        ));
1393
1394        // Should not retry
1395        assert!(matches!(
1396            AuthErrorCode::TokenMissing.default_retry_policy(),
1397            RetryPolicy::NoRetry
1398        ));
1399        assert!(matches!(
1400            AuthErrorCode::OriginMismatch.default_retry_policy(),
1401            RetryPolicy::NoRetry
1402        ));
1403    }
1404
1405    // Simulated handshake scenarios
1406
1407    #[tokio::test]
1408    async fn handshake_rejects_missing_token_with_proper_error() {
1409        use tokio_tungstenite::tungstenite::http::StatusCode;
1410
1411        let plugin = AllowAllAuthPlugin;
1412
1413        // Create a request without a token
1414        let request = Request::builder()
1415            .uri("/ws")
1416            .body(())
1417            .expect("request should build");
1418
1419        let auth_request = ConnectionAuthRequest::from_http_request(
1420            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1421            &request,
1422        );
1423
1424        // For this test, we'll use a plugin that requires tokens
1425        // Actually AllowAllAuthPlugin doesn't require tokens, so let's create a static token plugin
1426        let static_plugin = StaticTokenAuthPlugin::new(["valid-token".to_string()]);
1427        let decision = static_plugin.authorize(&auth_request).await;
1428
1429        assert!(!decision.is_allowed());
1430
1431        if let AuthDecision::Deny(deny) = decision {
1432            assert_eq!(deny.code, AuthErrorCode::TokenMissing);
1433            assert_eq!(deny.http_status, 401);
1434            assert!(deny.reason.contains("Missing"));
1435        } else {
1436            panic!("Expected Deny decision");
1437        }
1438    }
1439
1440    #[tokio::test]
1441    async fn handshake_rejects_expired_token_with_retry_hint() {
1442        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1443        use std::time::{SystemTime, UNIX_EPOCH};
1444
1445        let signing_key = arete_auth::SigningKey::generate();
1446        let verifying_key = signing_key.verifying_key();
1447        let signer = TokenSigner::new(signing_key, "test-issuer");
1448
1449        // Create an expired token
1450        let now = SystemTime::now()
1451            .duration_since(UNIX_EPOCH)
1452            .unwrap()
1453            .as_secs();
1454        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1455            .with_scope("read")
1456            .with_key_class(KeyClass::Secret)
1457            .build();
1458
1459        let mut expired_claims = claims;
1460        expired_claims.exp = now - 3600;
1461        expired_claims.iat = now - 7200;
1462        expired_claims.nbf = now - 7200;
1463
1464        let token = signer.sign(expired_claims).unwrap();
1465
1466        // Create verifier and plugin
1467        let verifier =
1468            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience");
1469        let plugin = SignedSessionAuthPlugin::new(verifier);
1470
1471        let request = Request::builder()
1472            .uri(format!("/ws?hs_token={}", token))
1473            .body(())
1474            .expect("request should build");
1475
1476        let auth_request = ConnectionAuthRequest::from_http_request(
1477            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1478            &request,
1479        );
1480
1481        let decision = plugin.authorize(&auth_request).await;
1482
1483        assert!(!decision.is_allowed());
1484
1485        if let AuthDecision::Deny(deny) = decision {
1486            assert_eq!(deny.code, AuthErrorCode::TokenExpired);
1487            assert_eq!(deny.http_status, 401);
1488            // Should suggest refreshing the token
1489            assert!(matches!(
1490                deny.retry_policy,
1491                RetryPolicy::RetryWithFreshToken
1492            ));
1493        } else {
1494            panic!("Expected Deny decision");
1495        }
1496    }
1497
1498    #[tokio::test]
1499    async fn handshake_rejects_invalid_signature_with_retry_hint() {
1500        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1501
1502        // Create two different key pairs
1503        let signing_key = arete_auth::SigningKey::generate();
1504        let wrong_key = arete_auth::SigningKey::generate();
1505
1506        // Sign with one key, verify with another
1507        let signer = TokenSigner::new(signing_key, "test-issuer");
1508        let wrong_verifying_key = wrong_key.verifying_key();
1509        let verifier =
1510            arete_auth::TokenVerifier::new(wrong_verifying_key, "test-issuer", "test-audience");
1511        let plugin = SignedSessionAuthPlugin::new(verifier);
1512
1513        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1514            .with_scope("read")
1515            .with_key_class(KeyClass::Secret)
1516            .build();
1517
1518        let token = signer.sign(claims).unwrap();
1519
1520        let request = Request::builder()
1521            .uri(format!("/ws?hs_token={}", token))
1522            .body(())
1523            .expect("request should build");
1524
1525        let auth_request = ConnectionAuthRequest::from_http_request(
1526            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1527            &request,
1528        );
1529
1530        let decision = plugin.authorize(&auth_request).await;
1531
1532        assert!(!decision.is_allowed());
1533
1534        if let AuthDecision::Deny(deny) = decision {
1535            assert_eq!(deny.code, AuthErrorCode::TokenInvalidSignature);
1536            assert_eq!(deny.http_status, 401);
1537            // Should suggest refreshing the token
1538            assert!(matches!(
1539                deny.retry_policy,
1540                RetryPolicy::RetryWithFreshToken
1541            ));
1542        } else {
1543            panic!("Expected Deny decision");
1544        }
1545    }
1546
1547    #[tokio::test]
1548    async fn handshake_rejects_origin_mismatch_without_retry() {
1549        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1550
1551        let signing_key = arete_auth::SigningKey::generate();
1552        let verifying_key = signing_key.verifying_key();
1553        let signer = TokenSigner::new(signing_key, "test-issuer");
1554
1555        let verifier =
1556            arete_auth::TokenVerifier::new(verifying_key, "test-issuer", "test-audience")
1557                .with_origin_validation();
1558        let plugin = SignedSessionAuthPlugin::new(verifier).with_origin_validation();
1559
1560        // Token bound to specific origin
1561        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
1562            .with_scope("read")
1563            .with_key_class(KeyClass::Secret)
1564            .with_origin("https://allowed.example.com")
1565            .build();
1566
1567        let token = signer.sign(claims).unwrap();
1568
1569        // Request from different origin
1570        let request = Request::builder()
1571            .uri(format!("/ws?hs_token={}", token))
1572            .header("Origin", "https://evil.example.com")
1573            .body(())
1574            .expect("request should build");
1575
1576        let auth_request = ConnectionAuthRequest::from_http_request(
1577            "127.0.0.1:8877".parse().expect("socket addr should parse"),
1578            &request,
1579        );
1580
1581        let decision = plugin.authorize(&auth_request).await;
1582
1583        assert!(!decision.is_allowed());
1584
1585        if let AuthDecision::Deny(deny) = decision {
1586            assert_eq!(deny.code, AuthErrorCode::OriginMismatch);
1587            assert_eq!(deny.http_status, 403);
1588            // Should NOT suggest retrying - this is a security issue
1589            assert!(matches!(deny.retry_policy, RetryPolicy::NoRetry));
1590        } else {
1591            panic!("Expected Deny decision");
1592        }
1593    }
1594
1595    // Test that AuthDeny can be converted to HTTP error response
1596    #[test]
1597    fn auth_deny_to_http_response() {
1598        let deny = AuthDeny::new(AuthErrorCode::RateLimitExceeded, "Too many requests")
1599            .with_suggested_action("Wait before retrying")
1600            .with_retry_policy(RetryPolicy::RetryAfter(Duration::from_secs(30)));
1601
1602        let response = deny.to_error_response();
1603
1604        // Verify the response is serializable
1605        let json = serde_json::to_string(&response).expect("Should serialize");
1606        assert!(json.contains("rate-limit-exceeded"));
1607        assert!(json.contains("Too many requests"));
1608        assert!(json.contains("Wait before retrying"));
1609        assert!(json.contains("\"retryable\":true"));
1610        assert!(json.contains("\"retry_after\":30"));
1611    }
1612
1613    // Test comprehensive error scenarios
1614    #[tokio::test]
1615    async fn comprehensive_auth_error_scenarios() {
1616        use arete_auth::{KeyClass, SessionClaims, TokenSigner};
1617
1618        let signing_key = arete_auth::SigningKey::generate();
1619        let verifying_key = signing_key.verifying_key();
1620        let signer = TokenSigner::new(signing_key, "test-issuer");
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
