sa_token_plugin_axum/
adapter.rs

1// Author: 金书记
2//
3//! Axum请求/响应适配器
4
5use std::collections::HashMap;
6use http::{Request, Response};
7use sa_token_adapter::context::{SaRequest, SaResponse, CookieOptions};
8use serde::Serialize;
9
10/// Axum请求适配器
11pub struct AxumRequestAdapter<'a, T> {
12    request: &'a Request<T>,
13}
14
15impl<'a, T> AxumRequestAdapter<'a, T> {
16    pub fn new(request: &'a Request<T>) -> Self {
17        Self { request }
18    }
19}
20
21impl<'a, T> SaRequest for AxumRequestAdapter<'a, T> {
22    fn get_header(&self, name: &str) -> Option<String> {
23        self.request.headers().get(name)
24            .and_then(|v| v.to_str().ok())
25            .map(|s| s.to_string())
26    }
27    
28    fn get_cookie(&self, name: &str) -> Option<String> {
29        self.request.headers().get("cookie")
30            .and_then(|v| v.to_str().ok())
31            .and_then(|cookies| parse_cookies(cookies).get(name).cloned())
32    }
33    
34    fn get_param(&self, name: &str) -> Option<String> {
35        self.request
36            .uri()
37            .query()
38            .and_then(|query| parse_query_string(query).get(name).cloned())
39    }
40    
41    fn get_path(&self) -> String {
42        self.request.uri().path().to_string()
43    }
44    
45    fn get_method(&self) -> String {
46        self.request.method().to_string()
47    }
48}
49
50/// Axum响应适配器
51pub struct AxumResponseAdapter<T> {
52    response: Response<T>,
53}
54
55impl<T> AxumResponseAdapter<T> {
56    pub fn new(response: Response<T>) -> Self {
57        Self { response }
58    }
59    
60    pub fn into_response(self) -> Response<T> {
61        self.response
62    }
63}
64
65impl<T> SaResponse for AxumResponseAdapter<T> {
66    fn set_header(&mut self, name: &str, value: &str) {
67        if let Ok(header_name) = http::header::HeaderName::from_bytes(name.as_bytes()) {
68            if let Ok(header_value) = http::header::HeaderValue::from_str(value) {
69                self.response.headers_mut().insert(header_name, header_value);
70            }
71        }
72    }
73    
74    fn set_cookie(&mut self, name: &str, value: &str, options: CookieOptions) {
75        let mut cookie = format!("{}={}", name, value);
76        
77        if let Some(domain) = options.domain {
78            cookie.push_str(&format!("; Domain={}", domain));
79        }
80        if let Some(path) = options.path {
81            cookie.push_str(&format!("; Path={}", path));
82        }
83        if let Some(max_age) = options.max_age {
84            cookie.push_str(&format!("; Max-Age={}", max_age));
85        }
86        if options.http_only {
87            cookie.push_str("; HttpOnly");
88        }
89        if options.secure {
90            cookie.push_str("; Secure");
91        }
92        
93        self.set_header("Set-Cookie", &cookie);
94    }
95    
96    fn set_status(&mut self, status: u16) {
97        *self.response.status_mut() = http::StatusCode::from_u16(status).unwrap_or(http::StatusCode::OK);
98    }
99    
100    fn set_json_body<U: Serialize>(&mut self, _body: U) -> Result<(), serde_json::Error> {
101        // 需要替换response body,在axum中比较复杂
102        Ok(())
103    }
104}
105
106/// 解析 Cookie 字符串
107fn parse_cookies(cookie_header: &str) -> HashMap<String, String> {
108    let mut cookies = HashMap::new();
109    for pair in cookie_header.split(';') {
110        let parts: Vec<&str> = pair.trim().splitn(2, '=').collect();
111        if parts.len() == 2 {
112            cookies.insert(parts[0].to_string(), parts[1].to_string());
113        }
114    }
115    cookies
116}
117
118/// 解析查询字符串
119fn parse_query_string(query: &str) -> HashMap<String, String> {
120    let mut params = HashMap::new();
121    for pair in query.split('&') {
122        let parts: Vec<&str> = pair.splitn(2, '=').collect();
123        if parts.len() == 2 {
124            params.insert(
125                urlencoding::decode(parts[0]).unwrap_or_default().to_string(),
126                urlencoding::decode(parts[1]).unwrap_or_default().to_string(),
127            );
128        }
129    }
130    params
131}