1use std::collections::BTreeMap;
18
19use hmac::{Hmac, Mac};
20use sha2::Sha256;
21use sutegi_json::Json;
22use sutegi_web::{Request, Response};
23
24type HmacSha256 = Hmac<Sha256>;
25
26pub struct Sessions {
28 secret: Vec<u8>,
29 cookie: String,
30 secure: bool,
31 max_age: Option<u64>,
32}
33
34impl Sessions {
35 pub fn new(secret: &[u8]) -> Sessions {
37 Sessions {
38 secret: secret.to_vec(),
39 cookie: "sutegi_session".to_string(),
40 secure: true,
41 max_age: Some(86_400),
42 }
43 }
44
45 pub fn cookie_name(mut self, name: &str) -> Sessions {
46 self.cookie = name.to_string();
47 self
48 }
49
50 pub fn insecure(mut self) -> Sessions {
52 self.secure = false;
53 self
54 }
55
56 pub fn max_age(mut self, secs: Option<u64>) -> Sessions {
57 self.max_age = secs;
58 self
59 }
60
61 fn sign(&self, msg: &[u8]) -> String {
62 let mut mac = HmacSha256::new_from_slice(&self.secret).expect("HMAC accepts any key length");
63 mac.update(msg);
64 to_hex(&mac.finalize().into_bytes())
65 }
66
67 pub fn token(&self, value: &str) -> String {
70 let hex = to_hex(value.as_bytes());
71 format!("{}.{}", hex, self.sign(value.as_bytes()))
72 }
73
74 pub fn verify_token(&self, token: &str) -> Option<String> {
76 let (hex, sig) = token.split_once('.')?;
77 let bytes = from_hex(hex)?;
78 if constant_time_eq(self.sign(&bytes).as_bytes(), sig.as_bytes()) {
79 String::from_utf8(bytes).ok()
80 } else {
81 None
82 }
83 }
84
85 pub fn load(&self, req: &Request) -> Session {
88 if let Some(raw) = req.cookie(&self.cookie) {
89 if let Some((payload_hex, sig)) = raw.split_once('.') {
90 if let Some(bytes) = from_hex(payload_hex) {
91 if constant_time_eq(self.sign(&bytes).as_bytes(), sig.as_bytes()) {
92 if let Ok(s) = std::str::from_utf8(&bytes) {
93 if let Ok(Json::Obj(map)) = Json::parse(s) {
94 return Session { data: map, dirty: false };
95 }
96 }
97 }
98 }
99 }
100 }
101 Session { data: BTreeMap::new(), dirty: false }
102 }
103
104 pub fn save(&self, session: &Session, resp: Response) -> Response {
106 let payload = Json::Obj(session.data.clone()).to_string();
107 let payload_hex = to_hex(payload.as_bytes());
108 let sig = self.sign(payload.as_bytes());
109 let mut cookie = format!(
110 "{}={}.{}; Path=/; HttpOnly; SameSite=Lax",
111 self.cookie, payload_hex, sig
112 );
113 if self.secure {
114 cookie.push_str("; Secure");
115 }
116 if let Some(age) = self.max_age {
117 cookie.push_str(&format!("; Max-Age={}", age));
118 }
119 resp.with_header("set-cookie", &cookie)
120 }
121
122 pub fn clear(&self, resp: Response) -> Response {
124 resp.with_header("set-cookie", &format!("{}=; Path=/; Max-Age=0", self.cookie))
125 }
126}
127
128pub struct Session {
130 data: BTreeMap<String, Json>,
131 dirty: bool,
132}
133
134impl Session {
135 pub fn get(&self, key: &str) -> Option<&Json> {
136 self.data.get(key)
137 }
138 pub fn get_str(&self, key: &str) -> Option<&str> {
139 self.data.get(key).and_then(|j| j.as_str())
140 }
141 pub fn set(&mut self, key: &str, value: Json) {
142 self.data.insert(key.to_string(), value);
143 self.dirty = true;
144 }
145 pub fn remove(&mut self, key: &str) {
146 self.data.remove(key);
147 self.dirty = true;
148 }
149 pub fn clear(&mut self) {
150 self.data.clear();
151 self.dirty = true;
152 }
153 pub fn is_empty(&self) -> bool {
154 self.data.is_empty()
155 }
156 pub fn is_dirty(&self) -> bool {
158 self.dirty
159 }
160}
161
162fn to_hex(bytes: &[u8]) -> String {
163 let mut s = String::with_capacity(bytes.len() * 2);
164 for b in bytes {
165 s.push_str(&format!("{:02x}", b));
166 }
167 s
168}
169
170fn from_hex(s: &str) -> Option<Vec<u8>> {
171 if s.len() % 2 != 0 {
172 return None;
173 }
174 (0..s.len())
175 .step_by(2)
176 .map(|i| u8::from_str_radix(&s[i..i + 2], 16).ok())
177 .collect()
178}
179
180fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
181 if a.len() != b.len() {
182 return false;
183 }
184 let mut diff = 0u8;
185 for i in 0..a.len() {
186 diff |= a[i] ^ b[i];
187 }
188 diff == 0
189}
190
191#[cfg(test)]
192mod tests {
193 use super::*;
194
195 fn req_with_cookie(name: &str, value: &str) -> Request {
196 Request {
197 method: sutegi_web::Method::Get,
198 path: "/".into(),
199 query: String::new(),
200 version: "HTTP/1.1".into(),
201 headers: vec![("Cookie".into(), format!("{}={}", name, value))],
202 body: vec![],
203 peer: None,
204 }
205 }
206
207 fn cookie_value(resp: &Response, _name: &str) -> String {
209 let header = resp
210 .headers
211 .iter()
212 .find(|(k, _)| k.eq_ignore_ascii_case("set-cookie"))
213 .map(|(_, v)| v.clone())
214 .unwrap();
215 let kv = header.split(';').next().unwrap();
216 kv.split_once('=').unwrap().1.to_string()
217 }
218
219 #[test]
220 fn roundtrip_and_tamper() {
221 let s = Sessions::new(b"super-secret-key").insecure();
222 let mut sess = s.load(&req_with_cookie("x", "")); assert!(sess.is_empty());
224 sess.set("user_id", Json::int(42));
225
226 let resp = s.save(&sess, Response::new(200));
227 let cookie = cookie_value(&resp, "sutegi_session");
228
229 let reloaded = s.load(&req_with_cookie("sutegi_session", &cookie));
231 assert_eq!(reloaded.get("user_id").and_then(Json::as_i64), Some(42));
232
233 let tampered = cookie.replacen(|c: char| c.is_ascii_hexdigit(), "0", 1);
235 let after = s.load(&req_with_cookie("sutegi_session", &tampered));
236 assert!(after.is_empty() || after.get("user_id").and_then(Json::as_i64) != Some(42));
237 }
238
239 #[test]
240 fn token_sign_verify() {
241 let s = Sessions::new(b"k");
242 let t = s.token("user:42");
243 assert_eq!(s.verify_token(&t).as_deref(), Some("user:42"));
244 assert!(s.verify_token("deadbeef.badsig").is_none());
245 }
246}