arete_server/websocket/
auth.rs

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

arete_server/websocket/
auth.rs
