Skip to main content

next_rs_middleware/
response.rs

1use std::collections::HashMap;
2
3#[derive(Debug, Clone)]
4pub enum MiddlewareResult {
5    Next,
6    Rewrite(String),
7    Redirect(RedirectResponse),
8    Response(NextResponse),
9}
10
11#[derive(Debug, Clone)]
12pub struct RedirectResponse {
13    pub url: String,
14    pub status: u16,
15    pub headers: HashMap<String, String>,
16}
17
18impl RedirectResponse {
19    pub fn temporary(url: impl Into<String>) -> Self {
20        Self {
21            url: url.into(),
22            status: 307,
23            headers: HashMap::new(),
24        }
25    }
26
27    pub fn permanent(url: impl Into<String>) -> Self {
28        Self {
29            url: url.into(),
30            status: 308,
31            headers: HashMap::new(),
32        }
33    }
34
35    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
36        self.headers.insert(key.into(), value.into());
37        self
38    }
39}
40
41#[derive(Debug, Clone)]
42pub struct NextResponse {
43    pub status: u16,
44    pub headers: HashMap<String, String>,
45    pub cookies: Vec<SetCookie>,
46    pub body: Option<String>,
47}
48
49#[derive(Debug, Clone)]
50pub struct SetCookie {
51    pub name: String,
52    pub value: String,
53    pub path: Option<String>,
54    pub domain: Option<String>,
55    pub max_age: Option<i64>,
56    pub http_only: bool,
57    pub secure: bool,
58    pub same_site: Option<SameSite>,
59}
60
61#[derive(Debug, Clone, Copy)]
62pub enum SameSite {
63    Strict,
64    Lax,
65    None,
66}
67
68impl NextResponse {
69    pub fn next() -> MiddlewareResult {
70        MiddlewareResult::Next
71    }
72
73    pub fn redirect(url: impl Into<String>) -> MiddlewareResult {
74        MiddlewareResult::Redirect(RedirectResponse::temporary(url))
75    }
76
77    pub fn redirect_permanent(url: impl Into<String>) -> MiddlewareResult {
78        MiddlewareResult::Redirect(RedirectResponse::permanent(url))
79    }
80
81    pub fn rewrite(url: impl Into<String>) -> MiddlewareResult {
82        MiddlewareResult::Rewrite(url.into())
83    }
84
85    pub fn new(status: u16) -> Self {
86        Self {
87            status,
88            headers: HashMap::new(),
89            cookies: Vec::new(),
90            body: None,
91        }
92    }
93
94    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
95        self.headers.insert(key.into(), value.into());
96        self
97    }
98
99    pub fn with_body(mut self, body: impl Into<String>) -> Self {
100        self.body = Some(body.into());
101        self
102    }
103
104    pub fn set_cookie(mut self, cookie: SetCookie) -> Self {
105        self.cookies.push(cookie);
106        self
107    }
108
109    pub fn into_result(self) -> MiddlewareResult {
110        MiddlewareResult::Response(self)
111    }
112}
113
114impl SetCookie {
115    pub fn new(name: impl Into<String>, value: impl Into<String>) -> Self {
116        Self {
117            name: name.into(),
118            value: value.into(),
119            path: None,
120            domain: None,
121            max_age: None,
122            http_only: false,
123            secure: false,
124            same_site: None,
125        }
126    }
127
128    pub fn with_path(mut self, path: impl Into<String>) -> Self {
129        self.path = Some(path.into());
130        self
131    }
132
133    pub fn with_max_age(mut self, seconds: i64) -> Self {
134        self.max_age = Some(seconds);
135        self
136    }
137
138    pub fn http_only(mut self) -> Self {
139        self.http_only = true;
140        self
141    }
142
143    pub fn secure(mut self) -> Self {
144        self.secure = true;
145        self
146    }
147
148    pub fn with_same_site(mut self, same_site: SameSite) -> Self {
149        self.same_site = Some(same_site);
150        self
151    }
152
153    pub fn to_header_value(&self) -> String {
154        let mut parts = vec![format!("{}={}", self.name, self.value)];
155
156        if let Some(path) = &self.path {
157            parts.push(format!("Path={}", path));
158        }
159        if let Some(domain) = &self.domain {
160            parts.push(format!("Domain={}", domain));
161        }
162        if let Some(max_age) = self.max_age {
163            parts.push(format!("Max-Age={}", max_age));
164        }
165        if self.http_only {
166            parts.push("HttpOnly".to_string());
167        }
168        if self.secure {
169            parts.push("Secure".to_string());
170        }
171        if let Some(same_site) = &self.same_site {
172            parts.push(format!(
173                "SameSite={}",
174                match same_site {
175                    SameSite::Strict => "Strict",
176                    SameSite::Lax => "Lax",
177                    SameSite::None => "None",
178                }
179            ));
180        }
181
182        parts.join("; ")
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn test_next_response_redirect() {
192        let result = NextResponse::redirect("/login");
193        if let MiddlewareResult::Redirect(redirect) = result {
194            assert_eq!(redirect.url, "/login");
195            assert_eq!(redirect.status, 307);
196        } else {
197            panic!("Expected Redirect");
198        }
199    }
200
201    #[test]
202    fn test_next_response_rewrite() {
203        let result = NextResponse::rewrite("/api/v2/users");
204        if let MiddlewareResult::Rewrite(url) = result {
205            assert_eq!(url, "/api/v2/users");
206        } else {
207            panic!("Expected Rewrite");
208        }
209    }
210
211    #[test]
212    fn test_response_with_headers() {
213        let response = NextResponse::new(200)
214            .with_header("X-Custom", "value")
215            .with_body("Hello");
216
217        assert_eq!(response.status, 200);
218        assert_eq!(response.headers.get("X-Custom"), Some(&"value".to_string()));
219        assert_eq!(response.body, Some("Hello".to_string()));
220    }
221
222    #[test]
223    fn test_set_cookie() {
224        let cookie = SetCookie::new("session", "abc123")
225            .with_path("/")
226            .with_max_age(3600)
227            .http_only()
228            .secure()
229            .with_same_site(SameSite::Strict);
230
231        let header = cookie.to_header_value();
232        assert!(header.contains("session=abc123"));
233        assert!(header.contains("Path=/"));
234        assert!(header.contains("Max-Age=3600"));
235        assert!(header.contains("HttpOnly"));
236        assert!(header.contains("Secure"));
237        assert!(header.contains("SameSite=Strict"));
238    }
239
240    #[test]
241    fn test_response_with_cookie() {
242        let response = NextResponse::new(200).set_cookie(SetCookie::new("token", "xyz"));
243
244        assert_eq!(response.cookies.len(), 1);
245        assert_eq!(response.cookies[0].name, "token");
246    }
247}