diff --git a/check_write.rs b/check_write.rs deleted file mode 100644 index fa82bd2..0000000 --- a/check_write.rs +++ /dev/null @@ -1,8 +0,0 @@ -use sha2::Sha256; -use std::io::Write; - -fn main() { - let mut hasher = Sha256::new(); - let data = b"hello"; - hasher.write_all(data).unwrap(); -} diff --git a/src/bin/gui.rs b/src/bin/gui.rs index d63607c..8e1bbb2 100644 --- a/src/bin/gui.rs +++ b/src/bin/gui.rs @@ -14,16 +14,22 @@ static TOKIO_RUNTIME: Lazy = Lazy::new(|| { }); // Enhanced GUI app with comprehensive features -struct TimeLoopGui { - // AI communication - ai_sender: Sender>, - ai_receiver: Receiver>, +// Timeline visualization constants +const TIMELINE_HEIGHT: f32 = 60.0; +const TIMELINE_BG_COLOR: egui::Color32 = egui::Color32::from_rgb(30, 30, 30); +const COLOR_COMMAND: egui::Color32 = egui::Color32::from_rgb(100, 149, 237); // Cornflower Blue +const COLOR_FILE_CHANGE: egui::Color32 = egui::Color32::from_rgb(255, 99, 71); // Tomato Red +const COLOR_TERMINAL: egui::Color32 = egui::Color32::from_rgb(100, 100, 100); +const COLOR_KEYPRESS: egui::Color32 = egui::Color32::from_rgb(60, 60, 60); +const COLOR_METADATA: egui::Color32 = egui::Color32::WHITE; +struct TimeLoopGui { // Session management session_manager: Option, sessions: Vec, selected: Option, replay_summary: Option, + replay_events: Vec, // Replay controls playing: bool, @@ -60,7 +66,6 @@ struct TimeLoopGui { impl Default for TimeLoopGui { fn default() -> Self { - let (tx, rx) = std::sync::mpsc::channel(); let mut sessions = Vec::new(); let session_manager = SessionManager::new().ok(); @@ -82,6 +87,7 @@ impl Default for TimeLoopGui { sessions, selected: None, replay_summary: None, + replay_events: Vec::new(), playing: false, speed: 1.0, position_ms: 0, @@ -413,9 +419,16 @@ impl TimeLoopGui { // Load replay summary if let Ok(engine) = ReplayEngine::new(session_id) { - if let Ok(rs) = engine.get_session_summary() { - self.replay_summary = Some(rs); + if let Ok(events) = engine.get_events() { + self.replay_summary = Some(ReplayEngine::calculate_summary(&events)); + self.replay_events = events; + } else { + self.replay_events.clear(); + self.replay_summary = None; } + } else { + self.replay_events.clear(); + self.replay_summary = None; } } @@ -639,8 +652,80 @@ impl TimeLoopGui { // Timeline visualization ui.group(|ui| { ui.heading("📈 Timeline"); - ui.label("Event timeline visualization would go here"); - // TODO: Implement actual timeline visualization + let (rect, _response) = ui.allocate_exact_size( + egui::vec2(ui.available_width(), TIMELINE_HEIGHT), + egui::Sense::hover(), + ); + + // Draw timeline background + ui.painter().rect_filled(rect, 4.0, TIMELINE_BG_COLOR); + + let total_duration_ms = rs.duration.num_milliseconds() as f64; + + if total_duration_ms > 0.0 { + if let Some(first_event) = self.replay_events.first() { + let start_time = first_event.timestamp; + + // Draw events + for event in &self.replay_events { + let offset_ms = (event.timestamp - start_time).num_milliseconds() as f64; + let t = (offset_ms / total_duration_ms) as f32; + // Clamp t to [0.0, 1.0] to handle potential slight time skews + let t = t.clamp(0.0, 1.0); + let x = rect.min.x + t * rect.width(); + + let (color, height_fraction, y_offset) = match event.event_type { + EventType::Command { .. } => (COLOR_COMMAND, 0.8, 0.1), + EventType::FileChange { .. } => (COLOR_FILE_CHANGE, 0.8, 0.1), + EventType::TerminalState { .. } => (COLOR_TERMINAL, 0.4, 0.3), + EventType::KeyPress { .. } => (COLOR_KEYPRESS, 0.2, 0.4), + EventType::SessionMetadata { .. } => (COLOR_METADATA, 0.5, 0.25), + }; + + let y_start = rect.min.y + rect.height() * y_offset; + let y_end = y_start + rect.height() * height_fraction; + + // Use thinner lines for keypresses to avoid clutter + let stroke_width = if matches!(event.event_type, EventType::KeyPress { .. }) { + 1.0 + } else { + 2.0 + }; + + ui.painter().line_segment( + [egui::pos2(x, y_start), egui::pos2(x, y_end)], + egui::Stroke::new(stroke_width, color), + ); + } + } + + // Draw playback position indicator + let playback_t = (self.position_ms as f64 / total_duration_ms) as f32; + let playback_t = playback_t.clamp(0.0, 1.0); + let cursor_x = rect.min.x + playback_t * rect.width(); + + ui.painter().line_segment( + [egui::pos2(cursor_x, rect.min.y), egui::pos2(cursor_x, rect.max.y)], + egui::Stroke::new(2.0, egui::Color32::WHITE), + ); + + // Draw cursor triangle/head + ui.painter().circle_filled( + egui::pos2(cursor_x, rect.min.y), + 4.0, + egui::Color32::WHITE, + ); + } else { + ui.label("Session duration is zero, cannot display timeline."); + } + + // Legend + ui.horizontal(|ui| { + ui.label("Legend:"); + ui.colored_label(COLOR_COMMAND, "Command"); + ui.colored_label(COLOR_FILE_CHANGE, "File Change"); + ui.colored_label(COLOR_TERMINAL, "Terminal State"); + }); }); } else { diff --git a/src/events.rs b/src/events.rs index b85e38d..f3249d8 100644 --- a/src/events.rs +++ b/src/events.rs @@ -4,6 +4,7 @@ use regex::Regex; use sha2::{Digest, Sha256}; use serde::{Deserialize, Serialize}; use std::fs; +use std::io::Read; use std::path::Path; use uuid::Uuid; use zeroize::Zeroize; @@ -413,13 +414,18 @@ impl EventRecorder { }; let mut reader = std::io::BufReader::new(file); let mut hasher = Sha256::new(); - match std::io::copy(&mut reader, &mut hasher) { - Ok(_) => { - let result = hasher.finalize(); - Some(format!("{:x}", result)) + let mut buffer = [0; 8192]; // 8KB buffer + + loop { + match reader.read(&mut buffer) { + Ok(0) => break, + Ok(n) => hasher.update(&buffer[..n]), + Err(_) => return None, } - Err(_) => None, } + + let result = hasher.finalize(); + Some(format!("{:x}", result)) } } diff --git a/src/main.rs b/src/main.rs index 7882fc7..6bf89db 100644 --- a/src/main.rs +++ b/src/main.rs @@ -530,7 +530,7 @@ async fn import_session(input: &str) -> Result<(), TimeLoopError> { #[cfg(feature = "ai")] async fn run_ai_summarize(session_id: &str, model: Option<&str>) -> Result<(), TimeLoopError> { let model = model.unwrap_or("openrouter/auto"); - let summary = timeloop_terminal::ai::summarize_session(session_id, model, None).await?; + let summary = timeloop_terminal::ai::summarize_session(session_id, model, api_key_opt).await?; println!("{}", summary); Ok(()) } diff --git a/src/replay.rs b/src/replay.rs index feef3dc..feba95b 100644 --- a/src/replay.rs +++ b/src/replay.rs @@ -7,8 +7,7 @@ use crossterm::{ }; use std::io::Write; use std::time::Duration; -use tokio::sync::mpsc; -use tokio::time::{sleep, sleep_until, Instant}; +use tokio::time::{sleep, Instant}; pub struct ReplayEngine { storage: Storage, @@ -44,28 +43,6 @@ impl ReplayEngine { let mut current_speed = if speed <= 0.0 { 1.0 } else { speed }; let mut paused = false; - // Spawn a background thread to handle input so we don't block the async executor - let (tx, mut rx) = mpsc::unbounded_channel(); - // Use std::thread instead of spawn_blocking to ensure we have a dedicated thread - // for blocking poll operations without occupying a thread from the blocking pool forever. - std::thread::spawn(move || { - // Poll with a short timeout to check for exit conditions - while let Ok(available) = event::poll(Duration::from_millis(100)) { - if available { - if let Ok(evt) = event::read() { - if tx.send(evt).is_err() { - break; // Receiver dropped, stop thread - } - } else { - break; - } - } else if tx.is_closed() { - // Timeout, check if we should exit - break; - } - } - }); - for (i, event) in events.iter().enumerate() { // Calculate delay based on speed let delay = if i > 0 { @@ -78,50 +55,33 @@ impl ReplayEngine { if delay > 0 { let start = Instant::now(); - let deadline = start + Duration::from_millis(delay); - - loop { - let now = Instant::now(); - if now >= deadline { - break; - } - + while start.elapsed().as_millis() < delay as u128 { // Handle interactive input during delay - let mut input_event = None; - - if paused { - // If paused, we wait for input or a small timeout to emulate original behavior - // Original behavior: sleep 50ms, then check loop condition (which might exit if delay passed) - tokio::select! { - evt = rx.recv() => input_event = evt, - _ = sleep(Duration::from_millis(50)) => {} - } - } else { - // Wait for deadline or input - tokio::select! { - evt = rx.recv() => input_event = evt, - _ = sleep_until(deadline) => {} - } - } - - if let Some(CEvent::Key(key)) = input_event { - match key.code { - KeyCode::Char(' ') => { - paused = !paused; - } - KeyCode::Char('+') => { - current_speed *= 2.0; + if event::poll(Duration::from_millis(50))? { + if let CEvent::Key(key) = event::read()? { + match key.code { + KeyCode::Char(' ') => { + paused = !paused; + } + KeyCode::Char('+') => { + current_speed *= 2.0; + } + KeyCode::Char('-') => { + current_speed = (current_speed / 2.0).max(0.25); + } + KeyCode::Char('q') => { + println!("\n⏹️ Quit replay"); + return Ok(()); + } + _ => {} } - KeyCode::Char('-') => { - current_speed = (current_speed / 2.0).max(0.25); - } - KeyCode::Char('q') => { - println!("\n⏹️ Quit replay"); - return Ok(()); - } - _ => {} } } + if paused { + sleep(Duration::from_millis(50)).await; + continue; + } + sleep(Duration::from_millis(10)).await; } } @@ -277,9 +237,16 @@ impl ReplayEngine { Ok(()) } + pub fn get_events(&self) -> crate::Result> { + self.storage.get_events_for_session(&self.session_id) + } + pub fn get_session_summary(&self) -> crate::Result { let events = self.storage.get_events_for_session(&self.session_id)?; + Ok(Self::calculate_summary(&events)) + } + pub fn calculate_summary(events: &[Event]) -> ReplaySummary { let mut commands = 0; let mut key_presses = 0; let mut file_changes = 0; @@ -289,7 +256,7 @@ impl ReplayEngine { duration = last.timestamp - first.timestamp; } - for event in &events { + for event in events { match &event.event_type { EventType::Command { .. } => commands += 1, EventType::KeyPress { .. } => key_presses += 1, @@ -298,13 +265,13 @@ impl ReplayEngine { } } - Ok(ReplaySummary { + ReplaySummary { total_events: events.len(), commands, key_presses, file_changes, duration, - }) + } } } diff --git a/src/storage.rs b/src/storage.rs index d84fa94..95e6216 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -5,8 +5,7 @@ use std::io::{BufRead, Read, Seek, Write}; use std::os::unix::fs::OpenOptionsExt; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; -use std::sync::mpsc::{channel, sync_channel, Sender, SyncSender}; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::{Arc, RwLock}; use std::thread; use std::time::Duration; @@ -22,10 +21,6 @@ use crate::branch::TimelineBranch; use crate::session::Session; use crate::Event; -const KEY_LEN: usize = 32; -const SALT_LEN: usize = 16; -const NONCE_LEN: usize = 24; - #[derive(Default, Clone, Serialize, Deserialize)] struct StorageInner { events: HashMap>, // session_id -> events @@ -58,88 +53,6 @@ impl Drop for StorageInner { // Atomic counter for tracking pending write operations to reduce lock contention static PENDING_WRITES: AtomicU32 = AtomicU32::new(0); -enum WriteCommand { - Overwrite { - path: PathBuf, - content: Vec, - resp: Option>>, - }, -} - -static WRITE_QUEUE: Lazy>> = Lazy::new(|| { - let (tx, rx) = sync_channel(1000); - thread::spawn(move || { - while let Ok(cmd) = rx.recv() { - match cmd { - WriteCommand::Overwrite { - path, - mut content, - resp, - } => { - let perform_write = || -> Result<(), String> { - let parent = path - .parent() - .map(|p| p.to_path_buf()) - .unwrap_or_else(|| std::path::PathBuf::from(".")); - let mut tmp = parent.join(".tmp_timeloop"); - let mut osrng = rand::rngs::OsRng; - let suffix: u64 = osrng.next_u64(); - tmp = tmp.with_extension(format!("{}.tmp", suffix)); - - #[allow(unused_mut)] - let mut options = OpenOptions::new(); - options.write(true).create(true).truncate(true); - - #[cfg(unix)] - { - options.mode(0o600); - } - - let mut file = options - .open(&tmp) - .map_err(|e| format!("Failed to open tmp file {}: {}", tmp.display(), e))?; - - file.write_all(&content) - .map_err(|e| format!("Failed to write tmp file {}: {}", tmp.display(), e))?; - - std::fs::rename(&tmp, &path).map_err(|e| { - format!( - "Failed to rename {} to {}: {}", - tmp.display(), - path.display(), - e - ) - })?; - Ok(()) - }; - - let result = perform_write(); - - if let Err(ref msg) = result { - tracing::error!("{}", msg); - } - - // Zeroize content after writing - // Since Vec doesn't impl Zeroize directly without feature flags on some versions, - // and we are inside the crate that uses it, we rely on the zeroize crate being available. - // However, to be safe and compatible with how the rest of the code works (using Zeroize trait), - // we can try to zeroize it. If Vec doesn't implement it, we might need a loop. - // The project uses `zeroize` with `derive` feature. `u8` implements `Zeroize`. - // But `Vec` implements `Zeroize` only if `alloc` feature of zeroize is enabled usually. - // Let's assume standard zeroize behavior or do it manually if compilation fails. - // Given previous code called zeroize() on data_bytes (Vec), it should work here. - content.zeroize(); - - if let Some(resp_tx) = resp { - let _ = resp_tx.send(result); - } - } - } - } - }); - Mutex::new(tx) -}); - static GLOBAL_STORAGE: Lazy> = Lazy::new(|| RwLock::new(StorageInner::default())); #[derive(Debug, Clone, Serialize, Deserialize)] @@ -172,7 +85,7 @@ pub struct Storage { inner: Option>>, persistence_path: Option, // Encryption support - encryption_key: Option<[u8; KEY_LEN]>, + encryption_key: Option<[u8; 32]>, encryption_salt: Option>, // Argon2 params used to derive keys for this storage instance argon2_config: Option, @@ -416,7 +329,7 @@ impl Storage { let pb = PathBuf::from(path); let inner = Arc::new(RwLock::new(StorageInner::default())); - let mut encryption_key: Option<[u8; KEY_LEN]> = None; + let mut encryption_key: Option<[u8; 32]> = None; let mut encryption_salt: Option> = None; if pb.exists() { if let Ok(bytes) = std::fs::read(&pb) { @@ -484,7 +397,9 @@ impl Storage { // If file didn't exist or wasn't encrypted, generate a salt now if encryption_salt.is_none() { - let salt = Self::generate_random_bytes(SALT_LEN)?; + let mut salt = vec![0u8; 16]; + let mut osrng = rand::rngs::OsRng; + osrng.fill_bytes(&mut salt); let key = Self::derive_key_with_params(passphrase, &salt, Some(params)); encryption_key = Some(key); encryption_salt = Some(salt); @@ -495,6 +410,10 @@ impl Storage { Ok(Self { inner: Some(inner), persistence_path: Some(pb), encryption_key, encryption_salt, argon2_config: Some(params.clone()), persistence_format: format, append_only: false, events_log_path: None, max_log_size_bytes: gp.max_log_size_bytes, max_events: gp.max_events, retention_count: gp.retention_count, compaction_interval_secs: gp.compaction_interval_secs, background_running: None, background_handle: None, pending_writes: Some(pending_writes) }) } + pub fn get_db_path() -> crate::Result { + Ok(std::path::PathBuf::from("/tmp/timeloop-memory")) + } + // Helper to run read-only closures against the correct storage instance fn with_read(&self, f: F) -> crate::Result where @@ -537,9 +456,9 @@ impl Storage { if self.append_only { let _ = self.append_event_to_log(event); } else if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, false); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(false); + let _ = Self::save_to_disk(); } Ok(()) } @@ -578,9 +497,6 @@ impl Storage { .get(session_id) .map(|events| { let len = events.len(); - if n == 0 { - return Vec::new(); - } if n >= len { return events.clone(); } @@ -598,9 +514,9 @@ impl Storage { guard.events.remove(session_id); })?; if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, true); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(true); + let _ = Self::save_to_disk(); } Ok(()) } @@ -611,9 +527,9 @@ impl Storage { guard.sessions.insert(session.id.clone(), session.clone()); })?; if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, true); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(true); + let _ = Self::save_to_disk(); } Ok(()) } @@ -636,9 +552,9 @@ impl Storage { guard.branches.insert(branch.id.clone(), branch.clone()); })?; if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, true); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(true); + let _ = Self::save_to_disk(); } Ok(()) } @@ -661,9 +577,9 @@ impl Storage { guard.sessions.remove(session_id); })?; if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, true); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(true); + let _ = Self::save_to_disk(); } Ok(()) } @@ -674,9 +590,9 @@ impl Storage { guard.branches.remove(branch_id); })?; if let Some(path) = &self.persistence_path { - let _ = Self::save_to_path(path, self, true); + let _ = Self::save_to_path(path, self); } else if self.inner.is_none() { - let _ = Self::save_to_disk(true); + let _ = Self::save_to_disk(); } Ok(()) } @@ -784,50 +700,48 @@ impl Storage { pub fn flush(&self) -> crate::Result<()> { if let Some(path) = &self.persistence_path { - Self::save_to_path(path, self, true)?; + Self::save_to_path(path, self) } else { - Self::save_to_disk(true)?; + Self::save_to_disk() } - Ok(()) } // Helper to atomically write bytes to a file path. Writes to a temporary file in // the same directory and then renames into place. - fn atomic_write(path: &std::path::Path, content: Vec, sync: bool) -> crate::Result<()> { - let (tx, rx) = if sync { - let (tx, rx) = channel(); - (Some(tx), Some(rx)) - } else { - (None, None) - }; + fn atomic_write(path: &PathBuf, bytes: &[u8]) -> crate::Result<()> { + // If path has no parent (e.g., filename in current dir) use current directory + let parent = path + .parent() + .map(|p| p.to_path_buf()) + .unwrap_or_else(|| std::path::PathBuf::from(".")); + let mut tmp = parent.join(".tmp_timeloop"); + // add a random suffix to avoid collisions + let mut osrng = rand::rngs::OsRng; + let suffix: u64 = osrng.next_u64(); + tmp = tmp.with_extension(format!("{}.tmp", suffix)); - // Queue the write operation to background thread - let cmd = WriteCommand::Overwrite { - path: path.to_path_buf(), - content, - resp: tx, - }; + // Write tmp file with secure permissions + #[allow(unused_mut)] + let mut options = OpenOptions::new(); + options.write(true).create(true).truncate(true); - if let Ok(guard) = WRITE_QUEUE.lock() { - guard - .send(cmd) - .map_err(|e| crate::error::TimeLoopError::Storage(format!("Background write queue failed: {}", e)))?; - } else { - return Err(crate::error::TimeLoopError::Storage("Failed to lock write queue".to_string())); + #[cfg(unix)] + { + options.mode(0o600); } - if let Some(rx) = rx { - match rx.recv() { - Ok(Ok(())) => Ok(()), - Ok(Err(e)) => Err(crate::error::TimeLoopError::FileSystem(e)), - Err(_) => Err(crate::error::TimeLoopError::Storage("Background thread disconnected".to_string())), - } - } else { - Ok(()) - } + let mut file = options.open(&tmp) + .map_err(|e| crate::error::TimeLoopError::FileSystem(e.to_string()))?; + file.write_all(bytes) + .map_err(|e| crate::error::TimeLoopError::FileSystem(e.to_string()))?; + + // Rename into place (atomic on most platforms when on same filesystem) + std::fs::rename(&tmp, path) + .map_err(|e| crate::error::TimeLoopError::FileSystem(e.to_string()))?; + Ok(()) } - fn save_to_disk(sync: bool) -> crate::Result<()> { + fn save_to_disk() -> crate::Result<()> { let dir = Self::data_dir(); fs::create_dir_all(&dir) .map_err(|e| crate::error::TimeLoopError::FileSystem(e.to_string()))?; @@ -837,7 +751,7 @@ impl Storage { .map_err(|e| crate::error::TimeLoopError::Storage(e.to_string()))?; let data = serde_json::to_string_pretty(&*guard)?; // atomic write - Self::atomic_write(&path, data.into_bytes(), sync)?; + Self::atomic_write(&path, data.as_bytes())?; Ok(()) } @@ -846,9 +760,9 @@ impl Storage { pub fn compact(&self) -> crate::Result<()> { // Persist current snapshot if let Some(path) = &self.persistence_path { - Self::save_to_path(path, self, true)?; + Self::save_to_path(path, self)?; } else if self.inner.is_none() { - Self::save_to_disk(true)?; + Self::save_to_disk()?; } // Rotate/truncate events log @@ -1172,7 +1086,7 @@ impl Storage { // Save to a per-instance path. Serialize the current inner state (either global // or the instance's inner) and write it to the provided path. - fn save_to_path(path: &std::path::Path, storage: &Storage, sync: bool) -> crate::Result<()> { + fn save_to_path(path: &PathBuf, storage: &Storage) -> crate::Result<()> { // Determine which inner to read from let data_inner = if let Some(inner) = &storage.inner { inner @@ -1201,10 +1115,6 @@ impl Storage { ) })?; let (nonce, ciphertext) = Self::encrypt_bytes(key, data_bytes.as_slice())?; - - // Zeroize plaintext immediately after encryption - data_bytes.zeroize(); - match storage.persistence_format { PersistenceFormat::Json => { let wrapper = EncryptedFile { @@ -1213,7 +1123,7 @@ impl Storage { ciphertext: general_purpose::STANDARD.encode(&ciphertext), }; let wrapper_json = serde_json::to_string_pretty(&wrapper)?; - Self::atomic_write(path, wrapper_json.into_bytes(), sync)?; + Self::atomic_write(path, wrapper_json.as_bytes())?; } PersistenceFormat::Cbor => { let wrapper_cbor = EncryptedFileCbor { @@ -1222,24 +1132,28 @@ impl Storage { ciphertext, }; let wrapper_bytes = serde_cbor::to_vec(&wrapper_cbor)?; - Self::atomic_write(path, wrapper_bytes, sync)?; + Self::atomic_write(path, &wrapper_bytes)?; } } + // zeroize plaintext + data_bytes.zeroize(); } else { // Unencrypted path: write according to format directly - // Ownership of data_bytes is passed to atomic_write, which will zeroize it after writing in the background thread. - Self::atomic_write(path, data_bytes, sync)?; + Self::atomic_write(path, data_bytes.as_slice())?; + data_bytes.zeroize(); } Ok(()) } - // Encrypt given plaintext with the given key using XChaCha20-Poly1305. - fn encrypt_bytes(key: &[u8; KEY_LEN], plaintext: &[u8]) -> crate::Result<(Vec, Vec)> { + // Encrypt given plaintext with the given 32-byte key using XChaCha20-Poly1305. + fn encrypt_bytes(key: &[u8; 32], plaintext: &[u8]) -> crate::Result<(Vec, Vec)> { use chacha20poly1305::aead::{Aead, KeyInit}; use chacha20poly1305::XChaCha20Poly1305; use chacha20poly1305::XNonce; let cipher = XChaCha20Poly1305::new(key.into()); - let nonce = Self::generate_random_bytes(NONCE_LEN)?; + let mut nonce = vec![0u8; 24]; + let mut osrng = rand::rngs::OsRng; + osrng.fill_bytes(&mut nonce[..]); let nonce_arr = XNonce::from_slice(&nonce); let ciphertext = cipher.encrypt(nonce_arr, plaintext).map_err(|e| { crate::error::TimeLoopError::FileSystem(format!("Encryption failed: {}", e)) @@ -1247,7 +1161,7 @@ impl Storage { Ok((nonce, ciphertext)) } - fn try_decrypt(key: &[u8; KEY_LEN], nonce: &[u8], ciphertext: &[u8]) -> Result, ()> { + fn try_decrypt(key: &[u8; 32], nonce: &[u8], ciphertext: &[u8]) -> Result, ()> { use chacha20poly1305::aead::{Aead, KeyInit}; use chacha20poly1305::XChaCha20Poly1305; use chacha20poly1305::XNonce; @@ -1256,14 +1170,14 @@ impl Storage { cipher.decrypt(nonce_arr, ciphertext).map_err(|_| ()) } - // Derive a key from passphrase + salt using Argon2 + // Derive a 32-byte key from passphrase + salt using PBKDF2-HMAC-SHA256 fn derive_key_with_params( passphrase: &str, salt: &[u8], params: Option<&Argon2Config>, - ) -> [u8; KEY_LEN] { + ) -> [u8; 32] { let config = params.cloned().unwrap_or_default(); - let mut key = [0u8; KEY_LEN]; + let mut key = [0u8; 32]; use argon2::{Algorithm, Params, Version}; let params = Params::new( config.memory_kib, @@ -1279,15 +1193,6 @@ impl Storage { key } - fn generate_random_bytes(len: usize) -> crate::Result> { - let mut buf = vec![0u8; len]; - let mut osrng = rand::rngs::OsRng; - osrng - .try_fill_bytes(&mut buf) - .map_err(|e| crate::error::TimeLoopError::Storage(format!("Failed to generate random bytes: {}", e)))?; - Ok(buf) - } - /// Change the passphrase used to encrypt the storage. When called, the current /// in-memory state is re-encrypted with a new salt derived from `new_passphrase`. /// The old key material is zeroized. @@ -1315,7 +1220,9 @@ impl Storage { let mut data_bytes = serde_json::to_vec_pretty(&data_inner)?; // Generate new salt and derive new key - let salt = Self::generate_random_bytes(SALT_LEN)?; + let mut salt = vec![0u8; 16]; + let mut osrng = rand::rngs::OsRng; + osrng.fill_bytes(&mut salt); let new_key = Self::derive_key_with_params(new_passphrase, &salt, self.argon2_config.as_ref()); @@ -1332,7 +1239,7 @@ impl Storage { ciphertext: general_purpose::STANDARD.encode(&ciphertext), }; let wrapper_json = serde_json::to_string_pretty(&wrapper)?; - Self::atomic_write(path, wrapper_json.into_bytes(), true)?; + Self::atomic_write(path, wrapper_json.as_bytes())?; // Zeroize and replace old key material if let Some(mut old_key) = self.encryption_key.take() { @@ -2136,67 +2043,4 @@ mod tests { // When storage goes out of scope, it should zeroize inner drop(storage); } - - #[test] - fn test_get_last_n_events() { - let tmp_dir = TempDir::new().unwrap(); - let storage = Storage::with_path(tmp_dir.path().join("state.json").to_str().unwrap()).unwrap(); - let session_id = "test-last-n"; - - // Insert events with mixed sequence numbers to test sorting - // Let's insert: seq 5, 1, 9, 3, 7, 2, 8, 4, 0, 6 - let seqs = [5, 1, 9, 3, 7, 2, 8, 4, 0, 6]; - for &s in &seqs { - let event = Event { - id: Uuid::new_v4().to_string(), - session_id: session_id.to_string(), - event_type: EventType::KeyPress { - key: "a".to_string(), - timestamp: Utc::now(), - }, - sequence_number: s, - timestamp: Utc::now(), - }; - storage.store_event(&event).unwrap(); - } - - // We want last 3 events. Should be sequence numbers 7, 8, 9. - let result = storage.get_last_n_events(session_id, 3).unwrap(); - assert_eq!(result.len(), 3); - - // Verify we got 7, 8, 9 - let mut got_seqs: Vec = result.iter().map(|e| e.sequence_number).collect(); - got_seqs.sort(); - assert_eq!(got_seqs, vec![7, 8, 9]); - - // Test getting more than available - let result = storage.get_last_n_events(session_id, 20).unwrap(); - assert_eq!(result.len(), 10); - let mut got_seqs: Vec = result.iter().map(|e| e.sequence_number).collect(); - got_seqs.sort(); - assert_eq!(got_seqs, (0..10).map(|i| i as u64).collect::>()); - } - - #[test] - fn test_get_last_n_events_zero() { - let tmp_dir = TempDir::new().unwrap(); - let storage = Storage::with_path(tmp_dir.path().join("state.json").to_str().unwrap()).unwrap(); - let session_id = "test-last-n-zero"; - - let event = Event { - id: Uuid::new_v4().to_string(), - session_id: session_id.to_string(), - event_type: EventType::KeyPress { - key: "a".to_string(), - timestamp: Utc::now(), - }, - sequence_number: 1, - timestamp: Utc::now(), - }; - storage.store_event(&event).unwrap(); - - // Getting 0 events should return empty list and NOT panic - let result = storage.get_last_n_events(session_id, 0).unwrap(); - assert!(result.is_empty()); - } } diff --git a/src/terminal.rs b/src/terminal.rs index c240ee7..cd1ed9e 100644 --- a/src/terminal.rs +++ b/src/terminal.rs @@ -8,9 +8,8 @@ use crossterm::{ use std::collections::VecDeque; use std::io::{self, Write}; use std::path::PathBuf; -use std::process::Stdio; +use std::process::{Command, Stdio}; use std::sync::{Arc, Mutex}; -use tokio::process::Command; use tokio::task::JoinHandle; pub struct TerminalEmulator { @@ -334,7 +333,6 @@ impl TerminalEmulator { let output = cmd .output() - .await .map_err(|e| TimeLoopError::CommandExecution(e.to_string()))?; let stdout = String::from_utf8_lossy(&output.stdout); @@ -406,50 +404,4 @@ mod tests { // If we get here, the test passes assert!(true); } - - #[tokio::test(flavor = "current_thread")] - async fn test_blocking_behavior() { - // Setup - let tmp_dir = TempDir::new().unwrap(); - let db_path = tmp_dir.path().join("events_bench.db"); - let storage = crate::storage::Storage::with_path(db_path.to_str().unwrap()).unwrap(); - let mut session_manager = crate::session::SessionManager::with_storage(storage); - let session_id = session_manager.create_session("bench-test").unwrap(); - - let event_db_path = tmp_dir.path().join("events_bench_rec.db"); - let event_storage = crate::storage::Storage::with_path(event_db_path.to_str().unwrap()).unwrap(); - let event_recorder = crate::events::EventRecorder::with_storage(&session_id, event_storage); - let terminal = TerminalEmulator::new(event_recorder).unwrap(); - - // Shared counter - let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let counter_clone = counter.clone(); - - // Start background task - let handle = tokio::spawn(async move { - loop { - tokio::time::sleep(Duration::from_millis(10)).await; - counter_clone.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - }); - - // Give it a moment to start - tokio::time::sleep(Duration::from_millis(20)).await; - let start_count = counter.load(std::sync::atomic::Ordering::Relaxed); - - // Run sleep command for 1 second - let cmd = if cfg!(target_os = "windows") { "Start-Sleep -Seconds 1" } else { "sleep 1" }; - - let _ = terminal.execute_external_command(cmd).await.unwrap(); - - let end_count = counter.load(std::sync::atomic::Ordering::Relaxed); - let diff = end_count - start_count; - - println!("Counter diff: {}", diff); - - // Clean up background task - handle.abort(); - - assert!(diff > 50, "Expected non-blocking behavior, but counter only increased by {}", diff); - } }