Skip to main content

xyz_rust/httpapi/
middleware.rs

1// 中间件积木(Go middleware.go 对应物):Bearer 校验、CORS 白名单、
2// Gzip 压缩。各自独立可组合(docs/adapters.md 档位 C)。
3//
4// 层顺序语义(tower:后 add 的层在外侧、先触发,等价 Go 的由外到内):
5//   CORS(预检在鉴权前)→ Bearer → Gzip → 路由。
6// 该顺序由 middleware::apply 组装,并有测试锁定。
7
8use std::collections::HashSet;
9use std::sync::Arc;
10use std::time::Duration;
11
12use axum::extract::Request;
13use axum::http::{Method, StatusCode, header};
14use axum::middleware::Next;
15use axum::response::{IntoResponse, Response};
16use axum::routing::Router;
17
18fn write_error(status: StatusCode, msg: &str) -> Response {
19    let body = serde_json::to_string(&serde_json::json!({ "error": msg }))
20        .unwrap_or_else(|_| "{\"error\":\"\"}".to_string());
21    let mut b = body;
22    b.push('\n');
23    (
24        status,
25        [(header::CONTENT_TYPE, "application/json; charset=utf-8")],
26        b,
27    )
28        .into_response()
29}
30
31/// Bearer 校验:Authorization: Bearer <tok> 必须命中令牌集之一,否则
32/// 401 + {"error":"unauthorized"} + WWW-Authenticate: Bearer。
33pub fn bearer_mw(
34    tokens: Arc<HashSet<String>>,
35) -> impl Fn(
36    Request,
37    Next,
38) -> std::pin::Pin<Box<dyn std::future::Future<Output = Response> + Send + 'static>>
39+ Clone
40+ Send
41+ 'static {
42    move |req: Request, next: Next| {
43        let tokens = Arc::clone(&tokens);
44        Box::pin(async move {
45            let auth = req
46                .headers()
47                .get(header::AUTHORIZATION)
48                .and_then(|v| v.to_str().ok())
49                .unwrap_or("");
50            // Go:prefix(Bearer + 空格)区分大小写,token 精确匹配。
51            let ok = auth.starts_with("Bearer ")
52                && tokens.contains(auth.strip_prefix("Bearer ").unwrap_or(""));
53            if !ok {
54                let mut resp = write_error(StatusCode::UNAUTHORIZED, "unauthorized");
55                let _ = resp.headers_mut().insert(
56                    header::WWW_AUTHENTICATE,
57                    header::HeaderValue::from_static("Bearer"),
58                );
59                return resp;
60            }
61            next.run(req).await
62        })
63    }
64}
65
66/// CORS 白名单:Origin 命中白名单(或 "*")时回写
67/// Access-Control-Allow-Origin;OPTIONS 预检在鉴权之前应答 204(浏览器
68/// 预检不带凭据)。
69pub fn cors_mw(
70    origins: Arc<HashSet<String>>,
71) -> impl Fn(
72    Request,
73    Next,
74) -> std::pin::Pin<Box<dyn std::future::Future<Output = Response> + Send + 'static>>
75+ Clone
76+ Send
77+ 'static {
78    move |req: Request, next: Next| {
79        let origins = Arc::clone(&origins);
80        Box::pin(async move {
81            let mut acao: Option<(String, bool)> = None;
82            if let Some(origin) = req
83                .headers()
84                .get(header::ORIGIN)
85                .and_then(|v| v.to_str().ok())
86            {
87                if origins.contains("*") {
88                    acao = Some(("*".to_string(), false));
89                } else if origins.contains(origin) {
90                    acao = Some((origin.to_string(), true));
91                }
92            }
93            // 预检必须回写 Access-Control-Allow-Origin(浏览器缺它即拒绝),
94            // 且应答发生在鉴权之前:浏览器的 OPTIONS 不带 Authorization。
95            if req.method() == Method::OPTIONS
96                && let Some((origin, vary)) = acao.take()
97            {
98                let mut resp = Response::new(axum::body::Body::empty());
99                *resp.status_mut() = StatusCode::NO_CONTENT;
100                let h = resp.headers_mut();
101                let _ = h.insert(
102                    header::ACCESS_CONTROL_ALLOW_ORIGIN,
103                    header::HeaderValue::from_str(&origin)
104                        .unwrap_or(header::HeaderValue::from_static("*")),
105                );
106                if vary {
107                    let _ = h.insert(header::VARY, header::HeaderValue::from_static("Origin"));
108                }
109                h.insert(
110                    header::ACCESS_CONTROL_ALLOW_METHODS,
111                    header::HeaderValue::from_static("GET, POST, PUT, PATCH, DELETE, OPTIONS"),
112                );
113                h.insert(
114                    header::ACCESS_CONTROL_ALLOW_HEADERS,
115                    header::HeaderValue::from_static("Content-Type, Authorization"),
116                );
117                h.insert(
118                    header::ACCESS_CONTROL_MAX_AGE,
119                    header::HeaderValue::from_static("86400"),
120                );
121                return resp;
122            }
123            let mut resp = next.run(req).await;
124            if let Some((origin, vary)) = acao {
125                let h = resp.headers_mut();
126                let _ = h.insert(
127                    header::ACCESS_CONTROL_ALLOW_ORIGIN,
128                    header::HeaderValue::from_str(&origin)
129                        .unwrap_or(header::HeaderValue::from_static("*")),
130                );
131                if vary {
132                    let existing = h
133                        .get(header::VARY)
134                        .and_then(|v| v.to_str().ok())
135                        .unwrap_or("")
136                        .to_string();
137                    let joined = if existing.is_empty() {
138                        "Origin".to_string()
139                    } else {
140                        format!("{existing}, Origin")
141                    };
142                    let _ = h.insert(
143                        header::VARY,
144                        header::HeaderValue::from_str(&joined)
145                            .unwrap_or(header::HeaderValue::from_static("Origin")),
146                    );
147                }
148            }
149            resp
150        })
151    }
152}
153
154/// 组装完整中间件链(由外到内):CORS → Bearer → Gzip → 路由。
155/// tower 的 layer 语义是「先 add 的靠内」,所以按 内→外 顺序加。
156/// 服务器上下文头(spec §11.6):X-App-Name / X-App-Version /
157/// X-XYZ-Version 三个身份头(enabled 时)+ 用户自定义静态头(始终写)。
158/// 覆盖全部路由(含 /healthz、/openapi.json 与挂载的 /mcp);命令/耗时头
159/// 是每路由的,在 EntryHandler 里写。
160pub fn server_context_mw(
161    app_name: String,
162    app_version: String,
163    custom: Vec<(String, String)>,
164    enabled: bool,
165) -> impl Fn(
166    Request,
167    Next,
168) -> std::pin::Pin<Box<dyn std::future::Future<Output = Response> + Send + 'static>>
169+ Clone
170+ Send
171+ 'static {
172    move |req: Request, next: Next| {
173        let app_name = app_name.clone();
174        let app_version = app_version.clone();
175        let custom = custom.clone();
176        Box::pin(async move {
177            let mut resp = next.run(req).await;
178            let h = resp.headers_mut();
179            if enabled {
180                for (name, value) in [
181                    ("x-app-name", app_name.as_str()),
182                    ("x-app-version", app_version.as_str()),
183                    ("x-xyz-version", crate::version::SDK_VERSION),
184                ] {
185                    if let Ok(v) = header::HeaderValue::from_str(value) {
186                        let _ = h.insert(header::HeaderName::from_static(name), v);
187                    }
188                }
189            }
190            for (k, v) in &custom {
191                if let (Ok(name), Ok(value)) = (
192                    header::HeaderName::from_bytes(k.as_bytes()),
193                    header::HeaderValue::from_str(v),
194                ) {
195                    let _ = h.insert(name, value);
196                }
197            }
198            resp
199        })
200    }
201}
202
203pub fn apply(
204    router: Router,
205    bearer_tokens: Vec<String>,
206    cors_origins: Vec<String>,
207    timeout: Duration,
208) -> Router {
209    let bearer: Arc<HashSet<String>> = Arc::new(bearer_tokens.into_iter().collect());
210    let cors: Arc<HashSet<String>> = Arc::new(cors_origins.into_iter().collect());
211    let mut r = router;
212    // 最内:Go 对任意大小响应都压缩(tower-http 默认谓词 SizeAbove(32),
213    // 对齐 Go 用 SizeAbove(0))。
214    r = r.layer(
215        tower_http::compression::CompressionLayer::new()
216            .compress_when(tower_http::compression::predicate::SizeAbove::new(0)),
217    );
218    if !bearer.is_empty() {
219        r = r.layer(axum::middleware::from_fn(bearer_mw(Arc::clone(&bearer))));
220    }
221    if !cors.is_empty() {
222        r = r.layer(axum::middleware::from_fn(cors_mw(Arc::clone(&cors))));
223    }
224    if timeout > Duration::ZERO {
225        r = r.layer(tower_http::timeout::TimeoutLayer::with_status_code(
226            StatusCode::REQUEST_TIMEOUT,
227            timeout,
228        ));
229    }
230    r
231}