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

arete_server/websocket/
auth.rs
