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

arete_server/websocket/
auth.rs
