8923bc0
use std::sync::LazyLock;
use actix_web::cookie::time::Duration;
use cache::redis::{Script, ScriptInvocation};
use cache::{Cache, CacheConnection};
use config::redis::RedisConfig;
use crate::error::SessionError;
use crate::key::SessionKey;
use crate::store::{SessionDiff, SessionState, SessionStore, generate_session_key};
const KEY_PREFIX: &str = "session";
const USER_BUCKET_PREFIX: &str = "user:session";
const SAVE_ATTEMPTS: usize = 3;
static LOAD_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
local key_type = redis.call('TYPE', KEYS[1])['ok']
if key_type == 'none' then
return {}
end
if key_type ~= 'hash' then
redis.call('DEL', KEYS[1])
return {}
end
local ttl = redis.call('TTL', KEYS[1])
local values = redis.call('HGETALL', KEYS[1])
table.insert(values, 1, tostring(ttl))
return values
"#,
)
});
static SAVE_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
if redis.call('EXISTS', KEYS[1]) == 1 then
return 0
end
for index = 5, #ARGV, 2 do
redis.call('HSET', KEYS[1], ARGV[index], ARGV[index + 1])
end
if redis.call('EXPIRE', KEYS[1], ARGV[1]) == 0 then
return -1
end
if ARGV[4] == '1' then
redis.call('SADD', KEYS[2], ARGV[3])
redis.call('EXPIRE', KEYS[2], ARGV[2])
end
return 1
"#,
)
});
static UPDATE_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
local key_type = redis.call('TYPE', KEYS[1])['ok']
if key_type ~= 'hash' then
if key_type ~= 'none' then
redis.call('DEL', KEYS[1])
end
return 0
end
local index = 7
local set_count = tonumber(ARGV[6])
for _ = 1, set_count do
redis.call('HSET', KEYS[1], ARGV[index], ARGV[index + 1])
index = index + 2
end
local removed_count = tonumber(ARGV[index])
index = index + 1
for _ = 1, removed_count do
redis.call('HDEL', KEYS[1], ARGV[index])
index = index + 1
end
redis.call('EXPIRE', KEYS[1], ARGV[1])
if ARGV[4] == '1' then
redis.call('SREM', KEYS[2], ARGV[3])
end
if ARGV[5] == '1' then
redis.call('SADD', KEYS[3], ARGV[3])
redis.call('EXPIRE', KEYS[3], ARGV[2])
end
return 1
"#,
)
});
static ROTATE_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
local key_type = redis.call('TYPE', KEYS[1])['ok']
if key_type ~= 'hash' then
if key_type ~= 'none' then
redis.call('DEL', KEYS[1])
end
return 0
end
if redis.call('RENAMENX', KEYS[1], KEYS[2]) == 0 then
return -1
end
local index = 8
local set_count = tonumber(ARGV[7])
for _ = 1, set_count do
redis.call('HSET', KEYS[2], ARGV[index], ARGV[index + 1])
index = index + 2
end
local removed_count = tonumber(ARGV[index])
index = index + 1
for _ = 1, removed_count do
redis.call('HDEL', KEYS[2], ARGV[index])
index = index + 1
end
if redis.call('EXPIRE', KEYS[2], ARGV[1]) == 0 then
return 2
end
if ARGV[5] == '1' then
redis.call('SREM', KEYS[3], ARGV[3])
end
if ARGV[6] == '1' then
redis.call('SADD', KEYS[4], ARGV[4])
redis.call('EXPIRE', KEYS[4], ARGV[2])
end
return 1
"#,
)
});
static UPDATE_TTL_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
if redis.call('TYPE', KEYS[1])['ok'] ~= 'hash' then
return 0
end
redis.call('EXPIRE', KEYS[1], ARGV[1])
if ARGV[3] == '1' then
redis.call('EXPIRE', KEYS[2], ARGV[2])
end
return 1
"#,
)
});
static DELETE_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
local deleted = redis.call('DEL', KEYS[1])
if ARGV[2] == '1' then
redis.call('SREM', KEYS[2], ARGV[1])
end
return deleted
"#,
)
});
static REVOKE_USER_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
local members = redis.call('SMEMBERS', KEYS[1])
for _, member in ipairs(members) do
redis.call('DEL', ARGV[1] .. member)
end
redis.call('DEL', KEYS[1])
return #members
"#,
)
});
#[derive(Clone)]
pub struct RedisStore {
cache: Cache,
}
impl RedisStore {
pub fn new(cache: Cache) -> Self {
Self { cache }
}
pub async fn connect(config: &RedisConfig) -> Result<Self, SessionError> {
Ok(Self::new(Cache::connect(config).await?))
}
pub fn cache(&self) -> &Cache {
&self.cache
}
fn storage_key(&self, session_key: &SessionKey) -> String {
self.cache.key(&[KEY_PREFIX, session_key.as_str()])
}
fn bucket_key(&self, user_id: Option<i64>) -> String {
self.cache.key(&[USER_BUCKET_PREFIX, &{
user_id.map(|id| id.to_string()).unwrap_or_default()
}])
}
fn session_key_prefix(&self) -> String {
self.cache.key(&[KEY_PREFIX, ""])
}
async fn connection(&self) -> Result<CacheConnection, SessionError> {
Ok(self.cache.acquire().await?)
}
fn rotated_key(&self, _session_key: &SessionKey) -> Result<SessionKey, SessionError> {
// 单机/主从无需 hash slot 标签。带标签的写法会把完整存储键（含 `:`）拼进会话键，
// 而会话键只允许 [A-Za-z0-9{}]，于是轮转必然失败并把登录整条请求打成 500。
if !self.cache.is_cluster() {
return Ok(generate_session_key());
}
// 集群下换键就会换 slot，而槽位标签只能保住同槽——两种需求无法同时满足。
Err(SessionError::config(
"session key rotation is not supported over a Redis Cluster connection",
))
}
}
pub async fn connect(config: &RedisConfig) -> Result<RedisStore, SessionError> {
RedisStore::connect(config).await
}
impl SessionStore for RedisStore {
async fn load(
&self,
session_key: &SessionKey,
) -> Result<Option<(SessionState, Duration)>, SessionError> {
let mut connection = self.connection().await?;
let mut invocation = LOAD_SCRIPT.prepare_invoke();
invocation.key(self.storage_key(session_key));
let entries: Vec<String> = invocation.invoke_async(&mut connection).await?;
let Some((ttl, entries)) = entries.split_first() else {
return Ok(None);
};
let ttl = ttl.parse::<i64>().unwrap_or(0);
if ttl <= 0 {
return Ok(None);
}
let Some(state) = ('inline_state_from_entries: {
let entries = entries.to_vec();
if entries.is_empty() {
break 'inline_state_from_entries None;
}
let mut state = SessionState::with_capacity(entries.len() / 2 + 1);
let mut entries = entries.into_iter();
while let (Some(field), Some(value)) = (entries.next(), entries.next()) {
state.insert(field, value);
}
Some(state)
}) else {
return Ok(None);
};
Ok(Some((state, Duration::seconds(ttl))))
}
async fn save(
&self,
session_state: SessionState,
ttl: &Duration,
) -> Result<SessionKey, SessionError> {
if session_state.is_empty() {
return Err(SessionError::config(
"the session state to save cannot be empty",
));
}
let user_id = {
let state = &session_state;
state
.get(crate::state::SESSION_USER_ID)
.and_then(|value| value.parse::<i64>().ok())
};
let bucket_ttl = ttl_seconds(ttl);
for _ in 0..SAVE_ATTEMPTS {
let session_key = generate_session_key();
let mut invocation = SAVE_SCRIPT.prepare_invoke();
invocation
.key(self.storage_key(&session_key))
.key(self.bucket_key(user_id));
invocation.arg(ttl_seconds(ttl));
invocation.arg(bucket_ttl);
invocation.arg(session_key.as_str());
invocation.arg(if user_id.is_some() { "1" } else { "0" });
for (field, value) in &session_state {
invocation.arg(field).arg(value);
}
let mut connection = self.connection().await?;
let outcome: i64 = invocation.invoke_async(&mut connection).await?;
match outcome {
1 => return Ok(session_key),
0 => continue,
_ => {
return Err(SessionError::config(
"the expiry of the session record cannot be set",
));
}
}
}
Err(SessionError::config(
"the generated session key was already in use, giving up the save attempt",
))
}
async fn update(
&self,
session_key: SessionKey,
diff: SessionDiff,
ttl: &Duration,
previous_user_id: Option<i64>,
user_id: Option<i64>,
) -> Result<SessionKey, SessionError> {
let mut invocation = UPDATE_SCRIPT.prepare_invoke();
invocation
.key(self.storage_key(&session_key))
.key(self.bucket_key(previous_user_id))
.key(self.bucket_key(user_id));
invocation.arg(ttl_seconds(ttl));
invocation.arg(ttl_seconds(ttl));
invocation.arg(session_key.as_str());
invocation.arg(if previous_user_id.is_some() { "1" } else { "0" });
invocation.arg(if user_id.is_some() { "1" } else { "0" });
append_diff_body(&mut invocation, &diff);
let mut connection = self.connection().await?;
let applied: i64 = invocation.invoke_async(&mut connection).await?;
if applied != 1 {
return Err(SessionError::Purged);
}
Ok(session_key)
}
async fn rotate(
&self,
session_key: SessionKey,
diff: SessionDiff,
ttl: &Duration,
previous_user_id: Option<i64>,
user_id: Option<i64>,
) -> Result<Option<SessionKey>, SessionError> {
for _ in 0..SAVE_ATTEMPTS {
let new_session_key = self.rotated_key(&session_key)?;
let mut invocation = ROTATE_SCRIPT.prepare_invoke();
invocation
.key(self.storage_key(&session_key))
.key(self.storage_key(&new_session_key))
.key(self.bucket_key(previous_user_id))
.key(self.bucket_key(user_id));
invocation.arg(ttl_seconds(ttl));
invocation.arg(ttl_seconds(ttl));
invocation.arg(session_key.as_str());
invocation.arg(new_session_key.as_str());
invocation.arg(if previous_user_id.is_some() { "1" } else { "0" });
invocation.arg(if user_id.is_some() { "1" } else { "0" });
append_diff_body(&mut invocation, &diff);
let mut connection = self.connection().await?;
let outcome: i64 = invocation.invoke_async(&mut connection).await?;
match outcome {
1 => return Ok(Some(new_session_key)),
2 => return Ok(None),
0 => return Err(SessionError::Purged),
-1 => continue,
_ => {
return Err(SessionError::config(
"unexpected result while rotating the session record",
));
}
}
}
Err(SessionError::config(
"the generated session key was already in use, giving up the rotation attempt",
))
}
async fn update_ttl(
&self,
session_key: &SessionKey,
ttl: &Duration,
user_id: Option<i64>,
) -> Result<(), SessionError> {
let mut invocation = UPDATE_TTL_SCRIPT.prepare_invoke();
invocation
.key(self.storage_key(session_key))
.key(self.bucket_key(user_id));
invocation.arg(ttl_seconds(ttl));
invocation.arg(ttl_seconds(ttl));
invocation.arg(if user_id.is_some() { "1" } else { "0" });
let mut connection = self.connection().await?;
let renewed: i64 = invocation.invoke_async(&mut connection).await?;
if renewed != 1 {
return Err(SessionError::Purged);
}
Ok(())
}
async fn delete(
&self,
session_key: &SessionKey,
user_id: Option<i64>,
) -> Result<(), SessionError> {
let mut invocation = DELETE_SCRIPT.prepare_invoke();
invocation
.key(self.storage_key(session_key))
.key(self.bucket_key(user_id));
invocation.arg(session_key.as_str());
invocation.arg(if user_id.is_some() { "1" } else { "0" });
let mut connection = self.connection().await?;
let _: i64 = invocation.invoke_async(&mut connection).await?;
Ok(())
}
async fn revoke_user(&self, user_id: i64) -> Result<usize, SessionError> {
let mut invocation = REVOKE_USER_SCRIPT.prepare_invoke();
invocation.key(self.bucket_key(Some(user_id)));
invocation.arg(self.session_key_prefix());
let mut connection = self.connection().await?;
let revoked: usize = invocation.invoke_async(&mut connection).await?;
Ok(revoked)
}
}
fn append_diff_body(invocation: &mut ScriptInvocation<'_>, diff: &SessionDiff) {
invocation.arg(diff.set.len() as i64);
for (field, value) in &diff.set {
invocation.arg(field).arg(value);
}
invocation.arg(diff.removed.len() as i64);
for field in &diff.removed {
invocation.arg(field);
}
}
fn ttl_seconds(ttl: &Duration) -> i64 {
ttl.whole_seconds().max(1)
}