f2e745a
use std::sync::LazyLock;
use actix_web::cookie::time::Duration;
use cache::redis::{AsyncCommands, Script, ScriptInvocation};
use cache::{Cache, CacheConnection};
use config::identity::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 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
return redis.call('HGETALL', KEYS[1])
"#,
)
});
static SAVE_SCRIPT: LazyLock<Script> = LazyLock::new(|| {
Script::new(
r#"
if redis.call('EXISTS', KEYS[1]) == 1 then
return 0
end
for index = 2, #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
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 = 3
local set_count = tonumber(ARGV[2])
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])
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 = 3
local set_count = tonumber(ARGV[2])
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
return 1
"#,
)
});
#[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()])
}
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 load() -> Result<RedisStore, SessionError> {
Ok(RedisStore::new(Cache::load().await?))
}
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>, 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?;
Ok(state_from_entries(entries))
}
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",
));
}
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));
invocation.arg(ttl_seconds(ttl));
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,
) -> Result<SessionKey, SessionError> {
let mut invocation = UPDATE_SCRIPT.prepare_invoke();
invocation.key(self.storage_key(&session_key));
append_diff(&mut invocation, &diff, ttl);
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,
) -> 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));
append_diff(&mut invocation, &diff, ttl);
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,
) -> Result<(), SessionError> {
let mut connection = self.connection().await?;
let renewed: bool = connection
.expire(self.storage_key(session_key), ttl_seconds(ttl))
.await?;
if !renewed {
return Err(SessionError::Purged);
}
Ok(())
}
async fn delete(&self, session_key: &SessionKey) -> Result<(), SessionError> {
let mut connection = self.connection().await?;
let _: usize = connection.del(self.storage_key(session_key)).await?;
Ok(())
}
}
fn state_from_entries(entries: Vec<String>) -> Option<SessionState> {
if entries.is_empty() {
return 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)
}
fn append_diff(invocation: &mut ScriptInvocation<'_>, diff: &SessionDiff, ttl: &Duration) {
invocation.arg(ttl_seconds(ttl));
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)
}