1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
use crate::{AxumSession, AxumSessionConfig, AxumSessionData, AxumSessionID, AxumSessionStore};
use axum_core::response::IntoResponse;
use axum_extra::middleware::Next;
use chrono::{Duration, Utc};
use http::{self, Request, StatusCode};
use std::collections::HashMap;
use tokio::sync::Mutex;
use tower_cookies::{Cookie, Cookies};
use uuid::Uuid;
///axum_session_manager_run Creates, Manages and Sets an AxumSession into the Request extensions.
///This will unload and load other Session data based on Access and a timer check.
///returns an Response when all the Futures After run.
pub async fn axum_session_manager_run<B>(
mut req: Request<B>,
next: Next<B>,
store: AxumSessionStore,
) -> impl IntoResponse {
let config = store.config.clone();
// We Extract the Tower_Cookies Extensions Variable so we can add Cookies to it. Some reason can only be done here..?
let cookies = match req.extensions().get::<Cookies>() {
Some(cookies) => cookies,
None => return Err(StatusCode::UNAUTHORIZED),
};
let session = AxumSession {
id: {
let id = if let Some(cookie) = cookies.get(&store.config.cookie_name) {
(
AxumSessionID(Uuid::parse_str(cookie.value()).expect("`Could not parse Uuid")),
true,
)
} else {
let store_ug = store.inner.read().await;
let new_id = loop {
let token = Uuid::new_v4();
if !store_ug.contains_key(&token.to_string()) {
break token;
}
};
(AxumSessionID(new_id), false)
};
// If a cookie did have an AxumSessionID then lets check if it still exists in the hash or Database
// If not make a new Session using the ID.
if id.1 {
let mut no_store = true;
{
if let Some(m) = store.inner.read().await.get(&id.0.to_string()) {
let mut inner = m.lock().await;
if inner.expires < Utc::now() || inner.destroy {
// Database Session expired, reuse the ID but drop data.
inner.data = HashMap::new();
}
// Session is extended by making a request with valid ID
inner.expires = Utc::now() + store.config.lifespan;
inner.autoremove = Utc::now() + store.config.memory_lifespan;
no_store = false;
}
}
if no_store {
let mut sess = store
.load_session(id.0.to_string())
.await
.ok()
.flatten()
.unwrap_or(AxumSessionData {
id: id.0 .0,
data: HashMap::new(),
expires: Utc::now() + Duration::hours(6),
destroy: false,
autoremove: Utc::now() + store.config.memory_lifespan,
});
if !sess.validate() || sess.destroy {
sess.data = HashMap::new();
sess.expires = Utc::now() + Duration::hours(6);
sess.autoremove = Utc::now() + store.config.memory_lifespan;
}
cookies.add(create_cookie(config, id.0 .0.to_string()));
store
.inner
.write()
.await
.insert(id.0 .0.to_string(), Mutex::new(sess));
}
} else {
// --- New ID was generated Lets make a session for it ---
cookies.add(create_cookie(config, id.0 .0.to_string()));
let sess = AxumSessionData {
id: id.0 .0,
data: HashMap::new(),
expires: Utc::now() + Duration::hours(6),
destroy: false,
autoremove: Utc::now() + store.config.memory_lifespan,
};
store
.inner
.write()
.await
.insert(id.0 .0.to_string(), Mutex::new(sess));
}
id.0
},
store: store.clone(),
};
let (last_sweep, last_database_sweep) = {
let timers = store.timers.read().await;
(timers.last_expiry_sweep, timers.last_database_expiry_sweep)
};
// This branch runs less often, and we already have write access,
// let's check if any sessions expired. We don't want to hog memory
// forever by abandoned sessions (e.g. when a client lost their cookie)
// Throttle by memory lifespan - e.g. sweep every hour
if last_sweep <= Utc::now() {
store.inner.write().await.retain(|_k, v| {
if let Ok(data) = v.try_lock() {
data.autoremove > Utc::now()
} else {
//the lock is busy so rather than just killing
//everything lets ignore it till next time.
true
}
});
store.timers.write().await.last_expiry_sweep = Utc::now() + store.config.memory_lifespan;
}
// Throttle by database lifespan - e.g. sweep every 6 hours
if last_database_sweep <= Utc::now() && store.is_persistent() {
store.cleanup().await.unwrap();
store.timers.write().await.last_database_expiry_sweep = Utc::now() + store.config.lifespan;
}
//Sets a clone of the Store in the Extensions for Direct usage and sets the Session for Direct usage
req.extensions_mut().insert(store.clone());
req.extensions_mut().insert(session.clone());
let response = next.run(req).await;
//run this After a response has returned so we save the most updated data to sql.
if store.is_persistent() {
if let Some(session_data) = session
.store
.inner
.read()
.await
.get(&session.id.0.to_string())
{
session
.store
.store_session(session_data.lock().await.clone())
.await
.unwrap()
}
}
Ok(response)
}
fn create_cookie<'a>(config: AxumSessionConfig, value: String) -> Cookie<'a> {
let mut cookie_builder = Cookie::build(config.cookie_name, value)
.path(config.cookie_path)
.secure(config.cookie_secure)
.http_only(config.cookie_http_only);
if let Some(domain) = config.cookie_domain {
cookie_builder = cookie_builder
.domain(domain)
.same_site(config.cookie_same_site);
}
if let Some(max_age) = config.cookie_max_age {
let time_duration = max_age.to_std().expect("Max Age out of bounds");
cookie_builder =
cookie_builder.max_age(time_duration.try_into().expect("Max Age out of bounds"));
}
cookie_builder.finish()
}