Skip to main content

serverkit/
cookie.rs

1use std::fmt;
2
3use crate::Headers;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct Cookie {
7    name: String,
8    value: String,
9    path: Option<String>,
10    domain: Option<String>,
11    max_age: Option<i64>,
12    same_site: Option<SameSite>,
13    http_only: bool,
14    secure: bool,
15}
16
17impl Cookie {
18    pub fn new(name: impl Into<String>, value: impl Into<String>) -> Self {
19        Self {
20            name: name.into(),
21            value: value.into(),
22            path: None,
23            domain: None,
24            max_age: None,
25            same_site: None,
26            http_only: false,
27            secure: false,
28        }
29    }
30
31    pub fn name(&self) -> &str {
32        &self.name
33    }
34
35    pub fn value(&self) -> &str {
36        &self.value
37    }
38
39    pub fn path(mut self, path: impl Into<String>) -> Self {
40        self.path = Some(path.into());
41        self
42    }
43
44    pub fn domain(mut self, domain: impl Into<String>) -> Self {
45        self.domain = Some(domain.into());
46        self
47    }
48
49    pub fn max_age(mut self, seconds: i64) -> Self {
50        self.max_age = Some(seconds);
51        self
52    }
53
54    pub fn same_site(mut self, same_site: SameSite) -> Self {
55        self.same_site = Some(same_site);
56        self
57    }
58
59    pub fn http_only(mut self, enabled: bool) -> Self {
60        self.http_only = enabled;
61        self
62    }
63
64    pub fn secure(mut self, enabled: bool) -> Self {
65        self.secure = enabled;
66        self
67    }
68
69    pub(crate) fn header_value(&self) -> String {
70        let mut value = format!("{}={}", self.name, self.value);
71
72        if let Some(path) = &self.path {
73            value.push_str("; Path=");
74            value.push_str(path);
75        }
76
77        if let Some(domain) = &self.domain {
78            value.push_str("; Domain=");
79            value.push_str(domain);
80        }
81
82        if let Some(max_age) = self.max_age {
83            value.push_str("; Max-Age=");
84            value.push_str(&max_age.to_string());
85        }
86
87        if let Some(same_site) = self.same_site {
88            value.push_str("; SameSite=");
89            value.push_str(same_site.as_str());
90        }
91
92        if self.http_only {
93            value.push_str("; HttpOnly");
94        }
95
96        if self.secure {
97            value.push_str("; Secure");
98        }
99
100        value
101    }
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105pub enum SameSite {
106    Strict,
107    Lax,
108    None,
109}
110
111impl SameSite {
112    fn as_str(self) -> &'static str {
113        match self {
114            Self::Strict => "Strict",
115            Self::Lax => "Lax",
116            Self::None => "None",
117        }
118    }
119}
120
121#[derive(Debug, Default, Clone, PartialEq, Eq)]
122pub struct Cookies {
123    entries: Vec<(String, String)>,
124}
125
126impl Cookies {
127    pub fn get(&self, name: &str) -> Option<&str> {
128        self.entries
129            .iter()
130            .find(|(actual, _)| actual == name)
131            .map(|(_, value)| value.as_str())
132    }
133
134    pub fn get_all<'cookies>(
135        &'cookies self,
136        name: &'cookies str,
137    ) -> impl Iterator<Item = &'cookies str> + 'cookies {
138        self.entries
139            .iter()
140            .filter(move |(actual, _)| actual == name)
141            .map(|(_, value)| value.as_str())
142    }
143
144    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
145        self.entries
146            .iter()
147            .map(|(name, value)| (name.as_str(), value.as_str()))
148    }
149
150    pub fn len(&self) -> usize {
151        self.entries.len()
152    }
153
154    pub fn is_empty(&self) -> bool {
155        self.entries.is_empty()
156    }
157
158    pub(crate) fn from_headers(headers: &Headers) -> Self {
159        let entries = headers
160            .get_all("cookie")
161            .filter_map(|value| std::str::from_utf8(value).ok())
162            .flat_map(|value| value.split(';'))
163            .filter_map(|pair| {
164                let (name, value) = pair.trim().split_once('=')?;
165                (!name.is_empty()).then(|| (name.to_owned(), unquote(value)))
166            })
167            .collect();
168
169        Self { entries }
170    }
171}
172
173fn unquote(value: &str) -> String {
174    value
175        .strip_prefix('"')
176        .and_then(|value| value.strip_suffix('"'))
177        .unwrap_or(value)
178        .replace("\\\"", "\"")
179        .replace("\\\\", "\\")
180}
181
182impl fmt::Display for Cookie {
183    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
184        formatter.write_str(&self.header_value())
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::{Cookie, Cookies, SameSite};
191    use crate::Headers;
192
193    #[test]
194    fn parses_repeated_cookie_headers() {
195        let mut headers = Headers::new();
196        headers.append("Cookie", "session=abc; theme=dark").unwrap();
197        headers.append("cookie", "tag=one").unwrap();
198        let cookies = Cookies::from_headers(&headers);
199
200        assert_eq!(cookies.get("session"), Some("abc"));
201        assert_eq!(cookies.get("theme"), Some("dark"));
202        assert_eq!(cookies.get("tag"), Some("one"));
203    }
204
205    #[test]
206    fn formats_set_cookie_attributes() {
207        let cookie = Cookie::new("session", "abc")
208            .path("/")
209            .max_age(60)
210            .same_site(SameSite::Lax)
211            .http_only(true)
212            .secure(true);
213
214        assert_eq!(
215            cookie.to_string(),
216            "session=abc; Path=/; Max-Age=60; SameSite=Lax; HttpOnly; Secure",
217        );
218    }
219}