1use super::entry::CachedPeer;
27use serde::{Deserialize, Serialize};
28use std::collections::HashMap;
29use std::fs::{self, File, OpenOptions};
30use std::io;
31use std::path::{Path, PathBuf};
32use std::time::SystemTime;
33use tracing::{debug, info, warn};
34use zeroize::Zeroize;
35
36#[derive(Debug, Serialize, Deserialize)]
38pub struct CacheData {
39 pub version: u32,
41
42 pub instance_id: String,
44
45 pub timestamp: u64,
47
48 #[serde(with = "peer_map_serde")]
50 pub peers: HashMap<[u8; 32], CachedPeer>,
51
52 pub checksum: u64,
54}
55
56impl CacheData {
57 pub const CURRENT_VERSION: u32 = 1;
59
60 pub fn new(instance_id: String) -> Self {
62 Self {
63 version: Self::CURRENT_VERSION,
64 instance_id,
65 timestamp: SystemTime::now()
66 .duration_since(SystemTime::UNIX_EPOCH)
67 .map(|d| d.as_secs())
68 .unwrap_or(0),
69 peers: HashMap::new(),
70 checksum: 0,
71 }
72 }
73
74 pub fn calculate_checksum(&self) -> u64 {
76 use std::collections::hash_map::DefaultHasher;
77 use std::hash::{Hash, Hasher};
78
79 let mut hasher = DefaultHasher::new();
80 self.version.hash(&mut hasher);
81 self.peers.len().hash(&mut hasher);
82
83 let mut ids: Vec<_> = self.peers.keys().collect();
85 ids.sort();
86 for id in ids {
87 id.hash(&mut hasher);
88 }
89
90 hasher.finish()
91 }
92
93 pub fn finalize(&mut self) {
95 self.timestamp = SystemTime::now()
96 .duration_since(SystemTime::UNIX_EPOCH)
97 .map(|d| d.as_secs())
98 .unwrap_or(0);
99 self.checksum = self.calculate_checksum();
100 }
101
102 pub fn verify(&self) -> bool {
104 self.checksum == self.calculate_checksum()
105 }
106}
107
108#[derive(Debug)]
110pub struct CachePersistence {
111 cache_file: PathBuf,
112 lock_file: PathBuf,
113 instance_id: String,
114 enable_locking: bool,
115}
116
117impl CachePersistence {
118 pub fn new(cache_dir: &Path, enable_locking: bool) -> io::Result<Self> {
120 Self::new_with_filename(cache_dir, "bootstrap_cache.json", enable_locking)
121 }
122
123 pub fn new_with_filename(
125 cache_dir: &Path,
126 filename: &str,
127 enable_locking: bool,
128 ) -> io::Result<Self> {
129 fs::create_dir_all(cache_dir)?;
130
131 let cache_file = cache_dir.join(filename);
132 let lock_file = cache_dir.join(format!("{}.lock", filename));
133 let instance_id = generate_instance_id();
134
135 Ok(Self {
136 cache_file,
137 lock_file,
138 instance_id,
139 enable_locking,
140 })
141 }
142
143 pub fn load(&self) -> io::Result<CacheData> {
145 if !self.cache_file.exists() {
146 debug!("No existing cache file, starting fresh");
147 return Ok(CacheData::new(self.instance_id.clone()));
148 }
149
150 let _lock = if self.enable_locking {
151 Some(self.acquire_shared_lock()?)
152 } else {
153 None
154 };
155
156 let data = fs::read_to_string(&self.cache_file)?;
157 let cache: CacheData = serde_json::from_str(&data)
158 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
159
160 if !cache.verify() {
162 warn!("Cache checksum mismatch, data may be corrupted");
163 return Ok(CacheData::new(self.instance_id.clone()));
165 }
166
167 if cache.version < CacheData::CURRENT_VERSION {
169 info!(
170 "Migrating cache from version {} to {}",
171 cache.version,
172 CacheData::CURRENT_VERSION
173 );
174 }
176
177 info!("Loaded {} peers from cache", cache.peers.len());
178 Ok(cache)
179 }
180
181 pub fn save(&self, cache: &mut CacheData) -> io::Result<()> {
183 let _lock = if self.enable_locking {
184 Some(self.acquire_exclusive_lock()?)
185 } else {
186 None
187 };
188
189 cache.instance_id.clone_from(&self.instance_id);
190 cache.finalize();
191
192 let temp_file = self.cache_file.with_extension("tmp");
194 let data = serde_json::to_string_pretty(cache)
195 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
196
197 fs::write(&temp_file, data)?;
198
199 fs::rename(&temp_file, &self.cache_file)?;
201
202 debug!("Saved {} peers to cache", cache.peers.len());
203 Ok(())
204 }
205
206 #[allow(dead_code)]
208 pub fn merge(&self, cache: &mut CacheData, other_path: &Path) -> io::Result<usize> {
209 let other_data = fs::read_to_string(other_path)?;
210 let other: CacheData = serde_json::from_str(&other_data)
211 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
212
213 if !other.verify() {
214 warn!("Merge source has invalid checksum, skipping");
215 return Ok(0);
216 }
217
218 let mut merged_count = 0;
219 for (id, peer) in other.peers {
220 cache
221 .peers
222 .entry(id)
223 .and_modify(|existing| {
224 if peer.last_seen > existing.last_seen {
226 *existing = peer.clone();
227 merged_count += 1;
228 }
229 })
230 .or_insert_with(|| {
231 merged_count += 1;
232 peer
233 });
234 }
235
236 info!(
237 "Merged {} peers from {}",
238 merged_count,
239 other_path.display()
240 );
241 Ok(merged_count)
242 }
243
244 #[allow(dead_code)]
246 pub fn cache_file(&self) -> &Path {
247 &self.cache_file
248 }
249
250 #[cfg(unix)]
251 fn acquire_shared_lock(&self) -> io::Result<FileLock> {
252 use std::os::unix::io::AsRawFd;
253
254 let file = OpenOptions::new()
255 .read(true)
256 .write(true)
257 .create(true)
258 .truncate(false)
259 .open(&self.lock_file)?;
260
261 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_SH | libc::LOCK_NB) };
263
264 if result != 0 {
265 let err = io::Error::last_os_error();
266 if err.kind() == io::ErrorKind::WouldBlock {
268 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_SH) };
270 if result != 0 {
271 return Err(io::Error::last_os_error());
272 }
273 } else {
274 return Err(err);
275 }
276 }
277
278 Ok(FileLock { file })
279 }
280
281 #[cfg(unix)]
282 fn acquire_exclusive_lock(&self) -> io::Result<FileLock> {
283 use std::os::unix::io::AsRawFd;
284
285 let file = OpenOptions::new()
286 .read(true)
287 .write(true)
288 .create(true)
289 .truncate(false)
290 .open(&self.lock_file)?;
291
292 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
294
295 if result != 0 {
296 let err = io::Error::last_os_error();
297 if err.kind() == io::ErrorKind::WouldBlock {
299 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX) };
300 if result != 0 {
301 return Err(io::Error::last_os_error());
302 }
303 } else {
304 return Err(err);
305 }
306 }
307
308 Ok(FileLock { file })
309 }
310
311 #[cfg(not(unix))]
312 fn acquire_shared_lock(&self) -> io::Result<FileLock> {
313 let file = OpenOptions::new()
315 .read(true)
316 .write(true)
317 .create(true)
318 .truncate(false)
319 .open(&self.lock_file)?;
320 Ok(FileLock { file })
321 }
322
323 #[cfg(not(unix))]
324 fn acquire_exclusive_lock(&self) -> io::Result<FileLock> {
325 let file = OpenOptions::new()
326 .read(true)
327 .write(true)
328 .create(true)
329 .truncate(false)
330 .open(&self.lock_file)?;
331 Ok(FileLock { file })
332 }
333}
334
335struct FileLock {
337 #[allow(dead_code)]
338 file: File,
339}
340
341#[cfg(unix)]
342impl Drop for FileLock {
343 fn drop(&mut self) {
344 use std::os::unix::io::AsRawFd;
345 unsafe {
346 libc::flock(self.file.as_raw_fd(), libc::LOCK_UN);
347 }
348 }
349}
350
351const ENCRYPTED_CACHE_VERSION: u8 = 1;
357
358pub struct EncryptedCachePersistence {
363 inner: CachePersistence,
364 encryption_key: [u8; 32],
365}
366
367impl EncryptedCachePersistence {
368 pub fn new(
375 cache_dir: &Path,
376 enable_locking: bool,
377 encryption_key: [u8; 32],
378 ) -> io::Result<Self> {
379 let inner =
380 CachePersistence::new_with_filename(cache_dir, "bootstrap_cache.enc", enable_locking)?;
381 Ok(Self {
382 inner,
383 encryption_key,
384 })
385 }
386
387 pub fn load(&self) -> io::Result<CacheData> {
389 if !self.inner.cache_file.exists() {
390 debug!("No existing encrypted cache file, starting fresh");
391 return Ok(CacheData::new(self.inner.instance_id.clone()));
392 }
393
394 let _lock = if self.inner.enable_locking {
395 Some(self.inner.acquire_shared_lock()?)
396 } else {
397 None
398 };
399
400 let encrypted_data = fs::read(&self.inner.cache_file)?;
401 let json_data = self.decrypt(&encrypted_data)?;
402
403 let cache: CacheData = serde_json::from_slice(&json_data)
404 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
405
406 if !cache.verify() {
407 warn!("Encrypted cache checksum mismatch, data may be corrupted");
408 return Ok(CacheData::new(self.inner.instance_id.clone()));
409 }
410
411 info!("Loaded {} peers from encrypted cache", cache.peers.len());
412 Ok(cache)
413 }
414
415 pub fn save(&self, cache: &mut CacheData) -> io::Result<()> {
417 let _lock = if self.inner.enable_locking {
418 Some(self.inner.acquire_exclusive_lock()?)
419 } else {
420 None
421 };
422
423 cache.instance_id.clone_from(&self.inner.instance_id);
424 cache.finalize();
425
426 let json_data =
427 serde_json::to_vec(cache).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
428
429 let encrypted_data = self.encrypt(&json_data)?;
430
431 let temp_file = self.inner.cache_file.with_extension("tmp");
433 fs::write(&temp_file, &encrypted_data)?;
434 fs::rename(&temp_file, &self.inner.cache_file)?;
435
436 debug!("Saved {} peers to encrypted cache", cache.peers.len());
437 Ok(())
438 }
439
440 pub fn exists(&self) -> bool {
442 self.inner.cache_file.exists()
443 }
444
445 fn encrypt(&self, plaintext: &[u8]) -> io::Result<Vec<u8>> {
447 use aws_lc_rs::aead::{
448 self, Aad, BoundKey, CHACHA20_POLY1305, Nonce, NonceSequence, UnboundKey,
449 };
450
451 let mut nonce_bytes = [0u8; 12];
453 aws_lc_rs::rand::fill(&mut nonce_bytes)
454 .map_err(|e| io::Error::other(format!("RNG failed: {e}")))?;
455
456 let unbound_key = UnboundKey::new(&CHACHA20_POLY1305, &self.encryption_key)
458 .map_err(|e| io::Error::other(format!("Key creation failed: {e}")))?;
459
460 struct SingleNonce(Option<[u8; 12]>);
461 impl NonceSequence for SingleNonce {
462 fn advance(&mut self) -> Result<Nonce, aws_lc_rs::error::Unspecified> {
463 self.0
464 .take()
465 .map(Nonce::assume_unique_for_key)
466 .ok_or(aws_lc_rs::error::Unspecified)
467 }
468 }
469
470 let mut sealing_key = aead::SealingKey::new(unbound_key, SingleNonce(Some(nonce_bytes)));
471
472 let mut in_out = plaintext.to_vec();
474 sealing_key
475 .seal_in_place_append_tag(Aad::empty(), &mut in_out)
476 .map_err(|e| io::Error::other(format!("Encryption failed: {e}")))?;
477
478 let mut result = Vec::with_capacity(1 + 12 + in_out.len());
480 result.push(ENCRYPTED_CACHE_VERSION);
481 result.extend_from_slice(&nonce_bytes);
482 result.extend_from_slice(&in_out);
483 Ok(result)
484 }
485
486 fn decrypt(&self, ciphertext: &[u8]) -> io::Result<Vec<u8>> {
488 use aws_lc_rs::aead::{
489 self, Aad, BoundKey, CHACHA20_POLY1305, Nonce, NonceSequence, UnboundKey,
490 };
491
492 if ciphertext.len() < 1 + 12 + 16 {
493 return Err(io::Error::new(
494 io::ErrorKind::InvalidData,
495 "Ciphertext too short",
496 ));
497 }
498
499 let version = ciphertext[0];
500 if version != ENCRYPTED_CACHE_VERSION {
501 return Err(io::Error::new(
502 io::ErrorKind::InvalidData,
503 format!("Unsupported encrypted cache version: {version}"),
504 ));
505 }
506
507 let nonce_bytes: [u8; 12] = ciphertext[1..13]
508 .try_into()
509 .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid nonce"))?;
510
511 let unbound_key = UnboundKey::new(&CHACHA20_POLY1305, &self.encryption_key)
513 .map_err(|e| io::Error::other(format!("Key creation failed: {e}")))?;
514
515 struct SingleNonce(Option<[u8; 12]>);
516 impl NonceSequence for SingleNonce {
517 fn advance(&mut self) -> Result<Nonce, aws_lc_rs::error::Unspecified> {
518 self.0
519 .take()
520 .map(Nonce::assume_unique_for_key)
521 .ok_or(aws_lc_rs::error::Unspecified)
522 }
523 }
524
525 let mut opening_key = aead::OpeningKey::new(unbound_key, SingleNonce(Some(nonce_bytes)));
526
527 let mut in_out = ciphertext[13..].to_vec();
529 let plaintext = opening_key
530 .open_in_place(Aad::empty(), &mut in_out)
531 .map_err(|_| {
532 io::Error::new(
533 io::ErrorKind::InvalidData,
534 "Decryption failed - wrong key or corrupted",
535 )
536 })?;
537
538 Ok(plaintext.to_vec())
539 }
540}
541
542impl Drop for EncryptedCachePersistence {
543 fn drop(&mut self) {
544 self.encryption_key.zeroize();
545 }
546}
547
548pub(super) fn generate_instance_id() -> String {
549 format!(
550 "{}_{:x}",
551 std::process::id(),
552 SystemTime::now()
553 .duration_since(SystemTime::UNIX_EPOCH)
554 .map(|d| d.as_millis())
555 .unwrap_or(0)
556 )
557}
558
559mod peer_map_serde {
561 use super::*;
562 use serde::ser::SerializeMap;
563
564 pub fn serialize<S>(
565 map: &HashMap<[u8; 32], CachedPeer>,
566 serializer: S,
567 ) -> Result<S::Ok, S::Error>
568 where
569 S: serde::Serializer,
570 {
571 let mut map_ser = serializer.serialize_map(Some(map.len()))?;
572 for (k, v) in map {
573 map_ser.serialize_entry(&hex::encode(k), v)?;
574 }
575 map_ser.end()
576 }
577
578 pub fn deserialize<'de, D>(deserializer: D) -> Result<HashMap<[u8; 32], CachedPeer>, D::Error>
579 where
580 D: serde::Deserializer<'de>,
581 {
582 use serde::de::MapAccess;
583
584 struct MapVisitor;
585
586 impl<'de> serde::de::Visitor<'de> for MapVisitor {
587 type Value = HashMap<[u8; 32], CachedPeer>;
588
589 fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
590 formatter.write_str("a map with hex-encoded 32-byte keys")
591 }
592
593 fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
594 where
595 M: MapAccess<'de>,
596 {
597 let mut map = HashMap::new();
598 while let Some((key, value)) = access.next_entry::<String, CachedPeer>()? {
599 let bytes = hex::decode(&key).map_err(serde::de::Error::custom)?;
600 if bytes.len() != 32 {
601 return Err(serde::de::Error::custom("key must be 32 bytes"));
602 }
603 let mut arr = [0u8; 32];
604 arr.copy_from_slice(&bytes);
605 map.insert(arr, value);
606 }
607 Ok(map)
608 }
609 }
610
611 deserializer.deserialize_map(MapVisitor)
612 }
613}
614
615#[cfg(test)]
616mod tests {
617 use super::*;
618 use crate::PeerId;
619 use crate::bootstrap_cache::entry::PeerSource;
620 use tempfile::TempDir;
621
622 #[test]
623 fn test_cache_data_new() {
624 let data = CacheData::new("test_instance".to_string());
625 assert_eq!(data.version, CacheData::CURRENT_VERSION);
626 assert_eq!(data.instance_id, "test_instance");
627 assert!(data.peers.is_empty());
628 }
629
630 #[test]
631 fn test_checksum() {
632 let mut data = CacheData::new("test".to_string());
633 data.finalize();
634
635 let checksum1 = data.checksum;
636 assert!(data.verify());
637
638 let peer = CachedPeer::new(
640 PeerId([1u8; 32]),
641 vec!["127.0.0.1:9000".parse().unwrap()],
642 PeerSource::Seed,
643 );
644 data.peers.insert(peer.peer_id.0, peer);
645 data.finalize();
646
647 let checksum2 = data.checksum;
648 assert_ne!(checksum1, checksum2);
649 assert!(data.verify());
650 }
651
652 #[test]
653 fn test_persistence_load_save() {
654 let temp_dir = TempDir::new().unwrap();
655 let persistence = CachePersistence::new(temp_dir.path(), false).unwrap();
656
657 let mut data = CacheData::new("test".to_string());
659 let peer = CachedPeer::new(
660 PeerId([42u8; 32]),
661 vec!["127.0.0.1:9000".parse().unwrap()],
662 PeerSource::Seed,
663 );
664 data.peers.insert(peer.peer_id.0, peer);
665 persistence.save(&mut data).unwrap();
666
667 let loaded = persistence.load().unwrap();
669 assert_eq!(loaded.peers.len(), 1);
670 assert!(loaded.peers.contains_key(&[42u8; 32]));
671 }
672
673 #[test]
674 fn test_persistence_no_file() {
675 let temp_dir = TempDir::new().unwrap();
676 let persistence = CachePersistence::new(temp_dir.path(), false).unwrap();
677
678 let data = persistence.load().unwrap();
680 assert!(data.peers.is_empty());
681 }
682
683 #[test]
684 fn test_merge() {
685 let temp_dir = TempDir::new().unwrap();
686 let persistence = CachePersistence::new(temp_dir.path(), false).unwrap();
687
688 let mut data1 = CacheData::new("first".to_string());
690 let peer1 = CachedPeer::new(
691 PeerId([1u8; 32]),
692 vec!["127.0.0.1:9001".parse().unwrap()],
693 PeerSource::Seed,
694 );
695 data1.peers.insert(peer1.peer_id.0, peer1);
696 persistence.save(&mut data1).unwrap();
697
698 let other_path = temp_dir.path().join("other_cache.json");
700 let mut data2 = CacheData::new("second".to_string());
701 let peer2 = CachedPeer::new(
702 PeerId([2u8; 32]),
703 vec!["127.0.0.1:9002".parse().unwrap()],
704 PeerSource::Seed,
705 );
706 data2.peers.insert(peer2.peer_id.0, peer2);
707 data2.finalize();
708 let json = serde_json::to_string(&data2).unwrap();
709 fs::write(&other_path, json).unwrap();
710
711 let merged = persistence.merge(&mut data1, &other_path).unwrap();
713 assert_eq!(merged, 1);
714 assert_eq!(data1.peers.len(), 2);
715 }
716
717 #[test]
722 fn test_encrypted_persistence_roundtrip() {
723 let temp_dir = TempDir::new().unwrap();
724 let key = [0x42u8; 32];
725 let persistence = EncryptedCachePersistence::new(temp_dir.path(), false, key).unwrap();
726
727 let mut data = CacheData::new("test".to_string());
729 let peer = CachedPeer::new(
730 PeerId([42u8; 32]),
731 vec!["127.0.0.1:9000".parse().unwrap()],
732 PeerSource::Seed,
733 );
734 data.peers.insert(peer.peer_id.0, peer);
735 persistence.save(&mut data).unwrap();
736
737 let loaded = persistence.load().unwrap();
739 assert_eq!(loaded.peers.len(), 1);
740 assert!(loaded.peers.contains_key(&[42u8; 32]));
741 }
742
743 #[test]
744 fn test_encrypted_persistence_wrong_key() {
745 let temp_dir = TempDir::new().unwrap();
746 let key1 = [0x42u8; 32];
747 let key2 = [0x43u8; 32];
748
749 let persistence1 = EncryptedCachePersistence::new(temp_dir.path(), false, key1).unwrap();
751 let mut data = CacheData::new("test".to_string());
752 let peer = CachedPeer::new(
753 PeerId([1u8; 32]),
754 vec!["127.0.0.1:9000".parse().unwrap()],
755 PeerSource::Seed,
756 );
757 data.peers.insert(peer.peer_id.0, peer);
758 persistence1.save(&mut data).unwrap();
759
760 let persistence2 = EncryptedCachePersistence::new(temp_dir.path(), false, key2).unwrap();
762 let result = persistence2.load();
763 assert!(result.is_err());
764 }
765
766 #[test]
767 fn test_encrypted_persistence_no_file() {
768 let temp_dir = TempDir::new().unwrap();
769 let key = [0x42u8; 32];
770 let persistence = EncryptedCachePersistence::new(temp_dir.path(), false, key).unwrap();
771
772 let data = persistence.load().unwrap();
774 assert!(data.peers.is_empty());
775 }
776
777 #[test]
778 fn test_encrypted_persistence_exists() {
779 let temp_dir = TempDir::new().unwrap();
780 let key = [0x42u8; 32];
781 let persistence = EncryptedCachePersistence::new(temp_dir.path(), false, key).unwrap();
782
783 assert!(!persistence.exists());
784
785 let mut data = CacheData::new("test".to_string());
786 persistence.save(&mut data).unwrap();
787
788 assert!(persistence.exists());
789 }
790
791 #[test]
792 fn test_encrypt_decrypt_roundtrip() {
793 let temp_dir = TempDir::new().unwrap();
794 let key = [0xAB; 32];
795 let persistence = EncryptedCachePersistence::new(temp_dir.path(), false, key).unwrap();
796
797 let plaintext = b"Hello, encrypted cache!";
798 let ciphertext = persistence.encrypt(plaintext).unwrap();
799
800 assert!(ciphertext.len() > plaintext.len());
802
803 let decrypted = persistence.decrypt(&ciphertext).unwrap();
804 assert_eq!(decrypted, plaintext);
805 }
806}