feat: implement authentication middleware with shared token state

Add transparent authentication middleware to ResilientClient using
the challenge-response protocol with proper concurrency control.
This implementation uses the shared auth types from the core library.

Signed-off-by: Sergio Correia <[email protected]>

Author: sergio-correia

keylime/src/resilient_client.rs

@@ -1,3 +1,5 @@
+use crate::auth::{AuthConfig, SessionToken};
+use anyhow;
 use async_trait::async_trait;
 use http::Extensions;
 use log::{debug, warn};
@@ -13,7 +15,9 @@ use reqwest_retry::{
     Retryable, RetryableStrategy,
 };
 use serde::Serialize;
+use std::sync::Arc;
 use std::time::Duration;
+use tokio::sync::{Mutex, RwLock};
 
 // We define a default maximum delay for retries, which in the pracitical sense
 // is set to 1 hour. This can be adjusted based on the application's needs.
@@ -82,6 +86,153 @@ impl RetryableStrategy for StopOnSuccessStrategy {
     }
 }
 
+/// Shared state for authentication tokens with proper concurrency control
+#[derive(Debug)]
+struct TokenState {
+    /// RwLock for the actual token - allows concurrent reads
+    token: RwLock<Option<SessionToken>>,
+    /// Mutex for refresh operations - ensures single writer
+    refresh_lock: Mutex<()>,
+    /// Authentication configuration
+    auth_config: AuthConfig,
+}
+
+impl TokenState {
+    fn new(auth_config: AuthConfig) -> Self {
+        Self {
+            token: RwLock::new(None),
+            refresh_lock: Mutex::new(()),
+            auth_config,
+        }
+    }
+
+    async fn get_valid_token(
+        &self,
+    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
+        // Fast path: try to read existing valid token
+        {
+            let token_guard = self.token.read().await;
+            if let Some(ref token) = *token_guard {
+                if token
+                    .is_valid(self.auth_config.token_refresh_buffer_minutes)
+                {
+                    debug!("Using existing valid token from middleware");
+                    return Ok(token.token.clone());
+                }
+            }
+        }
+
+        // Slow path: token is invalid or missing, need to refresh
+        self.refresh_token().await
+    }
+
+    async fn refresh_token(
+        &self,
+    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
+        // Acquire refresh lock to ensure only one refresh at a time
+        let _refresh_guard = self.refresh_lock.lock().await;
+
+        // Double-check: another request might have refreshed while we waited
+        {
+            let token_guard = self.token.read().await;
+            if let Some(ref token) = *token_guard {
+                if token
+                    .is_valid(self.auth_config.token_refresh_buffer_minutes)
+                {
+                    debug!("Token was refreshed by another request");
+                    return Ok(token.token.clone());
+                }
+            }
+        }
+
+        // TODO: Next, we'll integrate with the actual AuthenticationClient
+        // For now, this is a placeholder that will be replaced
+        warn!("Token refresh not yet implemented");
+        Err("Authentication not yet integrated".into())
+    }
+
+    async fn clear_token(&self) {
+        let mut token_guard = self.token.write().await;
+        *token_guard = None;
+        debug!("Authentication token cleared from shared state");
+    }
+}
+
+/// Middleware for transparent authentication using challenge-response protocol
+#[derive(Debug)]
+pub struct AuthenticationMiddleware {
+    token_state: Arc<TokenState>,
+}
+
+impl AuthenticationMiddleware {
+    pub fn new(auth_config: AuthConfig) -> Self {
+        let token_state = Arc::new(TokenState::new(auth_config));
+        Self { token_state }
+    }
+
+    fn is_auth_endpoint(&self, req: &reqwest::Request) -> bool {
+        let path = req.url().path();
+        // Skip authentication for auth endpoints to prevent infinite loops
+        path.contains("/sessions")
+    }
+}
+
+#[async_trait]
+impl Middleware for AuthenticationMiddleware {
+    async fn handle(
+        &self,
+        mut req: reqwest::Request,
+        extensions: &mut Extensions,
+        next: Next<'_>,
+    ) -> Result<Response, Error> {
+        // Skip authentication for auth endpoints to prevent infinite loops
+        if self.is_auth_endpoint(&req) {
+            debug!(
+                "Skipping auth for authentication endpoint: {}",
+                req.url().path()
+            );
+            return next.run(req, extensions).await;
+        }
+
+        // Add Authorization header if not present
+        if !req.headers().contains_key("Authorization") {
+            match self.token_state.get_valid_token().await {
+                Ok(token) => {
+                    debug!("Adding authentication token to request");
+                    req.headers_mut().insert(
+                        "Authorization",
+                        format!("Bearer {}", token).parse().map_err(|e| {
+                            Error::Middleware(anyhow::anyhow!(
+                                "Invalid token format: {}",
+                                e
+                            ))
+                        })?,
+                    );
+                }
+                Err(e) => {
+                    warn!("Failed to get auth token: {}", e);
+                    return Err(Error::Middleware(anyhow::anyhow!(
+                        "Authentication failed: {}",
+                        e
+                    )));
+                }
+            }
+        }
+
+        let response = next.run(req, extensions).await?;
+
+        // Handle 401 responses by clearing token
+        if response.status() == StatusCode::UNAUTHORIZED {
+            warn!("Received 401, clearing token for future requests");
+            self.token_state.clear_token().await;
+            // Note: We don't retry here to avoid infinite loops
+            // The retry will happen naturally on the next request
+        }
+
+        Ok(response)
+    }
+}
+
 /// A client that transparently handles retries with exponential backoff.
 #[derive(Debug, Clone)]
 pub struct ResilientClient {
@@ -120,6 +271,46 @@ impl ResilientClient {
         }
     }
 
