p2p_chat/storage/
known_mailboxes.rs

1//! This module defines the storage interface and implementation for managing
2//! known mailbox nodes, including their performance statistics.
3use crate::crypto::StorageEncryption;
4use anyhow::Result;
5use async_trait::async_trait;
6use libp2p::PeerId;
7use serde::{Deserialize, Serialize};
8use sled::Db;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11/// Represents a known mailbox node with its associated performance statistics.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct KnownMailbox {
14    /// The `PeerId` of the mailbox node.
15    pub peer_id: PeerId,
16    /// The timestamp of when the mailbox was last seen.
17    pub last_seen: i64,
18    /// The number of successful interactions with this mailbox.
19    pub success_count: u32,
20    /// The number of failed interactions with this mailbox.
21    pub failure_count: u32,
22}
23
24impl KnownMailbox {
25    /// Creates a new `KnownMailbox` entry.
26    ///
27    /// # Arguments
28    ///
29    /// * `peer_id` - The `PeerId` of the new mailbox.
30    pub fn new(peer_id: PeerId) -> Self {
31        Self {
32            peer_id,
33            last_seen: current_timestamp(),
34            success_count: 0,
35            failure_count: 0,
36        }
37    }
38
39    /// Updates the `last_seen` timestamp to the current time.
40    pub fn touch(&mut self) {
41        self.last_seen = current_timestamp();
42    }
43}
44
45/// Returns the current Unix timestamp in milliseconds.
46fn current_timestamp() -> i64 {
47    SystemTime::now()
48        .duration_since(UNIX_EPOCH)
49        .unwrap()
50        .as_secs() as i64
51}
52
53/// A trait for managing known mailbox nodes.
54#[async_trait]
55pub trait KnownMailboxesStore: Send + Sync {
56    /// Adds a new `KnownMailbox` to the store.
57    async fn add_mailbox(&self, mailbox: KnownMailbox) -> Result<()>;
58    /// Retrieves a `KnownMailbox` by its `PeerId`.
59    async fn get_mailbox(&self, peer_id: &PeerId) -> Result<Option<KnownMailbox>>;
60    /// Lists all known mailboxes.
61    async fn list_mailboxes(&self) -> Result<Vec<KnownMailbox>>;
62    /// Removes a `KnownMailbox` from the store.
63    async fn remove_mailbox(&self, peer_id: &PeerId) -> Result<()>;
64    /// Increments the success count for a mailbox.
65    async fn increment_success(&self, peer_id: &PeerId) -> Result<()>;
66    /// Increments the failure count for a mailbox.
67    async fn increment_failure(&self, peer_id: &PeerId) -> Result<()>;
68}
69
70/// A `KnownMailboxesStore` implementation using `sled` for storage.
71pub struct SledKnownMailboxesStore {
72    tree: sled::Tree,
73    encryption: Option<StorageEncryption>,
74}
75
76impl SledKnownMailboxesStore {
77    /// Creates a new `SledKnownMailboxesStore`.
78    ///
79    /// # Arguments
80    ///
81    /// * `db` - The `sled::Db` instance to use for storage.
82    /// * `encryption` - Optional `StorageEncryption` for encrypting data.
83    ///
84    /// # Errors
85    ///
86    /// Returns an error if the underlying `sled` tree cannot be opened.
87    pub fn new(db: Db, encryption: Option<StorageEncryption>) -> Result<Self> {
88        let tree = db.open_tree("known_mailboxes")?;
89        Ok(Self { tree, encryption })
90    }
91
92    /// Serializes a `KnownMailbox` and optionally encrypts it.
93    fn serialize_mailbox(&self, mailbox: &KnownMailbox) -> Result<Vec<u8>> {
94        let serialized = serde_json::to_vec(mailbox)?;
95
96        if let Some(ref encryption) = self.encryption {
97            encryption.encrypt_value(&serialized)
98        } else {
99            Ok(serialized)
100        }
101    }
102
103    /// Deserializes a `KnownMailbox` and optionally decrypts it.
104    fn deserialize_mailbox(&self, data: &[u8]) -> Result<KnownMailbox> {
105        let decrypted = if let Some(ref encryption) = self.encryption {
106            encryption.decrypt_value(data)?
107        } else {
108            data.to_vec()
109        };
110
111        Ok(serde_json::from_slice(&decrypted)?)
112    }
113}
114
115#[async_trait]
116impl KnownMailboxesStore for SledKnownMailboxesStore {
117    async fn add_mailbox(&self, mailbox: KnownMailbox) -> Result<()> {
118        let key = mailbox.peer_id.to_bytes();
119        let value = self.serialize_mailbox(&mailbox)?;
120        self.tree.insert(key, value)?;
121        self.tree.flush_async().await?;
122        Ok(())
123    }
124
125    async fn get_mailbox(&self, peer_id: &PeerId) -> Result<Option<KnownMailbox>> {
126        let key = peer_id.to_bytes();
127        match self.tree.get(key)? {
128            Some(data) => Ok(Some(self.deserialize_mailbox(&data)?)),
129            None => Ok(None),
130        }
131    }
132
133    async fn list_mailboxes(&self) -> Result<Vec<KnownMailbox>> {
134        let mut mailboxes = Vec::new();
135
136        for result in self.tree.iter() {
137            let (_key, value) = result?;
138            mailboxes.push(self.deserialize_mailbox(&value)?);
139        }
140
141        Ok(mailboxes)
142    }
143
144    async fn remove_mailbox(&self, peer_id: &PeerId) -> Result<()> {
145        let key = peer_id.to_bytes();
146        self.tree.remove(key)?;
147        self.tree.flush_async().await?;
148        Ok(())
149    }
150
151    async fn increment_success(&self, peer_id: &PeerId) -> Result<()> {
152        if let Some(mut mailbox) = self.get_mailbox(peer_id).await? {
153            mailbox.success_count += 1;
154            mailbox.failure_count = 0; // Reset consecutive failures
155            mailbox.touch();
156            self.add_mailbox(mailbox).await?;
157        }
158        Ok(())
159    }
160
161    async fn increment_failure(&self, peer_id: &PeerId) -> Result<()> {
162        if let Some(mut mailbox) = self.get_mailbox(peer_id).await? {
163            mailbox.failure_count += 1;
164            mailbox.touch();
165            self.add_mailbox(mailbox).await?;
166        }
167        Ok(())
168    }
169}
p2p_chat/storage/known_mailboxes.rs

p2p_chat/storage/
known_mailboxes.rs
