use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
use crate::engine::resume_data::{ResumeData, ResumeDataExt};
use crate::http::cookie_storage::CookieJar;
use crate::request::request_group::{DownloadOptions, DownloadStatus, GroupId, RequestGroup};
use crate::selector::server_stat_man::ServerStatMan;
const SESSION_OPTIONS_FILENAME: &str = "session_options.json";
const DEFAULT_AUTO_SAVE_INTERVAL_SECS: u64 = 60;
pub struct SessionPersistence {
session_dir: PathBuf,
auto_save_interval: Duration,
auto_save_enabled: bool,
cookie_jar: Option<CookieJar>,
}
impl SessionPersistence {
pub fn new(session_dir: &Path) -> Self {
Self {
session_dir: session_dir.to_path_buf(),
auto_save_interval: Duration::from_secs(DEFAULT_AUTO_SAVE_INTERVAL_SECS),
auto_save_enabled: true,
cookie_jar: None,
}
}
pub fn with_interval(mut self, interval_secs: u64) -> Self {
self.auto_save_interval = Duration::from_secs(interval_secs.max(10));
self
}
pub fn without_auto_save(mut self) -> Self {
self.auto_save_enabled = false;
self
}
pub fn with_cookie_jar(mut self, jar: CookieJar) -> Self {
self.cookie_jar = Some(jar);
self
}
pub fn cookie_jar_mut(&mut self) -> Option<&mut CookieJar> {
self.cookie_jar.as_mut()
}
pub fn cookie_jar(&self) -> Option<&CookieJar> {
self.cookie_jar.as_ref()
}
pub fn session_dir(&self) -> &Path {
&self.session_dir
}
pub async fn save_state(&self, groups: &[Arc<RwLock<RequestGroup>>]) -> Result<usize, String> {
tokio::fs::create_dir_all(&self.session_dir)
.await
.map_err(|e| {
format!(
"Failed to create session dir {}: {}",
self.session_dir.display(),
e
)
})?;
let mut saved = 0usize;
for group_lock in groups.iter() {
let group = group_lock.read().await;
match ResumeData::from_request_group(&group).await {
Ok(resume_data) => {
let file_name = format!("{}.aria2", resume_data.gid);
let path = self.session_dir.join(&file_name);
if let Err(e) = resume_data.save_to_file(&path) {
warn!(
gid = %resume_data.gid,
error = %e,
"Failed to save resume data for GID"
);
continue;
}
saved += 1;
debug!(
gid = %resume_data.gid,
path = %path.display(),
"Saved resume data"
);
}
Err(e) => {
debug!(
gid = %group.gid().value(),
error = %e,
"Skipping command that cannot be serialized"
);
}
}
}
self.save_global_options(groups).await?;
if let Some(ref jar) = self.cookie_jar {
let cookie_path = self.session_dir.join("cookies.json");
if let Err(e) = Self::save_cookie_jar_to_file(jar, &cookie_path).await {
warn!("Failed to persist cookies: {}", e);
} else {
debug!(path = %cookie_path.display(), "Cookies persisted to session");
}
}
info!(
saved,
dir = %self.session_dir.display(),
"Session state saved"
);
Ok(saved)
}
pub async fn load_state(
&mut self,
groups: &mut Vec<Arc<RwLock<RequestGroup>>>,
) -> Result<usize, String> {
if !self.session_dir.exists() {
debug!(
dir = %self.session_dir.display(),
"Session directory does not exist, nothing to load"
);
return Ok(0);
}
let mut loaded = 0usize;
let mut entries = tokio::fs::read_dir(&self.session_dir).await.map_err(|e| {
format!(
"Failed to read session dir {}: {}",
self.session_dir.display(),
e
)
})?;
while let Ok(Some(entry)) = entries.next_entry().await {
let path = entry.path();
let is_aria2 = path.extension().map(|e| e == "aria2").unwrap_or(false);
if !is_aria2 {
continue;
}
match ResumeData::load_from_file(&path) {
Ok(Some(resume_data)) => {
match Self::restore_command(&resume_data) {
Ok(group) => {
groups.push(Arc::new(RwLock::new(group)));
loaded += 1;
info!(
gid = %resume_data.gid,
status = %resume_data.status,
"Restored download from session"
);
}
Err(e) => {
warn!(
gid = %resume_data.gid,
error = %e,
"Failed to restore command from resume data"
);
}
}
}
Ok(None) => {
debug!(path = %path.display(), "Resume file was empty (skipped)");
}
Err(e) => {
warn!(
path = %path.display(),
error = %e,
"Corrupted or invalid .aria2 file, skipping gracefully"
);
}
}
}
let _ = self.load_global_options().await;
let cookie_path = self.session_dir.join("cookies.json");
if cookie_path.exists() {
match Self::load_cookie_jar_from_file(&cookie_path).await {
Ok(jar) => {
self.cookie_jar = Some(jar);
info!("Loaded cookies from session");
}
Err(e) => {
warn!("Failed to load cookies from session: {}", e);
}
}
}
info!(
loaded,
dir = %self.session_dir.display(),
"Session state loaded"
);
Ok(loaded)
}
fn restore_command(resume_data: &ResumeData) -> Result<RequestGroup, String> {
if resume_data.uris.is_empty() {
return Err("ResumeData has no URIs, cannot restore".to_string());
}
let uris: Vec<String> = resume_data.uris.iter().map(|u| u.uri.clone()).collect();
let mut options = DownloadOptions::default();
if let Some(ref output_path) = resume_data.output_path {
if let Some(parent) = Path::new(output_path).parent() {
options.dir = Some(parent.to_string_lossy().to_string());
}
if let Some(file_name) = Path::new(output_path).file_name() {
options.out = Some(file_name.to_string_lossy().to_string());
}
}
let gid = if !resume_data.gid.is_empty() {
GroupId::from_hex_string(&resume_data.gid).unwrap_or_else(GroupId::new_random)
} else {
GroupId::new_random()
};
let group = RequestGroup::new(gid, uris, options);
if resume_data.status == "paused" || resume_data.status == "waiting" {
}
if resume_data.completed_length > 0 {
group.set_resume_offset(resume_data.completed_length);
}
Ok(group)
}
pub fn start_auto_save(
&self,
groups: Arc<RwLock<Vec<Arc<RwLock<RequestGroup>>>>>,
) -> Option<tokio::task::JoinHandle<()>> {
if !self.auto_save_enabled {
debug!("Auto-save is disabled");
return None;
}
let session_dir = self.session_dir.clone();
let interval = self.auto_save_interval;
info!(
interval_secs = interval.as_secs(),
dir = %session_dir.display(),
"Starting auto-save task"
);
Some(tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
loop {
ticker.tick().await;
let groups_read = groups.read().await;
let persistence = SessionPersistence::new(&session_dir).without_auto_save();
match persistence.save_state(&groups_read).await {
Ok(count) => {
if count > 0 {
debug!(count, "Auto-save completed successfully");
}
}
Err(e) => {
warn!(error = %e, "Auto-save failed, will retry next interval");
}
}
}
}))
}
async fn save_global_options(
&self,
_groups: &[Arc<RwLock<RequestGroup>>],
) -> Result<(), String> {
let opts_path = self.session_dir.join(SESSION_OPTIONS_FILENAME);
let options_summary = serde_json::json!({
"version": "1.0",
"saved_at": chrono_timestamp_or_fallback(),
"note": "Global session options summary"
});
let json = serde_json::to_string_pretty(&options_summary)
.map_err(|e| format!("Failed to serialize session options: {}", e))?;
tokio::fs::write(&opts_path, json).await.map_err(|e| {
format!(
"Failed to write session options {}: {}",
opts_path.display(),
e
)
})?;
Ok(())
}
async fn load_global_options(&self) -> Result<(), String> {
let opts_path = self.session_dir.join(SESSION_OPTIONS_FILENAME);
if !opts_path.exists() {
return Ok(());
}
let content = tokio::fs::read_to_string(&opts_path).await.map_err(|e| {
format!(
"Failed to read session options {}: {}",
opts_path.display(),
e
)
})?;
let _parsed: serde_json::Value = serde_json::from_str(&content)
.map_err(|e| format!("Invalid JSON in session options: {}", e))?;
debug!(path = %opts_path.display(), "Loaded session options");
Ok(())
}
pub async fn cleanup(&self) -> Result<(), String> {
if !self.session_dir.exists() {
return Ok(());
}
let mut entries = tokio::fs::read_dir(&self.session_dir)
.await
.map_err(|e| format!("Failed to read session dir: {}", e))?;
while let Ok(Some(entry)) = entries.next_entry().await {
let path = entry.path();
if let Err(e) = tokio::fs::remove_file(&path).await {
warn!(path = %path.display(), error = %e, "Failed to remove session file");
}
}
info!(dir = %self.session_dir.display(), "Session directory cleaned up");
Ok(())
}
pub async fn save_server_stats(&self, stat_man: &ServerStatMan) -> Result<usize, String> {
let stat_file = self.session_dir.join("server-stat.json");
let saved = stat_man.save_to_file_async(&stat_file).await?;
if saved > 0 {
debug!(
count = saved,
path = %stat_file.display(),
"Server statistics saved"
);
}
Ok(saved)
}
pub async fn load_server_stats(&self, stat_man: &ServerStatMan) -> Result<usize, String> {
let stat_file = self.session_dir.join("server-stat.json");
if !stat_file.exists() {
debug!("No server statistics file found, starting fresh");
return Ok(0);
}
let loaded = stat_man.load_from_file_async(&stat_file).await?;
if loaded > 0 {
info!(
count = loaded,
path = %stat_file.display(),
"Server statistics loaded from previous session"
);
}
Ok(loaded)
}
pub async fn save_active_only(
&self,
groups: &[Arc<RwLock<RequestGroup>>],
) -> Result<usize, String> {
let mut count = 0;
for group in groups {
let g = group.read().await;
let status = g.status().await;
match status {
DownloadStatus::Active | DownloadStatus::Waiting => {
drop(g);
let group_read = group.read().await;
match ResumeData::from_request_group(&group_read).await {
Ok(resume_data) => {
drop(group_read);
let file_name = format!("{}.aria2", resume_data.gid);
let path = self.session_dir.join(&file_name);
if resume_data.save_to_file(&path).is_ok() {
count += 1;
debug!(gid = %resume_data.gid, "Saved active download");
} else {
warn!(gid = %resume_data.gid, "Failed to save active download");
}
}
Err(e) => {
debug!(error = %e, "Skipping active download that cannot be serialized");
}
}
}
_ => {} }
}
debug!(
saved = count,
total = groups.len(),
"save_active_only completed"
);
Ok(count)
}
pub async fn save_completed(
&self,
groups: &[Arc<RwLock<RequestGroup>>],
) -> Result<usize, String> {
let mut count = 0;
for group in groups {
let g = group.read().await;
let status = g.status().await;
if status.is_completed() || matches!(status, DownloadStatus::Complete) {
drop(g);
let group_read = group.read().await;
match ResumeData::from_request_group(&group_read).await {
Ok(resume_data) => {
drop(group_read);
let file_name = format!("{}.aria2", resume_data.gid);
let path = self.session_dir.join(&file_name);
if resume_data.save_to_file(&path).is_ok() {
count += 1;
debug!(gid = %resume_data.gid, "Saved completed download");
}
}
Err(e) => {
debug!(error = %e, "Skipping completed download that cannot be serialized");
}
}
}
}
debug!(
saved = count,
total = groups.len(),
"save_completed completed"
);
Ok(count)
}
async fn save_cookie_jar_to_file(jar: &CookieJar, path: &Path) -> Result<(), String> {
#[derive(Serialize)]
struct SerializableJar<'a> {
cookies: &'a [crate::http::cookie_storage::JarCookie],
}
let serializable = SerializableJar {
cookies: &jar.cookies,
};
let json = serde_json::to_string_pretty(&serializable).map_err(|e| e.to_string())?;
tokio::fs::write(path, json)
.await
.map_err(|e| format!("Failed to write cookie file: {}", e))
}
async fn load_cookie_jar_from_file(path: &Path) -> Result<CookieJar, String> {
let content = tokio::fs::read_to_string(path)
.await
.map_err(|e| format!("Failed to read cookie file: {}", e))?;
#[derive(Deserialize)]
struct SerializableJar {
cookies: Vec<crate::http::cookie_storage::JarCookie>,
}
let parsed: SerializableJar =
serde_json::from_str(&content).map_err(|e| format!("Invalid cookie JSON: {}", e))?;
let mut jar = CookieJar::new();
for cookie in parsed.cookies {
jar.store(cookie);
}
Ok(jar)
}
}
fn chrono_timestamp_or_fallback() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs().to_string())
.unwrap_or_else(|_| "unknown".to_string())
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DhtStateSnapshot {
pub nodes: Vec<DhtNodeInfo>,
pub token_secret: [u8; 20],
pub last_bootstrap_epoch_secs: Option<u64>,
pub total_nodes: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DhtNodeInfo {
pub id: [u8; 20],
pub addr: String,
pub last_seen_epoch_secs: u64,
}
impl DhtStateSnapshot {
pub fn empty() -> Self {
Self {
nodes: vec![],
token_secret: [0u8; 20],
last_bootstrap_epoch_secs: None,
total_nodes: 0,
}
}
pub fn new(
nodes: Vec<DhtNodeInfo>,
token_secret: [u8; 20],
last_bootstrap_epoch_secs: Option<u64>,
) -> Self {
let total_nodes = nodes.len();
Self {
nodes,
token_secret,
last_bootstrap_epoch_secs,
total_nodes,
}
}
pub fn to_json_string(&self) -> Result<String, String> {
serde_json::to_string(self).map_err(|e| e.to_string())
}
pub fn from_json_string(json: &str) -> Result<Self, String> {
serde_json::from_str(json).map_err(|e| e.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
fn create_test_session_dir() -> PathBuf {
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos()
% 1_000_000_000;
let dir =
std::env::temp_dir().join(format!("aria2_session_test_{}_{}", std::process::id(), ts));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).expect("Failed to create test session directory");
dir
}
fn create_test_groups(count: usize) -> Vec<Arc<RwLock<RequestGroup>>> {
let mut groups = Vec::new();
for i in 0..count {
let gid = GroupId::new(i as u64 + 1000);
let uri = format!("http://example.com/file{}.bin", i);
let options = DownloadOptions {
dir: Some("/downloads".to_string()),
split: Some(4),
..Default::default()
};
let group = Arc::new(RwLock::new(RequestGroup::new(gid, vec![uri], options)));
groups.push(group);
}
groups
}
#[tokio::test]
async fn test_session_save_creates_files() {
let session_dir = create_test_session_dir();
let persistence = SessionPersistence::new(&session_dir);
let groups = create_test_groups(3);
let saved_count = persistence
.save_state(&groups)
.await
.expect("Save should succeed");
assert_eq!(saved_count, 3, "Should save 3 commands");
let entries: Vec<_> = fs::read_dir(&session_dir)
.expect("Should read session dir")
.filter_map(|e| e.ok())
.collect();
let aria2_count = entries
.iter()
.filter(|e| {
e.path()
.extension()
.map(|ext| ext == "aria2")
.unwrap_or(false)
})
.count();
assert_eq!(aria2_count, 3, "Should have 3 .aria2 files");
for entry in entries.iter().filter(|e| {
e.path()
.extension()
.map(|ext| ext == "aria2")
.unwrap_or(false)
}) {
let content = fs::read_to_string(entry.path()).expect("Should read file");
let parsed: serde_json::Value =
serde_json::from_str(&content).expect("Should be valid JSON");
assert!(
parsed.get("gid").is_some(),
"Each .aria2 file should contain a GID field"
);
}
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_session_load_restores_commands() {
let session_dir = create_test_session_dir();
let mut persistence = SessionPersistence::new(&session_dir);
let original_groups = create_test_groups(2);
let saved = persistence
.save_state(&original_groups)
.await
.expect("Save should succeed");
assert_eq!(saved, 2, "Should save 2 commands");
let mut loaded_groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let loaded = persistence
.load_state(&mut loaded_groups)
.await
.expect("Load should succeed");
assert_eq!(loaded, 2, "Should restore 2 commands");
let mut found_uris: Vec<String> = Vec::new();
for group_lock in &loaded_groups {
let group = group_lock.read().await;
for uri in group.uris() {
found_uris.push(uri.clone());
}
}
assert!(
found_uris.iter().any(|u| u.contains("file0.bin")),
"Should restore first file URI"
);
assert!(
found_uris.iter().any(|u| u.contains("file1.bin")),
"Should restore second file URI"
);
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_session_save_empty_no_error() {
let session_dir = create_test_session_dir();
let persistence = SessionPersistence::new(&session_dir);
let empty_groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let result = persistence.save_state(&empty_groups).await;
assert!(result.is_ok(), "Saving empty session should not error");
let saved_count = result.unwrap();
assert_eq!(saved_count, 0, "Empty session should report 0 saved");
assert!(
session_dir.exists(),
"Session dir should be created even for empty save"
);
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_session_corrupted_file_skipped_gracefully() {
let session_dir = create_test_session_dir();
let corrupt_file = session_dir.join("corrupt-gid.aria2");
fs::write(&corrupt_file, "THIS IS NOT VALID JSON {{{{").expect("Should write corrupt file");
let valid_file = session_dir.join("valid-gid.aria2");
let valid_resume_data = ResumeData {
gid: "valid-gid-12345".to_string(),
uris: vec![crate::engine::resume_data::UriState {
uri: "http://example.com/valid-file.bin".to_string(),
tried: true,
used: false,
last_result: None,
speed_bytes_per_sec: None,
}],
total_length: 1024,
completed_length: 512,
uploaded_length: 0,
bitfield: vec![],
num_pieces: None,
piece_length: None,
status: "paused".to_string(),
error_message: None,
last_download_time: 0,
created_at: 0,
output_path: Some("/downloads/valid-file.bin".to_string()),
checksum: None,
options: std::collections::HashMap::new(),
resume_offset: Some(512),
bt_info_hash: None,
bt_saved_metadata_path: None,
};
valid_resume_data
.save_to_file(&valid_file)
.expect("Should write valid file");
let mut persistence = SessionPersistence::new(&session_dir);
let mut loaded_groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let result = persistence.load_state(&mut loaded_groups).await;
assert!(result.is_ok(), "Load should succeed despite corrupt file");
let loaded_count = result.unwrap();
assert_eq!(
loaded_count, 1,
"Should load 1 valid file (corrupt one skipped)"
);
assert_eq!(loaded_groups.len(), 1, "Should have 1 restored group");
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_session_load_nonexistent_dir_returns_zero() {
let nonexistent_dir =
PathBuf::from("/tmp/aria2_nonexistent_test_dir_that_should_not_exist_12345");
let mut persistence = SessionPersistence::new(&nonexistent_dir);
let mut groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let result = persistence.load_state(&mut groups).await;
assert!(result.is_ok(), "Nonexistent dir should return Ok");
assert_eq!(result.unwrap(), 0, "Nonexistent dir should return 0 loaded");
assert!(groups.is_empty(), "No groups should be added");
}
#[tokio::test]
async fn test_session_cleanup_removes_all_files() {
let session_dir = create_test_session_dir();
let persistence = SessionPersistence::new(&session_dir);
let groups = create_test_groups(2);
let _ = persistence.save_state(&groups).await.unwrap();
assert!(
session_dir.exists(),
"Session dir should exist before cleanup"
);
persistence.cleanup().await.expect("Cleanup should succeed");
if session_dir.exists() {
let remaining: Vec<_> = fs::read_dir(&session_dir)
.expect("Should read dir")
.filter_map(|e| e.ok())
.collect();
assert!(
remaining.is_empty(),
"All files should be removed after cleanup"
);
}
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_session_custom_interval() {
let session_dir = create_test_session_dir();
let persistence = SessionPersistence::new(&session_dir).with_interval(30);
assert_eq!(
persistence.auto_save_interval,
Duration::from_secs(30),
"Custom interval should be set"
);
let short_interval = SessionPersistence::new(&session_dir).with_interval(1);
assert!(
short_interval.auto_save_interval >= Duration::from_secs(10),
"Interval should be at least 10 seconds"
);
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_resume_data_roundtrip_via_persistence() {
let session_dir = create_test_session_dir();
let mut persistence = SessionPersistence::new(&session_dir);
let gid = GroupId::new(0xDEADBEEF);
let options = DownloadOptions {
dir: Some("/test/downloads".to_string()),
out: Some("special_file.iso".to_string()),
split: Some(16),
..Default::default()
};
let group = Arc::new(RwLock::new(RequestGroup::new(
gid,
vec!["http://example.com/special_file.iso".to_string()],
options,
)));
{
let g = group.write().await;
g.set_total_length_atomic(10485760); g.set_completed_length(5242880); }
let saved = persistence.save_state(&[group]).await.unwrap();
assert_eq!(saved, 1);
let mut loaded: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let loaded_count = persistence.load_state(&mut loaded).await.unwrap();
assert_eq!(loaded_count, 1);
let restored = loaded[0].read().await;
let uris = restored.uris();
assert_eq!(uris.len(), 1);
assert!(uris[0].contains("special_file.iso"));
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_selective_save_active_only() {
let session_dir = create_test_session_dir();
let persistence = SessionPersistence::new(&session_dir);
let mut groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let active_gid = GroupId::new(1001);
let active_group = Arc::new(RwLock::new(RequestGroup::new(
active_gid,
vec!["http://example.com/active.bin".to_string()],
DownloadOptions::default(),
)));
{
let mut g = active_group.write().await;
g.start().await.unwrap(); }
groups.push(active_group);
let waiting_gid = GroupId::new(1002);
let waiting_group = Arc::new(RwLock::new(RequestGroup::new(
waiting_gid,
vec!["http://example.com/waiting.bin".to_string()],
DownloadOptions::default(),
)));
groups.push(waiting_group);
let complete_gid = GroupId::new(1003);
let complete_group = Arc::new(RwLock::new(RequestGroup::new(
complete_gid,
vec!["http://example.com/complete.bin".to_string()],
DownloadOptions::default(),
)));
{
let mut g = complete_group.write().await;
g.complete().await.unwrap(); }
groups.push(complete_group);
let saved_count = persistence.save_active_only(&groups).await.unwrap();
assert_eq!(
saved_count, 2,
"save_active_only should save only active and waiting downloads"
);
let entries: Vec<_> = fs::read_dir(&session_dir)
.expect("Should read session dir")
.filter_map(|e| e.ok())
.filter(|e| {
e.path()
.extension()
.map(|ext| ext == "aria2")
.unwrap_or(false)
})
.collect();
assert_eq!(
entries.len(),
2,
"Should have exactly 2 .aria2 files for active+waiting"
);
let _ = fs::remove_dir_all(&session_dir);
}
#[test]
fn test_dht_snapshot_roundtrip() {
use crate::session::session_persistence::{DhtNodeInfo, DhtStateSnapshot};
let node1 = DhtNodeInfo {
id: [1u8; 20],
addr: "192.168.1.100:6881".to_string(),
last_seen_epoch_secs: 1700000000,
};
let node2 = DhtNodeInfo {
id: [2u8; 20],
addr: "10.0.0.5:6881".to_string(),
last_seen_epoch_secs: 1700000100,
};
let token_secret: [u8; 20] = [0xAB; 20];
let original = DhtStateSnapshot::new(vec![node1, node2], token_secret, Some(1699999000));
assert_eq!(original.total_nodes, 2);
assert_eq!(original.nodes.len(), 2);
assert!(original.last_bootstrap_epoch_secs.is_some());
let json = original
.to_json_string()
.expect("Serialization should succeed");
assert!(!json.is_empty(), "JSON output should not be empty");
assert!(
json.contains("192.168.1.100"),
"JSON should contain first node address"
);
assert!(
json.contains("10.0.0.5"),
"JSON should contain second node address"
);
let restored =
DhtStateSnapshot::from_json_string(&json).expect("Deserialization should succeed");
assert_eq!(restored.total_nodes, 2, "total_nodes should be preserved");
assert_eq!(restored.nodes.len(), 2, "nodes count should be preserved");
assert_eq!(
restored.token_secret, token_secret,
"token_secret should be preserved"
);
assert_eq!(
restored.last_bootstrap_epoch_secs,
Some(1699999000),
"last_bootstrap_epoch_secs should be preserved"
);
assert_eq!(
restored.nodes[0].id, [1u8; 20],
"First node ID should match"
);
assert_eq!(
restored.nodes[0].addr, "192.168.1.100:6881",
"First node address should match"
);
assert_eq!(
restored.nodes[0].last_seen_epoch_secs, 1700000000,
"First node timestamp should match"
);
assert_eq!(
restored.nodes[1].id, [2u8; 20],
"Second node ID should match"
);
assert_eq!(
restored.nodes[1].addr, "10.0.0.5:6881",
"Second node address should match"
);
let empty = DhtStateSnapshot::empty();
assert_eq!(empty.total_nodes, 0, "Empty snapshot should have 0 nodes");
assert!(
empty.nodes.is_empty(),
"Empty snapshot should have no nodes"
);
let empty_json = empty
.to_json_string()
.expect("Empty serialization should succeed");
let empty_restored = DhtStateSnapshot::from_json_string(&empty_json)
.expect("Empty deserialization should work");
assert_eq!(
empty_restored.total_nodes, 0,
"Restored empty should still be empty"
);
}
#[tokio::test]
async fn test_cookie_persist_integration() {
use crate::http::cookie_storage::{CookieJar, JarCookie};
let session_dir = create_test_session_dir();
let mut jar = CookieJar::new();
jar.store(JarCookie::new("session_id", "abc123", "example.com"));
jar.store(JarCookie::new("auth_token", "xyz789", "api.example.com"));
let persistence_with_cookies = SessionPersistence::new(&session_dir).with_cookie_jar(jar);
assert!(
persistence_with_cookies.cookie_jar().is_some(),
"Cookie jar should be set"
);
assert_eq!(
persistence_with_cookies.cookie_jar().unwrap().len(),
2,
"Should have 2 cookies before save"
);
let groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let _saved = persistence_with_cookies.save_state(&groups).await.unwrap();
let cookie_path = session_dir.join("cookies.json");
assert!(
cookie_path.exists(),
"cookies.json file should exist after save"
);
let mut persistence_new = SessionPersistence::new(&session_dir);
let mut loaded_groups: Vec<Arc<RwLock<RequestGroup>>> = Vec::new();
let _loaded = persistence_new
.load_state(&mut loaded_groups)
.await
.unwrap();
assert!(
persistence_new.cookie_jar().is_some(),
"Cookie jar should exist after load"
);
let loaded_jar = persistence_new.cookie_jar().unwrap();
assert_eq!(
loaded_jar.len(),
2,
"Should have loaded 2 cookies from file"
);
let example_cookies = loaded_jar.get_cookies_for_url("http://example.com/", false);
assert_eq!(
example_cookies.len(),
1,
"Should find 1 cookie for example.com"
);
assert_eq!(example_cookies[0].name, "session_id");
assert_eq!(example_cookies[0].value, "abc123");
let api_cookies = loaded_jar.get_cookies_for_url("http://api.example.com/api", false);
assert_eq!(
api_cookies.len(),
2,
"Should find 2 cookies for api.example.com (parent domain + exact)"
);
let auth_cookie = api_cookies
.iter()
.find(|c| c.name == "auth_token")
.expect("Should find auth_token cookie");
assert_eq!(auth_cookie.value, "xyz789");
let session_cookie = api_cookies
.iter()
.find(|c| c.name == "session_id")
.expect("Should find session_id cookie from parent domain");
assert_eq!(session_cookie.value, "abc123");
let _ = fs::remove_dir_all(&session_dir);
}
#[tokio::test]
async fn test_auto_save_with_custom_interval() {
let session_dir = create_test_session_dir();
let persistence_30s = SessionPersistence::new(&session_dir).with_interval(30);
assert_eq!(
persistence_30s.auto_save_interval,
Duration::from_secs(30),
"Custom 30s interval should be set"
);
let persistence_too_short = SessionPersistence::new(&session_dir).with_interval(1);
assert!(
persistence_too_short.auto_save_interval >= Duration::from_secs(10),
"Interval below minimum should be clamped to 10s"
);
let persistence_exact_min = SessionPersistence::new(&session_dir).with_interval(10);
assert_eq!(
persistence_exact_min.auto_save_interval,
Duration::from_secs(10),
"Exact minimum interval (10s) should be accepted"
);
let persistence_large = SessionPersistence::new(&session_dir).with_interval(300); assert_eq!(
persistence_large.auto_save_interval,
Duration::from_secs(300),
"Large interval (300s) should be accepted"
);
let persistence_default = SessionPersistence::new(&session_dir);
assert!(
persistence_default.auto_save_enabled,
"Auto-save should be enabled by default"
);
assert_eq!(
persistence_default.auto_save_interval,
Duration::from_secs(DEFAULT_AUTO_SAVE_INTERVAL_SECS),
"Default interval should be 60 seconds"
);
let persistence_disabled = SessionPersistence::new(&session_dir).without_auto_save();
assert!(
!persistence_disabled.auto_save_enabled,
"Auto-save should be disabled after without_auto_save()"
);
let _ = fs::remove_dir_all(&session_dir);
}
}