[feat]条件跨域请求约束添加

src/bootstrap/middleware/cors.rs

@@ -1,28 +1,110 @@
+use std::{collections::HashSet, sync::Arc};
+
 use axum_starter::{prepare, PrepareMiddlewareEffect};
 use http::{HeaderValue, Method};
 use tower_http::cors::CorsLayer;
 
+use std::{task::Poll};
+
+use tower::{Layer, Service};
+use tower_http::cors::{Any, Cors};
+
+
 pub trait CorsConfigTrait {
     fn allow_origins(&self) -> Vec<HeaderValue>;
 
     fn allow_methods(&self) -> Vec<Method>;
+
+    fn bypass_paths(&self) -> Arc<HashSet<String>>;
 }
 
-#[prepare(PrepareCors)]
-pub fn prepare_cors<T: CorsConfigTrait>(cfg: &T) -> CorsMiddleware {
-    CorsMiddleware(
-        CorsLayer::new()
-            .allow_origin(cfg.allow_origins())
-            .allow_methods(cfg.allow_methods()),
-    )
+#[derive(Clone)]
+pub struct ConditionCors<S> {
+    bypass_cors: Cors<S>,
+    default_cors: Cors<S>,
+    bypass_paths: Arc<HashSet<String>>,
 }
 
-pub struct CorsMiddleware(CorsLayer);
+impl<S, Req, Resp> Service<http::Request<Req>> for ConditionCors<S>
+where
+    S: Service<http::Request<Req>, Response = http::Response<Resp>>,
+    Req: Default,
+    Resp: Default,
+{
+    type Response = <Cors<S> as Service<http::Request<Req>>>::Response;
+
+    type Error = <Cors<S> as Service<http::Request<Req>>>::Error;
 
-impl<S> PrepareMiddlewareEffect<S> for CorsMiddleware {
-    type Middleware = CorsLayer;
+    type Future = <Cors<S> as Service<http::Request<Req>>>::Future;
+
+    fn poll_ready(
+        &mut self, cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        match (
+            self.bypass_cors.poll_ready(cx),
+            self.default_cors.poll_ready(cx),
+        ) {
+            (Poll::Ready(bypass), Poll::Ready(default)) => {
+                Poll::Ready(bypass.and_then(|_| default))
+            }
+            (Poll::Ready(_), Poll::Pending) => Poll::Pending,
+            (Poll::Pending, Poll::Ready(_)) => Poll::Pending,
+            (Poll::Pending, Poll::Pending) => Poll::Pending,
+        }
+    }
+
+    fn call(&mut self, req: http::Request<Req>) -> Self::Future {
+        let uri = req.uri().path();
+        if self.bypass_paths.contains(uri) {
+            self.bypass_cors.call(req)
+        } else {
+            self.default_cors.call(req)
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct ConditionCorsLayer {
+    bypass_cors: CorsLayer,
+    default_cors: CorsLayer,
+    bypass_paths: Arc<HashSet<String>>
+}
+
+impl<S: Clone> Layer<S> for ConditionCorsLayer {
+    type Service = ConditionCors<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        ConditionCors {
+            bypass_cors: self.bypass_cors.layer(inner.clone()),
+            default_cors: self.default_cors.layer(inner),
+            bypass_paths: self.bypass_paths.clone(),
+        }
+    }
+}
+
+pub struct ConditionCorsMiddlewarePrepare<T: CorsConfigTrait>(T);
+
+impl<S: Clone, T: CorsConfigTrait + 'static> PrepareMiddlewareEffect<S>
+    for ConditionCorsMiddlewarePrepare<T>
+{
+    type Middleware = ConditionCorsLayer;
 
     fn take(self, _: &mut axum_starter::StateCollector) -> Self::Middleware {
-        self.0
+        ConditionCorsLayer {
+            bypass_cors: CorsLayer::new()
+                .allow_origin(Any)
+                .allow_methods(Any),
+            default_cors: CorsLayer::new()
+                .allow_origin(self.0.allow_origins())
+                .allow_methods(self.0.allow_methods()),
+            bypass_paths: self.0.bypass_paths(),
+        }
     }
 }
+
+#[prepare(ConditionCorsPrepare)]
+pub fn prepare_condition_cors<T: CorsConfigTrait + Clone>(
+    config: &T,
+) -> ConditionCorsMiddlewarePrepare<T> {
+    ConditionCorsMiddlewarePrepare(config.clone())
+}

src/configs/cors_config.rs

@@ -1,3 +1,5 @@
+use std::{collections::HashSet, sync::Arc};
+
 use http::{HeaderValue, Method};
 use serde::{Deserialize, Deserializer};
 
@@ -9,12 +11,18 @@ pub struct CorsConfigImpl {
     allow_origins: Vec<HeaderValue>,
     #[serde(alias = "methods", deserialize_with = "de_methods")]
     allow_methods: Vec<Method>,
+    #[serde(alias = "paths")]
+    pass_paths: Arc<HashSet<String>>,
 }
 
 impl CorsConfigTrait for CorsConfigImpl {
     fn allow_origins(&self) -> Vec<HeaderValue> { self.allow_origins.clone() }
 
     fn allow_methods(&self) -> Vec<Method> { self.allow_methods.clone() }
+
+    fn bypass_paths(&self) -> Arc<HashSet<String>> {
+        self.pass_paths.clone()
+    }
 }
 
 fn de_origins<'de, D: Deserializer<'de>>(

src/main.rs

@@ -14,8 +14,7 @@ use bootstrap::{
         service_init::{graceful_shutdown, RouteV1, RouterFallback},
     },
     middleware::{
-        cors::PrepareCors, panic_report::PrepareCatchPanic,
-        tracing_request::PrepareRequestTracker,
+        cors::ConditionCorsPrepare, panic_report::PrepareCatchPanic, tracing_request::PrepareRequestTracker
     },
 };
 use ceobe_qiniu_upload::QiniuUpload;
@@ -95,7 +94,7 @@ async fn main_task() {
         // router
         .prepare_route(RouteV1)
         .prepare_route(RouterFallback)
-        .prepare_middleware::<Route, _>(PrepareCors::<_, CorsConfigImpl>)
+        .prepare_middleware::<Route, _>(ConditionCorsPrepare::<_, CorsConfigImpl>)
         .prepare_middleware::<Route, _>(
             PrepareCatchPanic::<_, QqChannelConfig>,
         )