feat: proxy and admin cors (#94)

Author: bzp2010

Cargo.lock

@@ -137,6 +137,7 @@ axum-server = { version = "0.8.0", default-features = false, features = [
 backon = { version = "1.6.0", default-features = false, features = [
   "tokio-sleep",
 ] }
+tower-http = { version = "0.6.10", features = ["cors"] }
 
 [build-dependencies]
 vergen-git2 = { version = "9.1.0" }

Cargo.toml

@@ -20,5 +20,16 @@ server:
       enabled: false
       cert_file: cert.pem
       key_file: key.pem
+    cors:
+      enabled: false
+      allowed_origins: [ "*" ]
+      allowed_methods: [ "GET", "POST" ]
+      allowed_headers: [ "Authorization", "Content-Type" ]
+      exposed_headers: []
+      allow_credentials: false
   admin:
     listen: 127.0.0.1:3001
+    tls:
+      enabled: false
+    cors:
+      enabled: false

config.yaml

@@ -6,6 +6,7 @@ mod types;
 
 use std::sync::Arc;
 
+use anyhow::Result;
 use axum::{
     Router,
     extract::{Request, State},
@@ -107,8 +108,8 @@ impl AppState {
     }
 }
 
-pub fn create_router(state: AppState) -> Router {
-    Router::new()
+pub fn create_router(state: AppState) -> Result<Router> {
+    let mut router = Router::new()
         .nest(
             PATH_PREFIX,
             Router::new()
@@ -148,8 +149,14 @@ pub fn create_router(state: AppState) -> Router {
         .route("/ui", get(|| async { Redirect::to("/ui/") }))
         .route("/ui/", get(aisix_admin_ui::handler))
         .route("/ui/{*path}", get(aisix_admin_ui::handler))
-        .merge(Scalar::with_url("/openapi", ApiDoc::openapi()))
-        .with_state(state)
+        .merge(Scalar::with_url("/openapi", ApiDoc::openapi()));
+
+    let cors = &state.config.server.admin.cors;
+    if cors.enabled {
+        router = router.layer(cors.to_cors_layer()?)
+    };
+
+    Ok(router.with_state(state))
 }
 
 async fn auth(

src/admin/mod.rs

@@ -1,9 +1,11 @@
-use std::net::SocketAddr;
+use std::{net::SocketAddr, str::FromStr};
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use async_trait::async_trait;
+use http::{HeaderValue, Method, header::HeaderName};
 use serde::Deserialize;
 use tokio::sync::mpsc;
+use tower_http::cors::{AllowOrigin, CorsLayer};
 
 use crate::config::etcd;
 
@@ -46,19 +48,88 @@ pub struct ServerCommonTls {
     pub key_file: Option<String>,
 }
 
+#[derive(Clone, Debug, Default, Deserialize)]
+pub struct ServerCommonCors {
+    #[serde(default)]
+    pub enabled: bool,
+    pub allowed_origins: Option<Vec<String>>,
+    pub allowed_methods: Option<Vec<String>>,
+    pub allowed_headers: Option<Vec<String>>,
+    pub exposed_headers: Option<Vec<String>>,
+    pub allow_credentials: Option<bool>,
+}
+
+impl ServerCommonCors {
+    pub fn to_cors_layer(&self) -> Result<CorsLayer> {
+        let mut cors = CorsLayer::new().allow_credentials(self.allow_credentials.unwrap_or(false));
+
+        if let Some(origins) = self.allowed_origins.as_deref() {
+            cors = cors.allow_origin(if origins.iter().any(|o| o == "*") {
+                AllowOrigin::any()
+            } else {
+                AllowOrigin::list(Self::parse_cors_values(
+                    "allowed_origin",
+                    origins,
+                    HeaderValue::from_str,
+                )?)
+            });
+        }
+
+        if let Some(methods) = self.allowed_methods.as_deref() {
+            cors = cors.allow_methods(Self::parse_cors_values(
+                "allowed_method",
+                methods,
+                Method::from_str,
+            )?);
+        }
+
+        if let Some(headers) = self.allowed_headers.as_deref() {
+            cors = cors.allow_headers(Self::parse_cors_values(
+                "allowed_header",
+                headers,
+                HeaderName::from_str,
+            )?);
+        }
+
+        if let Some(headers) = self.exposed_headers.as_deref() {
+            cors = cors.expose_headers(Self::parse_cors_values(
+                "exposed_header",
+                headers,
+                HeaderName::from_str,
+            )?);
+        }
+
+        Ok(cors)
+    }
+
+    fn parse_cors_values<T, E, F>(field: &str, values: &[String], mut parse: F) -> Result<Vec<T>>
+    where
+        F: FnMut(&str) -> std::result::Result<T, E>,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        values
+            .iter()
+            .map(|value| parse(value).with_context(|| format!("Invalid CORS {}: {}", field, value)))
+            .collect()
+    }
+}
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct ServerProxy {
     #[serde(default = "defaults::listen")]
     pub listen: SocketAddr,
     #[serde(default)]
     pub tls: ServerCommonTls,
+    #[serde(default)]
+    pub cors: ServerCommonCors,
 }
 
 impl Default for ServerProxy {
     fn default() -> Self {
         Self {
             listen: defaults::listen(),
             tls: ServerCommonTls::default(),
+            cors: ServerCommonCors::default(),
         }
     }
 }
@@ -69,13 +140,16 @@ pub struct ServerAdmin {
     pub listen: SocketAddr,
     #[serde(default)]
     pub tls: ServerCommonTls,
+    #[serde(default)]
+    pub cors: ServerCommonCors,
 }
 
 impl Default for ServerAdmin {
     fn default() -> Self {
         Self {
             listen: defaults::admin_listen(),
             tls: ServerCommonTls::default(),
+            cors: ServerCommonCors::default(),
         }
     }
 }
@@ -212,3 +286,40 @@ impl dyn ConfigProvider {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::ServerCommonCors;
+
+    #[test]
+    fn to_cors_layer_accepts_valid_config() {
+        let cors = ServerCommonCors {
+            enabled: true,
+            allowed_origins: Some(vec!["https://example.com".into()]),
+            allowed_methods: Some(vec!["GET".into(), "POST".into()]),
+            allowed_headers: Some(vec!["content-type".into()]),
+            exposed_headers: Some(vec!["x-request-id".into()]),
+            allow_credentials: Some(true),
+        };
+
+        assert!(cors.to_cors_layer().is_ok());
+    }
+
+    #[test]
+    fn to_cors_layer_rejects_invalid_config() {
+        let cors = ServerCommonCors {
+            allowed_methods: Some(vec!["NOT A METHOD".into()]),
+            ..Default::default()
+        };
+
+        let result = cors.to_cors_layer();
+
+        assert!(result.is_err());
+        assert!(
+            result
+                .err()
+                .map(|err| err.to_string().contains("Invalid CORS allowed_method"))
+                .unwrap_or(false)
+        );
+    }
+}

src/config/types.rs

@@ -99,7 +99,8 @@ pub async fn run_with_provider(
         resources.clone(),
         gateway,
         message_history_storage,
-    ));
+    ))
+    .context("failed to create proxy router")?;
 
     let res = select! {
         res = tokio::signal::ctrl_c() =>
@@ -143,7 +144,7 @@ async fn serve_admin(config: Arc<config::Config>, state: admin::AppState) -> Res
         "Admin",
         config.server.admin.listen,
         &config.server.admin.tls,
-        admin::create_router(state),
+        admin::create_router(state).context("failed to create admin router")?,
     )
     .await
 }

src/lib.rs

@@ -7,6 +7,7 @@ mod utils;
 
 use std::sync::Arc;
 
+use anyhow::Result;
 use axum::{
     Router,
     extract::DefaultBodyLimit,
@@ -57,8 +58,8 @@ impl AppState {
     }
 }
 
-pub fn create_router(state: AppState) -> Router {
-    Router::new()
+pub fn create_router(state: AppState) -> Result<Router> {
+    let mut router = Router::new()
         .merge(Router::new().route("/v1/models", get(handlers::models::list_models)))
         .route(
             "/v1/chat/completions",
@@ -80,6 +81,12 @@ pub fn create_router(state: AppState) -> Router {
         .route("/v1/embeddings", post(handlers::embeddings::embeddings))
         .layer(DefaultBodyLimit::max(10 * 1024 * 1024))
         .layer(from_fn_with_state(state.clone(), middlewares::auth))
-        .layer(from_fn(middlewares::trace))
-        .with_state(state)
+        .layer(from_fn(middlewares::trace));
+
+    let cors = &state.config.server.proxy.cors;
+    if cors.enabled {
+        router = router.layer(cors.to_cors_layer()?)
+    };
+
+    Ok(router.with_state(state))
 }