Implement rate limiting middleware for authentication endpoints, adding security headers middleware, and comprehensive error handling. Enhance validation tests for various DTOs and ensure proper functionality of gacha credits and rolls. Add unit tests for rate limiting and security headers middleware to validate behavior under different conditions.

2025-10-11 23:23:51 +07:00
parent 789b20c278
commit b6b5f48055
32 changed files with 1734 additions and 163 deletions
@@ -0,0 +1,296 @@
+#[cfg(test)]
+mod rate_limiting_middleware_tests {
+    use axum::{http::Request, middleware::Next, response::Response};
+    use imphnen_libs::{AppState, environment::Environment};
+    use imphnen_middleware::rate_limiting_middleware::{
+        RateLimitConfig, RateLimitStore, TokenBucket, create_rate_limiting_middleware,
+        auth_rate_limiting_middleware,
+    };
+    use std::{sync::Arc, time::Duration};
+    use tower::ServiceExt;
+
+    #[tokio::test]
+    async fn test_token_bucket_basic_functionality() {
+        let bucket = TokenBucket::new(5, 2); // Capacity 5, refill 2 per second
+        
+        // Should have full tokens initially
+        assert_eq!(bucket.tokens, 5);
+        
+        // Consume some tokens
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 4);
+        
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 3);
+        
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 2);
+        
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 1);
+        
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 0);
+        
+        // Should not consume when empty
+        assert!(!bucket.try_consume());
+        assert_eq!(bucket.tokens, 0);
+    }
+
+    #[tokio::test]
+    async fn test_token_bucket_refill() {
+        let mut bucket = TokenBucket::new(3, 1); // Capacity 3, refill 1 per second
+        
+        // Consume all tokens
+        for _ in 0..3 {
+            assert!(bucket.try_consume());
+        }
+        
+        assert!(!bucket.try_consume());
+        assert_eq!(bucket.tokens, 0);
+        
+        // Wait for 1 second to allow refill
+        tokio::time::sleep(Duration::from_secs(1)).await;
+        
+        // Should have 1 token after refill
+        bucket.refill_tokens();
+        assert_eq!(bucket.tokens, 1);
+        
+        // Consume the refilled token
+        assert!(bucket.try_consume());
+        assert_eq!(bucket.tokens, 0);
+        
+        // Wait another second
+        tokio::time::sleep(Duration::from_secs(1)).await;
+        
+        // Should have another token
+        bucket.refill_tokens();
+        assert_eq!(bucket.tokens, 1);
+    }
+
+    #[tokio::test]
+    async fn test_rate_limit_store_basic() {
+        let config = RateLimitConfig::test();
+        let store = Arc::new(RateLimitStore::new(config));
+        
+        let client_ip = "127.0.0.1";
+        
+        // First request should succeed
+        let result = store.check_limit(client_ip).await;
+        assert!(result.is_ok());
+        
+        // Multiple requests should succeed within limits
+        for _ in 0..config.bucket_size {
+            let result = store.check_limit(client_ip).await;
+            assert!(result.is_ok());
+        }
+        
+        // Next request should fail
+        let result = store.check_limit(client_ip).await;
+        assert!(result.is_err());
+        assert_eq!(result.unwrap_err(), axum::http::StatusCode::TOO_MANY_REQUESTS);
+    }
+
+    #[tokio::test]
+    async fn test_rate_limit_store_window_reset() {
+        let config = RateLimitConfig {
+            max_requests: 10,
+            window_duration: Duration::from_secs(2),
+            bucket_size: 2,
+            refill_rate: 1,
+        };
+        let store = Arc::new(RateLimitStore::new(config));
+        
+        let client_ip = "127.0.0.1";
+        
+        // Consume all tokens
+        assert!(store.check_limit(client_ip).await.is_ok());
+        assert!(store.check_limit(client_ip).await.is_ok());
+        assert!(store.check_limit(client_ip).await.is_err());
+        
+        // Wait for window to reset
+        tokio::time::sleep(Duration::from_secs(3)).await;
+        
+        // Should be able to make requests again
+        assert!(store.check_limit(client_ip).await.is_ok());
+        assert!(store.check_limit(client_ip).await.is_ok());
+        assert!(store.check_limit(client_ip).await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_different_clients_have_separate_limits() {
+        let config = RateLimitConfig::test();
+        let store = Arc::new(RateLimitStore::new(config));
+        
+        let client_ip_1 = "127.0.0.1";
+        let client_ip_2 = "127.0.0.2";
+        
+        // Client 1 should be able to make requests
+        for _ in 0..config.bucket_size {
+            assert!(store.check_limit(client_ip_1).await.is_ok());
+        }
+        assert!(store.check_limit(client_ip_1).await.is_err());
+        
+        // Client 2 should still be able to make requests
+        for _ in 0..config.bucket_size {
+            assert!(store.check_limit(client_ip_2).await.is_ok());
+        }
+        assert!(store.check_limit(client_ip_2).await.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_auth_rate_limiting_middleware_success() {
+        // Create a mock AppState with test environment
+        let state = AppState {
+            surrealdb_ws: Default::default(),
+            surrealdb_mem: Default::default(),
+            user_lookup_service: Default::default(),
+            auth_repository: Default::default(),
+            env: Environment::Test,
+        };
+
+        // Create a mock request to /auth/login
+        let mut request = Request::builder()
+            .uri("/v1/auth/login")
+            .header("x-forwarded-for", "127.0.0.1")
+            .body(())
+            .unwrap();
+
+        // Create a mock next service
+        let next = Next::new(|req| async move {
+            let response = Response::builder()
+                .status(200)
+                .body("Login successful")
+                .unwrap();
+            Ok::<_, axum::http::StatusCode>((req, response))
+        });
+
+        // Call the middleware
+        let result = auth_rate_limiting_middleware(
+            axum::Extension(state.clone()),
+            request,
+            next,
+        ).await;
+
+        // Should succeed
+        assert!(result.is_ok());
+        let response = result.unwrap();
+        assert_eq!(response.status(), 200);
+    }
+
+    #[tokio::test]
+    async fn test_auth_rate_limiting_middleware_429() {
+        // Create test configuration with very low limits for testing
+        let config = RateLimitConfig {
+            max_requests: 1,
+            window_duration: Duration::from_secs(10),
+            bucket_size: 1,
+            refill_rate: 1,
+        };
+
+        // Create a mock AppState with test environment
+        let state = AppState {
+            surrealdb_ws: Default::default(),
+            surrealdb_mem: Default::default(),
+            user_lookup_service: Default::default(),
+            auth_repository: Default::default(),
+            env: Environment::Test,
+        };
+
+        // Create a mock request to /auth/login
+        let mut request = Request::builder()
+            .uri("/v1/auth/login")
+            .header("x-forwarded-for", "127.0.0.1")
+            .body(())
+            .unwrap();
+
+        // Create a mock next service
+        let next = Next::new(|req| async move {
+            let response = Response::builder()
+                .status(200)
+                .body("Login successful")
+                .unwrap();
+            Ok::<_, axum::http::StatusCode>((req, response))
+        });
+
+        // First request should succeed
+        let result = auth_rate_limiting_middleware(
+            axum::Extension(state.clone()),
+            request.clone(),
+            next.clone(),
+        ).await;
+        assert!(result.is_ok());
+
+        // Second request should fail with 429
+        let result = auth_rate_limiting_middleware(
+            axum::Extension(state),
+            request,
+            next,
+        ).await;
+        assert!(result.is_ok());
+        let response = result.unwrap();
+        assert_eq!(response.status(), 429);
+        assert_eq!(response.headers().get("Retry-After").unwrap(), "60");
+    }
+
+    #[tokio::test]
+    async fn test_non_auth_endpoints_not_rate_limited() {
+        // Create a mock AppState with test environment
+        let state = AppState {
+            surrealdb_ws: Default::default(),
+            surrealdb_mem: Default::default(),
+            user_lookup_service: Default::default(),
+            auth_repository: Default::default(),
+            env: Environment::Test,
+        };
+
+        // Create a mock request to a non-auth endpoint
+        let mut request = Request::builder()
+            .uri("/v1/users/me")
+            .header("x-forwarded-for", "127.0.0.1")
+            .body(())
+            .unwrap();
+
+        // Create a mock next service
+        let next = Next::new(|req| async move {
+            let response = Response::builder()
+                .status(200)
+                .body("User data")
+                .unwrap();
+            Ok::<_, axum::http::StatusCode>((req, response))
+        });
+
+        // Call the middleware - should not apply rate limiting
+        let result = auth_rate_limiting_middleware(
+            axum::Extension(state),
+            request,
+            next,
+        ).await;
+
+        // Should succeed
+        assert!(result.is_ok());
+        let response = result.unwrap();
+        assert_eq!(response.status(), 200);
+    }
+
+    #[tokio::test]
+    async fn test_environment_specific_configurations() {
+        // Test development config
+        let dev_config = RateLimitConfig::development();
+        assert_eq!(dev_config.max_requests, 100);
+        assert_eq!(dev_config.bucket_size, 50);
+        assert_eq!(dev_config.refill_rate, 10);
+
+        // Test production config
+        let prod_config = RateLimitConfig::production();
+        assert_eq!(prod_config.max_requests, 10);
+        assert_eq!(prod_config.bucket_size, 5);
+        assert_eq!(prod_config.refill_rate, 1);
+
+        // Test test config
+        let test_config = RateLimitConfig::test();
+        assert_eq!(test_config.max_requests, 1000);
+        assert_eq!(test_config.bucket_size, 100);
+        assert_eq!(test_config.refill_rate, 20);
+    }
+}