use crate::error::{NodeTokenError, StorageResult};
use crate::protocol::types::{NodeCapabilities, NodeId, NodeSessionId};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::{Path, PathBuf};
use tracing::{debug, info};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[allow(dead_code)] pub struct SessionData {
pub node_id: NodeId,
pub session_id: NodeSessionId,
pub session_token: String,
pub capabilities: NodeCapabilities,
pub poll_timeout_secs: u64,
}
#[allow(dead_code)] pub struct LocalStorage {
data_dir: PathBuf,
session_file: PathBuf,
}
impl LocalStorage {
#[allow(dead_code)] pub fn new(data_dir: Option<&str>) -> StorageResult<Self> {
let data_dir = match data_dir {
Some(dir) => PathBuf::from(dir),
None => Self::default_data_dir()?,
};
let session_file = data_dir.join("session.json");
if !data_dir.exists() {
fs::create_dir_all(&data_dir).map_err(|e| {
NodeTokenError::Storage(format!("Failed to create data directory: {}", e))
})?;
info!("Created data directory: {:?}", data_dir);
}
debug!(
"LocalStorage initialized: data_dir={:?}, session_file={:?}",
data_dir, session_file
);
Ok(Self {
data_dir,
session_file,
})
}
#[allow(dead_code)] fn default_data_dir() -> StorageResult<PathBuf> {
let data_dir = dirs::data_dir()
.ok_or_else(|| {
NodeTokenError::Storage("Failed to get system data directory".to_string())
})?
.join("node-token");
Ok(data_dir)
}
#[allow(dead_code)] pub fn save_session(&self, session: &SessionData) -> StorageResult<()> {
debug!(
"Saving session to file: node_id={}, session_id={}",
session.node_id, session.session_id
);
let json = serde_json::to_string_pretty(session).map_err(|e| {
NodeTokenError::Storage(format!("Failed to serialize session data: {}", e))
})?;
fs::write(&self.session_file, &json)
.map_err(|e| NodeTokenError::Storage(format!("Failed to write session file: {}", e)))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = fs::metadata(&self.session_file)
.map_err(|e| {
NodeTokenError::Storage(format!("Failed to read file metadata: {}", e))
})?
.permissions();
perms.set_mode(0o600); fs::set_permissions(&self.session_file, perms).map_err(|e| {
NodeTokenError::Storage(format!("Failed to set file permissions: {}", e))
})?;
debug!("Set session file permissions to 600");
}
#[cfg(not(unix))]
{
warn!("File permissions setting is not supported on this platform");
}
info!("Session saved successfully to {:?}", self.session_file);
Ok(())
}
#[allow(dead_code)] pub fn load_session(&self) -> StorageResult<Option<SessionData>> {
if !self.session_file.exists() {
debug!("Session file does not exist: {:?}", self.session_file);
return Ok(None);
}
let json = fs::read_to_string(&self.session_file)
.map_err(|e| NodeTokenError::Storage(format!("Failed to read session file: {}", e)))?;
let session: SessionData = serde_json::from_str(&json).map_err(|e| {
NodeTokenError::Storage(format!("Failed to deserialize session data: {}", e))
})?;
debug!(
"Session loaded from file: node_id={}, session_id={}",
session.node_id, session.session_id
);
Ok(Some(session))
}
#[allow(dead_code)] pub fn clear_session(&self) -> StorageResult<()> {
if self.session_file.exists() {
fs::remove_file(&self.session_file).map_err(|e| {
NodeTokenError::Storage(format!("Failed to remove session file: {}", e))
})?;
info!("Session cleared from {:?}", self.session_file);
} else {
debug!("Session file does not exist, nothing to clear");
}
Ok(())
}
#[allow(dead_code)] pub fn data_dir(&self) -> &Path {
&self.data_dir
}
#[allow(dead_code)] pub fn session_file(&self) -> &Path {
&self.session_file
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::protocol::types::NodeModelCapability;
use std::fs;
use tempfile::TempDir;
fn create_test_session() -> SessionData {
SessionData {
node_id: uuid::Uuid::new_v4(),
session_id: uuid::Uuid::new_v4(),
session_token: "test-session-token-secret".to_string(),
capabilities: NodeCapabilities {
runtime: "ollama".to_string(),
models: vec![
NodeModelCapability {
model: "deepseek-chat".to_string(),
},
NodeModelCapability {
model: "llama3".to_string(),
},
],
},
poll_timeout_secs: 30,
}
}
#[test]
fn test_storage_creation() {
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
assert!(storage.data_dir().exists());
assert_eq!(storage.session_file(), temp_dir.path().join("session.json"));
}
#[test]
fn test_storage_default_dir() {
let result = LocalStorage::default_data_dir();
assert!(result.is_ok());
let default_dir = result.unwrap();
assert!(default_dir.ends_with("node-token"));
}
#[test]
fn test_save_and_load_session() {
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
let session = create_test_session();
storage.save_session(&session).unwrap();
assert!(storage.session_file().exists());
let loaded = storage.load_session().unwrap();
assert!(loaded.is_some());
let loaded = loaded.unwrap();
assert_eq!(loaded.node_id, session.node_id);
assert_eq!(loaded.session_id, session.session_id);
assert_eq!(loaded.session_token, session.session_token);
assert_eq!(loaded.capabilities.runtime, session.capabilities.runtime);
assert_eq!(
loaded.capabilities.models.len(),
session.capabilities.models.len()
);
}
#[test]
fn test_load_nonexistent_session() {
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
let loaded = storage.load_session().unwrap();
assert!(loaded.is_none());
}
#[test]
fn test_clear_session() {
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
let session = create_test_session();
storage.save_session(&session).unwrap();
assert!(storage.session_file().exists());
storage.clear_session().unwrap();
assert!(!storage.session_file().exists());
storage.clear_session().unwrap();
}
#[test]
fn test_session_file_permissions() {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
let session = create_test_session();
storage.save_session(&session).unwrap();
let metadata = fs::metadata(storage.session_file()).unwrap();
let perms = metadata.permissions();
let mode = perms.mode() & 0o777;
assert_eq!(mode, 0o600, "File permissions should be 600");
}
}
#[test]
fn test_corrupted_session_file() {
let temp_dir = TempDir::new().unwrap();
let storage = LocalStorage::new(Some(temp_dir.path().to_str().unwrap())).unwrap();
fs::write(storage.session_file(), "not valid json").unwrap();
let result = storage.load_session();
assert!(result.is_err());
}
#[test]
fn test_session_data_serialization() {
let session = create_test_session();
let json = serde_json::to_string(&session).unwrap();
let deserialized: SessionData = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.node_id, session.node_id);
assert_eq!(deserialized.session_id, session.session_id);
assert_eq!(deserialized.session_token, session.session_token);
}
}