1use async_trait::async_trait;
2use easy_trace::prelude::{debug, error, info};
3use serde_json::{Value};
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::Mutex;
7use crate::service::{Service, WebSocketWrite};
8use std::io::{Read, Write};
9use std::net::TcpStream;
10use std::path::Path;
11use sharing::service::{Action, SafeHandle};
12use sharing::crypto;
13
14const PERSISTENCE_PASSWORD: &str = "R4nd0mlyG3n3r4t3dP4ssw0rd";
15
16pub const PERSISTENCE_VERSION: &str = "1.0.0";
18
19pub struct PersistenceService {
21 pub cache: Arc<Mutex<HashMap<String, Value>>>,
23 pub file_handle: Arc<Mutex<Option<SafeHandle>>>,
25 pub file_path: String,
27 pub using_service_protection: bool,
29}
30
31impl PersistenceService {
32 pub fn new() -> Self {
34 info!("PersistenceService initialized (version: {})", PERSISTENCE_VERSION);
35
36 let file_path = sharing::paths::get_persistance_path();
37 let service = PersistenceService {
38 cache: Arc::new(Mutex::new(HashMap::new())),
39 file_handle: Arc::new(Mutex::new(None)),
40 file_path: file_path.display().to_string(),
41 using_service_protection: false,
42 };
43
44 service
45 }
46
47 pub async fn initialize(&mut self) -> Result<(), String> {
49 #[cfg(target_os = "windows")]
54 let handle_result = self.try_get_handle_from_service();
55 #[cfg(not(target_os = "windows"))]
56 let handle_result: Result<SafeHandle, String> =
57 Err("Skipping service handle on Unix (FDs are per-process)".to_string());
58
59 match handle_result {
60 Ok(handle) => {
61 info!("Got handle from service.exe successfully");
62 self.using_service_protection = true;
63 *self.file_handle.lock().await = Some(handle);
64 }
65 Err(e) => {
66 info!("Using direct file open: {}", e);
67
68 match self.try_open_file_directly_with_handle() {
70 Ok(handle) => {
71 info!("Got direct file handle successfully");
72 *self.file_handle.lock().await = Some(handle);
73 }
74 Err(e) => {
75 error!("Failed to open file directly: {}", e);
76 info!("PersistenceService will work in memory-only mode");
77 }
78 }
79 }
80 }
81
82 if let Err(e) = self.try_load_data().await {
84 error!("Failed to load existing data: {}", e);
85 }
86
87 Ok(())
88 }
89
90 #[cfg(target_os = "windows")]
95 fn try_get_handle_from_service(&self) -> Result<SafeHandle, String> {
96 info!("Requesting file handle from service.exe for: {}", self.file_path);
97
98 let mut stream = match TcpStream::connect(sharing::CHANNEL_NAME) {
99 Ok(stream) => stream,
100 Err(e) => return Err(format!("Could not connect to service.exe: {}", e)),
101 };
102
103 info!("Connected to service.exe pipe successfully");
104
105 let action = Action::GetHandle(self.file_path.clone());
107 match stream.write_all(action.to_string().as_bytes()) {
108 Ok(_) => {},
109 Err(e) => return Err(format!("Failed to send GET_HANDLE request: {}", e)),
110 }
111
112 info!("GET_HANDLE request sent, waiting for response...");
113
114 let mut buffer = [0; 1024];
116 let bytes_read = match stream.read(&mut buffer) {
117 Ok(bytes) => bytes,
118 Err(e) => return Err(format!("Failed to read response: {}", e)),
119 };
120
121 let response = String::from_utf8_lossy(&buffer[..bytes_read]);
122 info!("Received response from service.exe: {}", response);
123
124 if response.starts_with("HANDLE_OK:") {
125 let handle_str = response.strip_prefix("HANDLE_OK:").unwrap_or("");
127 match handle_str.parse::<usize>() {
128 Ok(handle_value) => {
129 let handle = SafeHandle::new(handle_value as *mut std::ffi::c_void);
130 info!("File handle received from service.exe: {:?}", handle_value);
131 Ok(handle)
132 }
133 Err(_) => Err("Invalid handle format in response".to_string()),
134 }
135 } else if response.starts_with("HANDLE_ERROR:") {
136 let error_msg = response.strip_prefix("HANDLE_ERROR:").unwrap_or("Unknown error");
137 Err(format!("Service.exe error: {}", error_msg))
138 } else {
139 Err("Unexpected response from service.exe".to_string())
140 }
141 }
142
143 #[cfg(not(target_os = "windows"))]
146 fn try_open_file_directly_with_handle(&self) -> Result<SafeHandle, String> {
147 use std::fs::OpenOptions;
148 use std::os::unix::io::AsRawFd;
149
150 info!("Opening persistence file directly: {}", self.file_path);
151
152 if let Some(parent) = Path::new(&self.file_path).parent() {
154 match std::fs::create_dir_all(parent) {
155 Ok(_) => {},
156 Err(e) => return Err(format!("Failed to create directory: {}", e)),
157 }
158 }
159
160
161 let file = match OpenOptions::new()
162 .read(true)
163 .write(true)
164 .create(true)
165 .open(&self.file_path) {
166 Ok(file) => file,
167 Err(e) => return Err(format!("Failed to open file: {}", e)),
168 };
169
170 let handle = file.as_raw_fd() as *mut std::ffi::c_void;
171 std::mem::forget(file); Ok(SafeHandle::new(handle))
174 }
175 #[cfg(target_os = "windows")]
176 fn try_open_file_directly_with_handle(&self) -> Result<SafeHandle, String> {
177 info!("Opening persistence file directly: {}", self.file_path);
178
179 if let Some(parent) = Path::new(&self.file_path).parent() {
181 match std::fs::create_dir_all(parent) {
182 Ok(_) => {},
183 Err(e) => return Err(format!("Failed to create directory: {}", e)),
184 }
185 }
186
187 use std::fs::OpenOptions;
188 use std::os::windows::io::AsRawHandle;
189
190 let file = match OpenOptions::new()
191 .read(true)
192 .write(true)
193 .create(true)
194 .open(&self.file_path) {
195 Ok(file) => file,
196 Err(e) => return Err(format!("Failed to open file: {}", e)),
197 };
198
199 let handle = file.as_raw_handle() as *mut std::ffi::c_void;
200 std::mem::forget(file); Ok(SafeHandle::new(handle))
203 }
204
205 async fn read_file_data(&self) -> Result<Vec<u8>, String> {
207 let file_handle = self.file_handle.lock().await;
208
209 if let Some(safe_handle) = *file_handle {
210 if safe_handle.is_valid() {
211 info!("Reading file using {} handle",
212 if self.using_service_protection { "service" } else { "direct" });
213 drop(file_handle); return self.read_with_handle(&safe_handle);
215 }
216 }
217 drop(file_handle);
218
219 info!("Reading file using fallback fs::read");
221 match std::fs::read(&self.file_path) {
222 Ok(data) => Ok(data),
223 Err(e) => Err(format!("Failed to read persistence file: {}", e)),
224 }
225 }
226
227 fn read_with_handle(&self, handle: &SafeHandle) -> Result<Vec<u8>, String> {
231 #[cfg(target_os = "windows")]
232 {
233 use winapi::um::fileapi::{ReadFile, SetFilePointer, GetFileSize};
234 use winapi::um::errhandlingapi::GetLastError;
235 use winapi::um::winbase::FILE_BEGIN;
236 use winapi::um::handleapi::INVALID_HANDLE_VALUE;
237
238 let handle_value = handle.get_handle();
239
240 if handle_value.is_null() || handle_value == INVALID_HANDLE_VALUE {
242 return Err("Invalid handle value".to_string());
243 }
244
245 info!("Windows: Using file handle: {:?}", handle_value);
246
247 let seek_result = unsafe {
249 SetFilePointer(handle_value, 0, std::ptr::null_mut(), FILE_BEGIN)
250 };
251
252 if seek_result == 0xFFFFFFFF {
253 info!("SetFilePointer failed, checking for error");
254 let error_code = unsafe { GetLastError() };
255 return Err(format!("SetFilePointer failed with error: {}", error_code));
256 }
257
258 let file_size = unsafe { GetFileSize(handle_value, std::ptr::null_mut()) };
260
261 if file_size == 0xFFFFFFFF {
262 let error_code = unsafe { GetLastError() };
263 info!("GetFileSize failed, checking for error");
264 return Err(format!("GetFileSize failed with error: {} (handle may be invalid)", error_code));
265 }
266
267 info!("File size from handle: {} bytes", file_size);
268
269 if file_size == 0 {
270 info!("File is empty, returning empty vector");
271 return Ok(Vec::new());
272 }
273
274 let mut buffer = vec![0u8; file_size as usize];
276 let mut bytes_read: u32 = 0;
277
278 info!("Reading {} bytes from file handle: {:?}", file_size, handle_value);
279
280 let result = unsafe {
281 ReadFile(
282 handle_value,
283 buffer.as_mut_ptr() as *mut std::ffi::c_void,
284 file_size,
285 &mut bytes_read,
286 std::ptr::null_mut(),
287 )
288 };
289
290 if result == 0 {
291 let error_code = unsafe { GetLastError() };
292 return Err(format!("ReadFile failed with error: {}", error_code));
293 }
294
295 buffer.truncate(bytes_read as usize);
296 info!("Successfully read {} bytes", bytes_read);
297 Ok(buffer)
298 }
299
300 #[cfg(target_os = "macos")]
301 {
302 use std::os::unix::io::FromRawFd;
303 use std::io::Read;
304
305 let fd = handle.get_handle() as i32;
306
307 if fd < 0 {
308 return Err("Invalid file descriptor".to_string());
309 }
310
311 info!("macOS: Using file descriptor: {}", fd);
312
313 let mut file = unsafe { std::fs::File::from_raw_fd(fd) };
315
316 use std::io::Seek;
318 if let Err(e) = file.seek(std::io::SeekFrom::Start(0)) {
319 return Err(format!("Failed to seek to start: {}", e));
320 }
321
322 let mut buffer = Vec::new();
324 match file.read_to_end(&mut buffer) {
325 Ok(bytes_read) => {
326 info!("macOS: Successfully read {} bytes", bytes_read);
327
328 std::mem::forget(file);
330
331 Ok(buffer)
332 }
333 Err(e) => Err(format!("Failed to read from file descriptor: {}", e))
334 }
335 }
336
337 #[cfg(target_os = "linux")]
338 {
339 use std::os::unix::io::FromRawFd;
340 use std::io::Read;
341
342 let fd = handle.get_handle() as i32;
343
344 if fd < 0 {
345 return Err("Invalid file descriptor".to_string());
346 }
347
348 info!("Linux: Using file descriptor: {}", fd);
349
350 let mut file = unsafe { std::fs::File::from_raw_fd(fd) };
352
353 use std::io::Seek;
355 if let Err(e) = file.seek(std::io::SeekFrom::Start(0)) {
356 return Err(format!("Failed to seek to start: {}", e));
357 }
358
359 let mut buffer = Vec::new();
361 match file.read_to_end(&mut buffer) {
362 Ok(bytes_read) => {
363 info!("Linux: Successfully read {} bytes", bytes_read);
364
365 std::mem::forget(file);
367
368 Ok(buffer)
369 }
370 Err(e) => Err(format!("Failed to read from file descriptor: {}", e))
371 }
372 }
373
374 #[cfg(not(any(target_os = "windows", target_os = "macos", target_os = "linux")))]
375 {
376 error!("Platform not supported for handle-based file reading");
378 Err("Platform not supported for handle-based file reading".to_string())
379 }
380 }
381 async fn try_load_data(&self) -> Result<(), String> {
383 let data = self.read_file_data().await?;
384
385 if data.is_empty() {
386 info!("Persistence file is empty, starting with empty cache");
387 return Ok(());
388 }
389
390 let decrypted_data = match crypto::decrypt(&data, PERSISTENCE_PASSWORD) {
392 Ok(decrypted) => decrypted,
393 Err(e) => {
394 error!("Failed to decrypt persistence file: {}", e);
395 info!("Starting with empty cache due to decryption failure");
396 return Ok(()); }
398 };
399
400 info!("Successfully decrypted persistence file");
401
402 let loaded_data = match serde_json::from_slice::<HashMap<String, Value>>(&decrypted_data) {
404 Ok(data) => data,
405 Err(e) => {
406 error!("Failed to parse decrypted data as JSON: {}", e);
407 info!("Starting with empty cache");
408 return Ok(()); }
410 };
411
412 let mut cache = self.cache.lock().await;
414 *cache = loaded_data;
415 let count = cache.len();
416 drop(cache);
417
418 info!("Successfully loaded {} keys from encrypted persistence file", count);
419 Ok(())
420 }
421
422 async fn save_to_file(&self) -> Result<(), String> {
424 let cache = self.cache.lock().await;
425
426 let data_bytes = match serde_json::to_vec_pretty(&*cache) {
428 Ok(bytes) => bytes,
429 Err(e) => return Err(format!("Failed to serialize cache: {}", e)),
430 };
431
432 let encrypted_data = match crypto::encrypt(&data_bytes, PERSISTENCE_PASSWORD) {
433 Ok(encrypted) => encrypted,
434 Err(e) => return Err(format!("Failed to encrypt persistence data: {}", e)),
435 };
436
437 let file_handle = self.file_handle.lock().await;
439 if let Some(safe_handle) = *file_handle {
440 if safe_handle.is_valid() {
441 drop(file_handle); match self.write_with_handle(&safe_handle, &encrypted_data).await {
443 Ok(_) => {
444 info!("Saved {} keys using {} handle",
445 cache.len(),
446 if self.using_service_protection { "service" } else { "direct" });
447 return Ok(());
448 }
449 Err(e) => {
450 error!("Failed to write with handle: {}", e);
451 }
453 }
454 }
455 }
456
457 match tokio::fs::write(&self.file_path, encrypted_data).await {
459 Ok(_) => {
460 info!("Saved {} keys using fallback write", cache.len());
461 Ok(())
462 }
463 Err(e) => Err(format!("Failed to write persistence file: {}", e)),
464 }
465 }
466
467 async fn write_with_handle(&self, handle: &SafeHandle, data: &[u8]) -> Result<(), String> {
470 #[cfg(target_os = "windows")]
471 {
472 use winapi::um::fileapi::{WriteFile, SetFilePointer, SetEndOfFile};
473 use winapi::um::errhandlingapi::GetLastError;
474 use winapi::um::winbase::FILE_BEGIN;
475
476 let handle_value = handle.get_handle();
477
478 info!("Windows: Writing {} bytes to handle: {:?}", data.len(), handle_value);
479
480 unsafe {
482 SetFilePointer(handle_value, 0, std::ptr::null_mut(), FILE_BEGIN);
483 }
484
485 let mut bytes_written: u32 = 0;
486 let result = unsafe {
487 WriteFile(
488 handle_value,
489 data.as_ptr() as *const std::ffi::c_void,
490 data.len() as u32,
491 &mut bytes_written,
492 std::ptr::null_mut(),
493 )
494 };
495
496 if result == 0 {
497 let error_code = unsafe { GetLastError() };
498 return Err(format!("WriteFile failed with error: {}", error_code));
499 }
500
501 unsafe {
503 SetEndOfFile(handle_value);
504 }
505
506 info!("Windows: Successfully wrote {} bytes", bytes_written);
507 Ok(())
508 }
509
510 #[cfg(target_os = "macos")]
511 {
512 use std::os::unix::io::FromRawFd;
513 use std::io::{Write, Seek};
514
515 let fd = handle.get_handle() as i32;
516
517 if fd < 0 {
518 return Err("Invalid file descriptor".to_string());
519 }
520
521 info!("macOS: Writing {} bytes to fd: {}", data.len(), fd);
522
523 let mut file = unsafe { std::fs::File::from_raw_fd(fd) };
525
526 if let Err(e) = file.seek(std::io::SeekFrom::Start(0)) {
528 return Err(format!("Failed to seek to start: {}", e));
529 }
530
531 if let Err(e) = file.set_len(0) {
533 return Err(format!("Failed to truncate file: {}", e));
534 }
535
536 match file.write_all(data) {
538 Ok(_) => {
539 if let Err(e) = file.flush() {
541 return Err(format!("Failed to flush file: {}", e));
542 }
543
544 info!("macOS: Successfully wrote {} bytes", data.len());
545
546 std::mem::forget(file);
548
549 Ok(())
550 }
551 Err(e) => Err(format!("Failed to write to file descriptor: {}", e))
552 }
553 }
554
555 #[cfg(target_os = "linux")]
556 {
557 use std::os::unix::io::FromRawFd;
558 use std::io::{Write, Seek};
559
560 let fd = handle.get_handle() as i32;
561
562 if fd < 0 {
563 return Err("Invalid file descriptor".to_string());
564 }
565
566 info!("Linux: Writing {} bytes to fd: {}", data.len(), fd);
567
568 let mut file = unsafe { std::fs::File::from_raw_fd(fd) };
570
571 if let Err(e) = file.seek(std::io::SeekFrom::Start(0)) {
573 return Err(format!("Failed to seek to start: {}", e));
574 }
575
576 if let Err(e) = file.set_len(0) {
577 return Err(format!("Failed to truncate file: {}", e));
578 }
579
580 match file.write_all(data) {
582 Ok(_) => {
583 if let Err(e) = file.flush() {
584 return Err(format!("Failed to flush file: {}", e));
585 }
586
587 info!("Linux: Successfully wrote {} bytes", data.len());
588
589 std::mem::forget(file);
591
592 Ok(())
593 }
594 Err(e) => Err(format!("Failed to write to file descriptor: {}", e))
595 }
596 }
597
598 #[cfg(not(any(target_os = "windows", target_os = "macos", target_os = "linux")))]
599 {
600 error!("Platform not supported for handle-based file writing");
601 Err("Platform not supported for handle-based file writing".to_string())
602 }
603 }
604
605 async fn get_value(&self, key: &str) -> Option<Value> {
607 let cache = self.cache.lock().await;
608 cache.get(key).cloned()
609 }
610
611 pub async fn set_value(&self, key: String, value: Value) -> Result<(), String> {
613 {
614 let mut cache = self.cache.lock().await;
615 cache.insert(key.clone(), value);
616 }
617
618 match self.save_to_file().await {
620 Ok(_) => {
621 info!("Set value for key: {}", key);
622 Ok(())
623 }
624 Err(e) => Err(e),
625 }
626 }
627
628 pub async fn delete_value(&self, key: &str) -> Result<bool, String> {
630 let removed = {
631 let mut cache = self.cache.lock().await;
632 cache.remove(key).is_some()
633 };
634
635 if removed {
636 match self.save_to_file().await {
638 Ok(_) => {
639 info!("Deleted value for key: {}", key);
640 Ok(true)
641 }
642 Err(e) => Err(e),
643 }
644 } else {
645 info!("Key not found for deletion: {}", key);
646 Ok(false)
647 }
648 }
649
650 async fn close_handle_with_service(&self) -> Result<(), String> {
652 let service_handle = self.file_handle.lock().await;
653 if service_handle.is_none() {
654 info!("No service handle to close");
655 return Ok(());
656 }
657 drop(service_handle); info!("Closing handle with service.exe for: {}", self.file_path);
660
661 let mut stream = match TcpStream::connect(sharing::CHANNEL_NAME) {
662 Ok(stream) => stream,
663 Err(e) => return Err(format!("Could not connect to service.exe for CLOSE_HANDLE: {}", e)),
664 };
665
666 let action = Action::CloseHandle(self.file_path.clone());
668 match stream.write_all(action.to_string().as_bytes()) {
669 Ok(_) => {},
670 Err(e) => return Err(format!("Failed to send CLOSE_HANDLE request: {}", e)),
671 }
672
673 let mut buffer = [0; 1024];
675 let bytes_read = match stream.read(&mut buffer) {
676 Ok(bytes) => bytes,
677 Err(e) => return Err(format!("Failed to read CLOSE_HANDLE response: {}", e)),
678 };
679
680 let response = String::from_utf8_lossy(&buffer[..bytes_read]);
681 info!("CLOSE_HANDLE response: {}", response);
682 Ok(())
683 }
684
685 async fn process_action(&self, action: Value, _write: WebSocketWrite) -> (i32, String) {
687 let action_obj = match action.get("ACTION").and_then(|a| a.as_object()) {
689 Some(obj) => obj,
690 None => {
691 error!("Missing or invalid ACTION object");
692 return (1, "Missing ACTION object".to_string());
693 }
694 };
695
696 let command = action_obj.get("command").and_then(|c| c.as_str()).unwrap_or("UNKNOWN");
697 let key = action_obj.get("key").and_then(|k| k.as_str()).unwrap_or("");
698 let value = action_obj.get("value");
699
700 info!("Processing persistence command: {} for key: {}", command, key);
701
702 match command.to_uppercase().as_str() {
703 "GET" => {
704 if key.is_empty() {
705 error!("Missing key parameter for GET command");
706 return (1, "Missing key parameter".to_string());
707 }
708
709 match self.get_value(key).await {
710 Some(value) => {
711 info!("Retrieved value for key: {}", key);
712 (0, serde_json::to_string(&value).unwrap_or("null".to_string()))
713 }
714 None => {
715 info!("Key not found: {}", key);
716 (0, "Key not found".to_string())
717 }
718 }
719 }
720 "SET" => {
721 if key.is_empty() {
722 error!("Missing key parameter for SET command");
723 return (1, "Missing key parameter".to_string());
724 }
725
726 let value = match value {
727 Some(v) => v.clone(),
728 None => {
729 error!("Missing value parameter for SET command");
730 return (1, "Missing value parameter".to_string());
731 }
732 };
733
734 match self.set_value(key.to_string(), value.clone()).await {
735 Ok(_) => {
736 info!("Set value for key: {}", key);
737 (0, format!("Value set successfully, size of value: {}", value.to_string().len()))
738 }
739 Err(e) => {
740 error!("Failed to set value: {}", e);
741 (1, format!("Failed to set value: {}", e))
742 }
743 }
744 }
745 "DEL" => {
746 if key.is_empty() {
747 error!("Missing key parameter for DELETE command");
748 return (1, "Missing key parameter".to_string());
749 }
750
751 match self.delete_value(key).await {
752 Ok(was_deleted) => {
753 if was_deleted {
754 info!("Deleted key: {}", key);
755 (0, "Key deleted successfully".to_string())
756 } else {
757 info!("Key not found for deletion: {}", key);
758 (0, "Key not found".to_string())
759 }
760 }
761 Err(e) => {
762 error!("Failed to delete value: {}", e);
763 (1, format!("Failed to delete value: {}", e))
764 }
765 }
766 }
767 _ => {
768 error!("Unknown persistence command: {}", command);
769 (1, "Unknown command".to_string())
770 }
771 }
772 }
773}
774
775#[async_trait]
776impl Service for PersistenceService {
777 async fn run(&self, action: Value, write: WebSocketWrite) -> (i32, String) {
779 debug!("PersistenceService: Running action: {:?}", action);
780 self.process_action(action, write).await
781 }
782
783 fn as_any(&self) -> &dyn std::any::Any {
785 self
786 }
787
788 fn stop_service(&self) {
790 info!("PersistenceService: Stopping service");
791
792 let rt = tokio::runtime::Runtime::new().unwrap();
794 rt.block_on(async {
795 if let Err(e) = self.close_handle_with_service().await {
796 error!("Failed to close handle with service: {}", e);
797 }
798 });
799
800 info!("PersistenceService stopped");
801 }
802
803 fn get_version(&self) -> String {
805 PERSISTENCE_VERSION.to_string()
806 }
807}