+    /// Creates a new client with optional authentication middleware
+    pub fn new_with_auth(
+        client: Option<Client>,
+        auth_config: Option<AuthConfig>,
+        initial_delay: std::time::Duration,
+        max_retries: u32,
+        success_codes: &[StatusCode],
+        max_delay: Option<std::time::Duration>,
+    ) -> Self {
+        let base_client = client.unwrap_or_default();
+        let final_max_delay = max_delay.unwrap_or(DEFAULT_MAX_DELAY);
+
+        let retry_policy = ExponentialBackoff::builder()
+            .retry_bounds(initial_delay, final_max_delay)
+            .jitter(Jitter::None)
+            .build_with_max_retries(max_retries);
+
+        let mut builder = ClientBuilder::new(base_client).with(
+            RetryTransientMiddleware::new_with_policy_and_strategy(
+                retry_policy,
+                StopOnSuccessStrategy {
+                    success_codes: success_codes.to_vec(),
+                },
+            ),
+        );
+
+        // Add authentication middleware if config is provided
+        if let Some(auth_cfg) = auth_config {
+            debug!("Adding authentication middleware to client");
+            let auth_middleware = AuthenticationMiddleware::new(auth_cfg);
+            builder = builder.with(auth_middleware);
+        }
+
+        let client_with_middleware = builder.with(LoggingMiddleware).build();
+
+        Self {
+            client: client_with_middleware,
+        }
+    }
+
     /// Generates a six-character lowercase alphanumeric request ID.
     fn generate_request_id() -> String {
         const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789";
@@ -502,4 +693,79 @@ mod tests {
         assert!(response.is_ok());
         assert_eq!(response.unwrap().status(), StatusCode::OK); //#[allow_ci]
     }
+
+    #[tokio::test]
+    async fn test_resilient_client_with_auth_config() {
+        use crate::auth::AuthConfig;
+
+        let auth_config = AuthConfig {
+            verifier_base_url: "https://verifier.example.com".to_string(),
+            agent_id: "test-agent".to_string(),
+            avoid_tpm: true,
+            timeout_ms: 5000,
+            token_refresh_buffer_minutes: 5,
+            max_auth_retries: 3,
+        };
+
+        // Test with authentication
+        let _client_with_auth = ResilientClient::new_with_auth(
+            None,
+            Some(auth_config),
+            std::time::Duration::from_millis(10),
+            3,
+            &[StatusCode::OK],
+            None,
+        );
+
+        // Verify the client was created successfully
+    }
+
+    #[tokio::test]
+    async fn test_resilient_client_without_auth_config() {
+        // Test without authentication (should behave like the original client)
+        let _client_without_auth = ResilientClient::new_with_auth(
+            None,
+            None, // No auth config
+            std::time::Duration::from_millis(10),
+            3,
+            &[StatusCode::OK],
+            None,
+        );
+
+        // Verify the client was created successfully
+    }
+
+    #[tokio::test]
+    async fn test_authentication_middleware_path_detection() {
+        use crate::auth::AuthConfig;
+
+        let auth_config = AuthConfig {
+            verifier_base_url: "https://verifier.example.com".to_string(),
+            agent_id: "test-agent".to_string(),
+            avoid_tpm: true,
+            timeout_ms: 5000,
+            token_refresh_buffer_minutes: 5,
+            max_auth_retries: 3,
+        };
+
+        let middleware = AuthenticationMiddleware::new(auth_config);
+
+        // Mock a request to a sessions endpoint (should be detected as auth endpoint)
+        let mock_request = reqwest::Request::new(
+            Method::POST,
+            "https://verifier.example.com/v3.0/sessions"
+                .parse()
+                .unwrap(), //#[allow_ci]
+        );
+        assert!(middleware.is_auth_endpoint(&mock_request));
+
+        // Mock a request to a non-auth endpoint
+        let mock_request2 = reqwest::Request::new(
+            Method::GET,
+            "https://verifier.example.com/v3.0/agents/123/attestations"
+                .parse()
+                .unwrap(), //#[allow_ci]
+        );
+        assert!(!middleware.is_auth_endpoint(&mock_request2));
+    }
 }