Router error handling errors :((

Author: virjilakrum

Cargo.toml

@@ -5,49 +5,39 @@ edition = "2021"
 
 [dependencies]
 tokio = { version = "1.36", features = ["full"] }
-virt = { version = "0.4.1", features = ["snapshot"] }
+virt = "0.4.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tracing = "0.1"
 tracing-subscriber = "0.3"
 anyhow = "1.0"  
 async-trait = "0.1"
 config = "0.15.6"
-axum = { version = "0.8", features = ["macros"] }
+axum = { version = "0.8.0", features = ["macros"] }
 hyper = { version = "1.0", features = ["full"] }
 tower = { version = "0.5.2", features = ["limit", "util"] }
-tower-http = { version = "0.6.2", features = ["trace", "limit"] }
+tower-http = { version = "0.6.2", features = ["trace", "limit", "add-extension"] }
 clap = { version = "4.4", features = ["derive"] }
 colored = "3.0"
 thiserror = "2.0.11"
 chrono = "0.4"
 uuid = { version = "1.8.0", features = ["v4"] }
-libvirt = "0.1.0"
-governor = { version = "0.8.0", features = ["std", "nohashmap"] }
+governor = { version = "0.8", features = ["dashmap"] }
+jsonwebtoken = "8.3.0"
 
-# Platform-specific dependencies
 [target.'cfg(target_os = "linux")'.dependencies]
-libvirt = "0.1.0"
-nvml-wrapper = "0.10.0"
-glob = "0.3"
+nvml-wrapper = { version = "0.10.0", optional = true }
 
 [target.'cfg(target_os = "macos")'.dependencies]
-core-graphics = "0.23.2"
-metal = { version = "0.26.0", features = ["private"] }
+core-graphics = { version = "0.24.0", optional = true }
+metal = { version = "0.27.0", features = ["private"], optional = true }
 
 [target.'cfg(target_os = "windows")'.dependencies]
-winapi = { version = "0.3", features = ["dxgi", "d3dcommon"] }
-dxgi = "0.4"
-
-[lib]
-name = "gpu_share_vm_manager"
-path = "src/lib.rs"
-
-[[bin]]
-name = "gpu-share-vm-manager"
-path = "src/main.rs"
+dxgi = { version = "0.3.0-alpha4", optional = true }
+winapi = { version = "0.3", features = ["dxgi", "d3dcommon"], optional = true }
+windows = { version = "0.48", features = ["Win32_Graphics_Dxgi"] }
 
 [features]
-metal = ["dep:metal", "dep:core-graphics"]
-windows = ["dep:dxgi", "dep:winapi"]
-linux = ["dep:nvml-wrapper", "dep:libvirt"]
+default = ["metal"]
+metal = ["dep:core-graphics", "dep:metal"]
+windows = ["dep:dxgi", "winapi"]

src/api/middleware/mod.rs

@@ -0,0 +1 @@
+pub mod rate_limit; 

src/api/middleware/rate_limit.rs

@@ -2,11 +2,16 @@ use axum::{
     http::StatusCode,
     response::{IntoResponse, Response},
 };
-use std::{num::NonZeroU32, time::Duration};
-use tower::{
-    layer::util::{Stack, LayerFn},
-    Limit, RateLimitLayer,
+pub use governor::{
+    clock::QuantaClock,
+    middleware::NoOpMiddleware,
+    state::keyed::DashMapStateStore as DashMapStore,
+    Quota, RateLimiter,
 };
+use std::{num::NonZeroU32, sync::Arc, time::Duration};
+use tower::limit::RateLimitLayer;
+use std::error::Error as StdError;
+use std::fmt;
 
 /// Rate limiting configuration for API endpoints
 #[derive(Debug, Clone)]
@@ -18,40 +23,63 @@ pub struct RateLimitConfig {
 impl RateLimitConfig {
     /// Creates a new rate limiter layer based on configuration
     pub fn layer(&self) -> RateLimitLayer {
-        let window = Duration::from_secs(self.per_seconds);
-        RateLimitLayer::new(self.requests.get(), window)
+        let rate = self.requests.get() as u64;
+        let per = Duration::from_secs(self.per_seconds);
+        RateLimitLayer::new(rate, per)
     }
 }
 
 /// Global rate limiting configuration
+#[derive(Clone)]
 pub struct GlobalRateLimit {
     /// General API rate limits
-    pub api: RateLimitConfig,
+    pub api: Arc<RateLimiter<String, DashMapStore<String>, QuantaClock, NoOpMiddleware>>,
     /// Stricter limits for GPU operations
-    pub gpu_operations: RateLimitConfig,
+    pub gpu_operations: Arc<RateLimiter<String, DashMapStore<String>, QuantaClock, NoOpMiddleware>>,
     /// Authentication-specific limits
-    pub auth: RateLimitConfig,
+    pub auth: Arc<RateLimiter<String, DashMapStore<String>, QuantaClock, NoOpMiddleware>>,
 }
 
 impl Default for GlobalRateLimit {
     fn default() -> Self {
+        let clock = QuantaClock::default();
         Self {
-            api: RateLimitConfig {
-                requests: NonZeroU32::new(100).unwrap(),
-                per_seconds: 60,
-            },
-            gpu_operations: RateLimitConfig {
-                requests: NonZeroU32::new(30).unwrap(),
-                per_seconds: 60,
-            },
-            auth: RateLimitConfig {
-                requests: NonZeroU32::new(10).unwrap(),
-                per_seconds: 60,
-            },
+            api: Arc::new(
+                RateLimiter::dashmap_with_clock(
+                    Quota::per_second(NonZeroU32::new(5).unwrap()).allow_burst(NonZeroU32::new(10).unwrap()),
+                    clock.clone(),
+                )
+            ),
+            gpu_operations: Arc::new(
+                RateLimiter::dashmap_with_clock(
+                    Quota::per_minute(NonZeroU32::new(3).unwrap()).allow_burst(NonZeroU32::new(5).unwrap()),
+                    clock.clone(),
+                )
+            ),
+            auth: Arc::new(
+                RateLimiter::dashmap_with_clock(
+                    Quota::per_minute(NonZeroU32::new(10).unwrap()).allow_burst(NonZeroU32::new(15).unwrap()),
+                    clock,
+                )
+            ),
         }
     }
 }
 
+impl GlobalRateLimit {
+    pub fn api_quota(&self) -> Quota {
+        Quota::per_second(NonZeroU32::new(5).unwrap()).allow_burst(NonZeroU32::new(10).unwrap())
+    }
+
+    pub fn gpu_quota(&self) -> Quota {
+        Quota::per_minute(NonZeroU32::new(3).unwrap()).allow_burst(NonZeroU32::new(5).unwrap())
+    }
+
+    pub fn auth_quota(&self) -> Quota {
+        Quota::per_minute(NonZeroU32::new(10).unwrap()).allow_burst(NonZeroU32::new(15).unwrap())
+    }
+}
+
 /// Custom rate limit exceeded response
 #[derive(Debug)]
 pub struct RateLimitExceeded;
@@ -68,11 +96,72 @@ impl IntoResponse for RateLimitExceeded {
 
 /// Layer factory for rate limiting with custom response
 pub fn rate_limit_layer(
-    config: RateLimitConfig,
-) -> Stack<LayerFn<fn(Limit) -> Limit>, RateLimitLayer> {
-    let layer = config.layer();
-    tower::ServiceBuilder::new()
-        .layer(layer)
-        .map_err(|_| RateLimitExceeded)
-        .into_inner()
-} 
+    _limiter: Arc<RateLimiter<String, DashMapStore<String>, QuantaClock, NoOpMiddleware>>,
+) -> RateLimitLayer {
+    // Sabit rate limit değerleri
+    let rate = 100;
+    let per = Duration::from_secs(1);
+    RateLimitLayer::new(rate, per)
+}
+
+// Enhanced error handling for rate limits
+impl StdError for RateLimitExceeded {}
+
+impl fmt::Display for RateLimitExceeded {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Rate limit exceeded")
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum::body::Body;
+    use axum::http::Request;
+    use tower::{Service, ServiceExt};
+
+    #[tokio::test]
+    async fn test_rate_limiting() {
+        let config = RateLimitConfig {
+            requests: NonZeroU32::new(2).unwrap(),
+            per_seconds: 1,
+        };
+
+        let mut service = tower::ServiceBuilder::new()
+            .layer(config.layer())
+            .service(tower::service_fn(|_| async {
+                Ok::<_, std::convert::Infallible>(Response::new(Body::empty()))
+            }));
+
+
+        let response = service
+            .ready()
+            .await
+            .unwrap()
+            .call(Request::new(Body::empty()))
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+
+        let response = service
+            .ready()
+            .await
+            .unwrap()
+            .call(Request::new(Body::empty()))
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+
+        let response = service
+            .ready()
+            .await
+            .unwrap()
+            .call(Request::new(Body::empty()))
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
+    }
+}

src/api/mod.rs

@@ -1,3 +1,4 @@
+pub mod middleware;
 pub mod routes;
 
 pub use routes::{create_router, AppState};