use crate::sync::version_vector::{ConflictInfo, VersionVector};
use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;
pub const MAX_SESSIONS_PER_PRINCIPAL: usize = 64;
pub const MAX_TOTAL_SESSIONS: usize = 10_000;
pub const MAX_SUBSCRIPTIONS: usize = 128;
pub const MAX_SUBSCRIPTION_LEN: usize = 256;
pub const MAX_FILTER_QUERY_LEN: usize = 4096;
pub const MAX_DEVICE_ID_LEN: usize = 128;
const CLEANUP_INTERVAL: std::time::Duration = std::time::Duration::from_secs(600);
pub fn validate_device_id(device_id: &str) -> Result<(), String> {
if device_id.is_empty() || device_id.len() > MAX_DEVICE_ID_LEN {
return Err(format!(
"device_id must be 1 to {} characters",
MAX_DEVICE_ID_LEN
));
}
if !device_id
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.'))
{
return Err("device_id may contain only letters, digits, '-', '_' and '.'".to_string());
}
Ok(())
}
type HmacSha256 = Hmac<Sha256>;
fn sign_hmac(data: &str, secret: &[u8]) -> String {
let mut mac = HmacSha256::new_from_slice(secret).expect("HMAC can take key of any size");
mac.update(data.as_bytes());
hex::encode(mac.finalize().into_bytes())
}
fn verify_hmac(data: &str, signature: &str, secret: &[u8]) -> bool {
let expected = sign_hmac(data, secret);
if expected.len() != signature.len() {
return false;
}
let mut result = 0u8;
for (a, b) in expected.bytes().zip(signature.bytes()) {
result |= a ^ b;
}
result == 0
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SyncSession {
pub session_id: String,
pub device_id: String,
pub user_id: Option<String>,
pub api_key: String,
pub last_vector: VersionVector,
pub last_sequence: u64,
pub filter_query: Option<String>,
pub subscriptions: Vec<String>,
pub created_at: u64,
pub last_activity: u64,
pub is_online: bool,
pub capabilities: ClientCapabilities,
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ClientCapabilities {
pub delta_sync: bool,
pub crdt_types: bool,
pub compression: bool,
pub max_batch_size: usize,
}
impl SyncSession {
pub fn new(
session_id: impl Into<String>,
device_id: impl Into<String>,
api_key: impl Into<String>,
) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
Self {
session_id: session_id.into(),
device_id: device_id.into(),
user_id: None,
api_key: api_key.into(),
last_vector: VersionVector::new(),
last_sequence: 0,
filter_query: None,
subscriptions: Vec::new(),
created_at: now,
last_activity: now,
is_online: true,
capabilities: ClientCapabilities::default(),
}
}
pub fn new_secure(
device_id: impl Into<String>,
api_key: impl Into<String>,
secret: &[u8],
) -> Self {
let device_id = device_id.into();
let api_key = api_key.into();
let nonce = uuid::Uuid::new_v4().to_string();
let data = format!("{}{}{}", device_id, nonce, api_key);
let signature = sign_hmac(&data, secret);
let session_id = format!("{}-{}-{}", device_id, nonce, signature);
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
Self {
session_id,
device_id,
user_id: None,
api_key,
last_vector: VersionVector::new(),
last_sequence: 0,
filter_query: None,
subscriptions: Vec::new(),
created_at: now,
last_activity: now,
is_online: true,
capabilities: ClientCapabilities::default(),
}
}
pub fn verify_session_id(session_id: &str, api_key: &str, secret: &[u8]) -> bool {
if session_id.len() < 65 {
return false;
}
let signature_start = session_id.len() - 64;
if session_id.as_bytes().get(signature_start.saturating_sub(1)) != Some(&b'-') {
return false;
}
let signature = &session_id[signature_start..];
let prefix = &session_id[..signature_start.saturating_sub(1)];
if prefix.len() < 37 {
return false;
}
let nonce_start = prefix.len() - 36;
if prefix.as_bytes().get(nonce_start.saturating_sub(1)) != Some(&b'-') {
return false;
}
let nonce = &prefix[nonce_start..];
let device_id = &prefix[..nonce_start.saturating_sub(1)];
let data = format!("{}{}{}", device_id, nonce, api_key);
verify_hmac(&data, signature, secret)
}
pub fn extract_device_id(session_id: &str) -> Option<String> {
if session_id.len() < 65 + 37 {
return None;
}
let signature_start = session_id.len() - 64;
let prefix = &session_id[..signature_start.saturating_sub(1)];
if prefix.len() < 37 {
return None;
}
let nonce_start = prefix.len() - 36;
let device_id = &prefix[..nonce_start.saturating_sub(1)];
Some(device_id.to_string())
}
pub fn update_vector(&mut self, vector: &VersionVector) {
self.last_vector = vector.clone();
self.last_activity = current_timestamp();
}
pub fn update_sequence(&mut self, sequence: u64) {
self.last_sequence = sequence;
self.last_activity = current_timestamp();
}
pub fn subscribe(&mut self, collection: impl Into<String>) {
let coll = collection.into();
if !self.subscriptions.contains(&coll) {
self.subscriptions.push(coll);
}
self.last_activity = current_timestamp();
}
pub fn unsubscribe(&mut self, collection: &str) {
self.subscriptions.retain(|c| c != collection);
self.last_activity = current_timestamp();
}
pub fn set_online(&mut self, online: bool) {
self.is_online = online;
if online {
self.last_activity = current_timestamp();
}
}
pub fn is_expired(&self, max_inactive_ms: u64) -> bool {
let now = current_timestamp();
now.saturating_sub(self.last_activity) > max_inactive_ms
}
pub fn to_value(&self) -> serde_json::Value {
serde_json::to_value(self).unwrap_or_default()
}
pub fn from_value(value: &serde_json::Value) -> Option<Self> {
serde_json::from_value(value.clone()).ok()
}
}
pub struct SyncSessionManager {
sessions: Arc<RwLock<HashMap<String, SyncSession>>>,
device_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
max_inactive_ms: u64,
cleanup_started: AtomicBool,
}
impl SyncSessionManager {
pub fn new() -> Self {
Self {
sessions: Arc::new(RwLock::new(HashMap::new())),
device_index: Arc::new(RwLock::new(HashMap::new())),
max_inactive_ms: 7 * 24 * 60 * 60 * 1000, cleanup_started: AtomicBool::new(false),
}
}
pub fn with_expiration(max_inactive_ms: u64) -> Self {
Self {
sessions: Arc::new(RwLock::new(HashMap::new())),
device_index: Arc::new(RwLock::new(HashMap::new())),
max_inactive_ms,
cleanup_started: AtomicBool::new(false),
}
}
pub fn spawn_cleanup_task(self: &Arc<Self>) {
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
if self
.cleanup_started
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return;
}
let weak = Arc::downgrade(self);
handle.spawn(async move {
let mut tick = tokio::time::interval(CLEANUP_INTERVAL);
tick.tick().await; loop {
tick.tick().await;
let Some(manager) = weak.upgrade() else {
break;
};
let removed = manager.cleanup_expired().await;
if removed > 0 {
tracing::debug!("sync sessions: dropped {} expired", removed);
}
}
});
}
pub async fn register_session_bounded(
&self,
session: SyncSession,
max_per_owner: usize,
max_total: usize,
) -> Result<(), String> {
let session_id = session.session_id.clone();
let device_id = session.device_id.clone();
let mut evicted: Vec<SyncSession> = Vec::new();
let mut sessions = self.sessions.write().await;
if let Some(owner) = session.user_id.as_deref() {
let mut owned: Vec<(u64, String)> = sessions
.values()
.filter(|s| s.user_id.as_deref() == Some(owner))
.map(|s| (s.last_activity, s.session_id.clone()))
.collect();
if owned.len() >= max_per_owner {
owned.sort();
let excess = owned.len() + 1 - max_per_owner.max(1);
for (_, id) in owned.into_iter().take(excess) {
if let Some(s) = sessions.remove(&id) {
evicted.push(s);
}
}
}
}
if sessions.len() >= max_total {
let expired: Vec<String> = sessions
.values()
.filter(|s| s.is_expired(self.max_inactive_ms))
.map(|s| s.session_id.clone())
.collect();
for id in expired {
if let Some(s) = sessions.remove(&id) {
evicted.push(s);
}
}
}
let result = if sessions.len() >= max_total {
Err(format!(
"too many sync sessions on this server (limit {})",
max_total
))
} else {
sessions.insert(session_id.clone(), session);
Ok(())
};
drop(sessions);
let mut index = self.device_index.write().await;
for s in &evicted {
if let Some(ids) = index.get_mut(&s.device_id) {
ids.retain(|id| id != &s.session_id);
if ids.is_empty() {
index.remove(&s.device_id);
}
}
}
if result.is_ok() {
index.entry(device_id).or_default().push(session_id);
}
result
}
pub async fn register_session(&self, session: SyncSession) {
let session_id = session.session_id.clone();
let device_id = session.device_id.clone();
let mut sessions = self.sessions.write().await;
sessions.insert(session_id.clone(), session);
drop(sessions);
let mut index = self.device_index.write().await;
index
.entry(device_id)
.or_insert_with(Vec::new)
.push(session_id);
}
pub async fn get_session(&self, session_id: &str) -> Option<SyncSession> {
let sessions = self.sessions.read().await;
sessions.get(session_id).cloned()
}
pub async fn get_device_sessions(&self, device_id: &str) -> Vec<SyncSession> {
let index = self.device_index.read().await;
let session_ids = index.get(device_id).cloned().unwrap_or_default();
drop(index);
let sessions = self.sessions.read().await;
session_ids
.iter()
.filter_map(|id| sessions.get(id).cloned())
.collect()
}
pub async fn update_session(&self, session: &SyncSession) {
let mut sessions = self.sessions.write().await;
if sessions.contains_key(&session.session_id) {
sessions.insert(session.session_id.clone(), session.clone());
}
}
pub async fn update_session_vector(&self, session_id: &str, vector: &VersionVector) {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.update_vector(vector);
}
}
pub async fn update_session_sequence(&self, session_id: &str, sequence: u64) {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.update_sequence(sequence);
}
}
pub async fn remove_session(&self, session_id: &str) {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.remove(session_id) {
drop(sessions);
let mut index = self.device_index.write().await;
if let Some(ids) = index.get_mut(&session.device_id) {
ids.retain(|id| id != session_id);
if ids.is_empty() {
index.remove(&session.device_id);
}
}
}
}
pub async fn remove_device_sessions(&self, device_id: &str) {
let index = self.device_index.read().await;
let session_ids = index.get(device_id).cloned().unwrap_or_default();
drop(index);
for session_id in session_ids {
self.remove_session(&session_id).await;
}
}
pub async fn set_session_online(&self, session_id: &str, online: bool) {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.set_online(online);
}
}
pub async fn get_active_sessions(&self) -> Vec<SyncSession> {
let sessions = self.sessions.read().await;
sessions.values().filter(|s| s.is_online).cloned().collect()
}
pub async fn get_subscribers(&self, collection: &str) -> Vec<SyncSession> {
let sessions = self.sessions.read().await;
sessions
.values()
.filter(|s| s.subscriptions.contains(&collection.to_string()))
.cloned()
.collect()
}
pub async fn cleanup_expired(&self) -> usize {
let expired: Vec<String> = {
let sessions = self.sessions.read().await;
sessions
.values()
.filter(|s| s.is_expired(self.max_inactive_ms))
.map(|s| s.session_id.clone())
.collect()
};
for session_id in &expired {
self.remove_session(session_id).await;
}
expired.len()
}
pub async fn session_count(&self) -> usize {
let sessions = self.sessions.read().await;
sessions.len()
}
pub async fn online_count(&self) -> usize {
let sessions = self.sessions.read().await;
sessions.values().filter(|s| s.is_online).count()
}
}
impl Default for SyncSessionManager {
fn default() -> Self {
Self::new()
}
}
fn current_timestamp() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisterSessionRequest {
pub device_id: String,
pub api_key: String,
pub capabilities: Option<ClientCapabilities>,
pub subscriptions: Option<Vec<String>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisterSessionResponse {
pub session_id: String,
pub server_capabilities: ClientCapabilities,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncPullRequest {
pub session_id: String,
pub client_vector: VersionVector,
pub filter: Option<String>,
pub limit: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncPullResponse {
pub changes: Vec<SyncChange>,
pub server_vector: VersionVector,
pub has_more: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncPushRequest {
pub session_id: String,
pub changes: Vec<SyncChange>,
pub client_vector: VersionVector,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncPushResponse {
pub server_vector: VersionVector,
pub conflicts: Vec<ConflictInfo>,
pub accepted: usize,
pub rejected: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncChange {
pub database: String,
pub collection: String,
pub document_key: String,
pub operation: ChangeOperation,
pub document_data: Option<serde_json::Value>,
pub parent_vectors: Vec<VersionVector>,
pub vector: VersionVector,
pub timestamp: u64,
pub is_delta: bool,
pub delta_patch: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ChangeOperation {
Insert,
Update,
Delete,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_session_lifecycle() {
let manager = SyncSessionManager::new();
let session = SyncSession::new("sess-1", "device-1", "api-key-1");
manager.register_session(session.clone()).await;
assert_eq!(manager.session_count().await, 1);
let retrieved = manager.get_session("sess-1").await;
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap().device_id, "device-1");
let mut vector = VersionVector::new();
vector.increment("device-1");
manager.update_session_vector("sess-1", &vector).await;
let updated = manager.get_session("sess-1").await.unwrap();
assert_eq!(updated.last_vector.get("device-1"), 1);
manager.remove_session("sess-1").await;
assert_eq!(manager.session_count().await, 0);
}
#[tokio::test]
async fn test_subscriptions() {
let manager = SyncSessionManager::new();
let mut session = SyncSession::new("sess-1", "device-1", "api-key-1");
session.subscribe("orders");
session.subscribe("products");
manager.register_session(session).await;
let subscribers = manager.get_subscribers("orders").await;
assert_eq!(subscribers.len(), 1);
let subscribers = manager.get_subscribers("users").await;
assert_eq!(subscribers.len(), 0);
}
#[tokio::test]
async fn test_bounded_registration_evicts_oldest_of_owner() {
let manager = SyncSessionManager::new();
for i in 0..3 {
let mut s = SyncSession::new(format!("sess-{}", i), "dev", "k");
s.user_id = Some("alice".to_string());
s.last_activity = i as u64 + 1;
manager.register_session_bounded(s, 3, 100).await.unwrap();
}
let mut s = SyncSession::new("sess-new", "dev", "k");
s.user_id = Some("alice".to_string());
manager.register_session_bounded(s, 3, 100).await.unwrap();
assert_eq!(manager.session_count().await, 3);
assert!(manager.get_session("sess-0").await.is_none());
assert!(manager.get_session("sess-new").await.is_some());
assert_eq!(manager.get_device_sessions("dev").await.len(), 3);
}
#[tokio::test]
async fn test_bounded_registration_total_cap() {
let manager = SyncSessionManager::new();
for i in 0..2 {
let mut s = SyncSession::new(format!("sess-{}", i), "dev", "k");
s.user_id = Some(format!("user-{}", i));
manager.register_session_bounded(s, 10, 2).await.unwrap();
}
let mut s = SyncSession::new("sess-x", "dev", "k");
s.user_id = Some("mallory".to_string());
assert!(manager.register_session_bounded(s, 10, 2).await.is_err());
assert_eq!(manager.session_count().await, 2);
let manager = SyncSessionManager::with_expiration(1000);
let mut old = SyncSession::new("old", "dev", "k");
old.last_activity = 0;
manager.register_session(old).await;
let s = SyncSession::new("fresh", "dev", "k");
manager.register_session_bounded(s, 10, 1).await.unwrap();
assert!(manager.get_session("old").await.is_none());
}
#[test]
fn test_validate_device_id() {
assert!(validate_device_id("iphone-12_a.b").is_ok());
assert!(validate_device_id("").is_err());
assert!(validate_device_id(&"a".repeat(MAX_DEVICE_ID_LEN + 1)).is_err());
assert!(validate_device_id("node b").is_err());
assert!(validate_device_id("10.0.0.1:6745").is_err());
}
#[test]
fn test_session_expiration() {
let mut session = SyncSession::new("sess-1", "device-1", "api-key-1");
session.last_activity = 0;
assert!(session.is_expired(1000));
assert!(!session.is_expired(u64::MAX));
}
#[test]
fn test_secure_session_creation() {
let secret = b"test-cluster-secret-key-12345678";
let session = SyncSession::new_secure("device-123", "api-key-abc", secret);
assert_eq!(session.device_id, "device-123");
assert_eq!(session.api_key, "api-key-abc");
assert!(session.session_id.starts_with("device-123-"));
assert!(session.session_id.len() > 64);
}
#[test]
fn test_verify_session_id_valid() {
let secret = b"test-cluster-secret-key-12345678";
let session = SyncSession::new_secure("device-123", "api-key-abc", secret);
assert!(SyncSession::verify_session_id(
&session.session_id,
"api-key-abc",
secret
));
}
#[test]
fn test_verify_session_id_wrong_api_key() {
let secret = b"test-cluster-secret-key-12345678";
let session = SyncSession::new_secure("device-123", "api-key-abc", secret);
assert!(!SyncSession::verify_session_id(
&session.session_id,
"wrong-api-key",
secret
));
}
#[test]
fn test_verify_session_id_wrong_secret() {
let secret = b"test-cluster-secret-key-12345678";
let wrong_secret = b"wrong-cluster-secret-key-1234567";
let session = SyncSession::new_secure("device-123", "api-key-abc", secret);
assert!(!SyncSession::verify_session_id(
&session.session_id,
"api-key-abc",
wrong_secret
));
}
#[test]
fn test_verify_session_id_tampered() {
let secret = b"test-cluster-secret-key-12345678";
let session = SyncSession::new_secure("device-123", "api-key-abc", secret);
let mut tampered = session.session_id.clone();
if let Some(last_char) = tampered.pop() {
let new_char = if last_char == 'a' { 'b' } else { 'a' };
tampered.push(new_char);
}
assert!(!SyncSession::verify_session_id(
&tampered,
"api-key-abc",
secret
));
}
#[test]
fn test_extract_device_id() {
let secret = b"test-cluster-secret-key-12345678";
let session = SyncSession::new_secure("my-device-id", "api-key-abc", secret);
let extracted = SyncSession::extract_device_id(&session.session_id);
assert_eq!(extracted, Some("my-device-id".to_string()));
}
#[test]
fn test_extract_device_id_invalid() {
assert!(SyncSession::extract_device_id("short").is_none());
assert!(SyncSession::extract_device_id("invalid-session-id").is_none());
}
}