1use crate::dht::{DHT, DhtKey};
20use crate::identity::enhanced::EnhancedIdentity;
21use aes_gcm::{
22 Aes256Gcm, Key as AesKey, Nonce,
23 aead::{Aead, KeyInit},
24};
25use serde::{Deserialize, Serialize};
26use sha2::{Digest, Sha256};
27use std::time::{Duration, SystemTime};
28use thiserror::Error;
29
30#[derive(Debug, Error)]
32pub enum StorageError {
33 #[error("DHT error: {0}")]
34 DhtError(String),
35
36 #[error("Encryption error: {0}")]
37 EncryptionError(String),
38
39 #[error("Serialization error: {0}")]
40 SerializationError(#[from] bincode::Error),
41
42 #[error("Key not found: {0}")]
43 KeyNotFound(String),
44
45 #[error("Invalid data format")]
46 InvalidFormat,
47}
48
49type Result<T> = std::result::Result<T, StorageError>;
50
51pub mod keys {
53 pub fn profile(user_id: &str) -> String {
55 format!("profile:{user_id}")
56 }
57
58 pub fn devices(user_id: &str) -> String {
60 format!("devices:{user_id}")
61 }
62
63 pub fn chat_channel(channel_id: &str) -> String {
65 format!("chat:channel:{channel_id}")
66 }
67
68 pub fn chat_message(channel_id: &str, msg_id: &str) -> String {
70 format!("chat:msg:{channel_id}:{msg_id}")
71 }
72
73 pub fn chat_index(channel_id: &str, timestamp: u64) -> String {
75 format!("chat:idx:{channel_id}:{timestamp}")
76 }
77
78 pub fn discuss_topic(topic_id: &str) -> String {
80 format!("discuss:topic:{topic_id}")
81 }
82
83 pub fn discuss_reply(topic_id: &str, reply_id: &str) -> String {
85 format!("discuss:reply:{topic_id}:{reply_id}")
86 }
87
88 pub fn project(project_id: &str) -> String {
90 format!("project:{project_id}")
91 }
92
93 pub fn document_meta(doc_id: &str) -> String {
95 format!("doc:meta:{doc_id}")
96 }
97
98 pub fn file_chunk(file_id: &str, chunk_num: u32) -> String {
100 format!("file:chunk:{file_id}:{chunk_num:08}")
101 }
102
103 pub fn organization(org_id: &str) -> String {
105 format!("org:{org_id}")
106 }
107
108 pub fn public_channel_list() -> String {
110 "public:channels".to_string()
111 }
112
113 pub fn user_channels(user_id: &str) -> String {
115 format!("user:channels:{user_id}")
116 }
117}
118
119pub mod ttl {
121 use std::time::Duration;
122
123 pub const PROFILE: Duration = Duration::from_secs(365 * 24 * 60 * 60); pub const MESSAGE: Duration = Duration::from_secs(90 * 24 * 60 * 60); pub const FILE_CHUNK: Duration = Duration::from_secs(365 * 24 * 60 * 60); pub const TEMP: Duration = Duration::from_secs(24 * 60 * 60); pub const PRESENCE: Duration = Duration::from_secs(5 * 60); }
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct EncryptedData {
142 pub ciphertext: Vec<u8>,
144
145 pub nonce: Vec<u8>,
147
148 pub key_id: String,
150
151 pub timestamp: SystemTime,
153
154 pub metadata: Option<serde_json::Value>,
156}
157
158pub struct StorageManager {
160 dht: DHT,
162
163 master_key: [u8; 32],
165}
166
167impl StorageManager {
168 pub fn new(dht: DHT, identity: &EnhancedIdentity) -> Result<Self> {
170 let mut hasher = Sha256::new();
172 hasher.update(identity.base_identity.user_id.as_bytes()); let master_key: [u8; 32] = hasher.finalize().into();
174
175 Ok(Self { dht, master_key })
176 }
177
178 pub async fn store_encrypted<T: Serialize>(
180 &mut self,
181 key: &str,
182 data: &T,
183 _ttl: Duration,
184 metadata: Option<serde_json::Value>,
185 ) -> Result<()> {
186 let plaintext = bincode::serialize(data)?;
188
189 let encrypted = self.encrypt(&plaintext)?;
191
192 let wrapper = EncryptedData {
194 ciphertext: encrypted.0,
195 nonce: encrypted.1.to_vec(),
196 key_id: "v1".to_string(),
197 timestamp: SystemTime::now(),
198 metadata,
199 };
200
201 let wrapper_bytes = bincode::serialize(&wrapper)?;
203
204 let hash = blake3::hash(key.as_bytes());
206 let dht_key = *hash.as_bytes();
207
208 self.dht
209 .store(&DhtKey::from_bytes(dht_key), wrapper_bytes)
210 .await
211 .map_err(|e| StorageError::DhtError(e.to_string()))?;
212
213 Ok(())
214 }
215
216 pub async fn get_encrypted<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T> {
218 let hash = blake3::hash(key.as_bytes());
220 let dht_key = *hash.as_bytes();
221 let value = self
222 .dht
223 .retrieve(&DhtKey::from_bytes(dht_key))
224 .await
225 .map_err(|e| StorageError::DhtError(e.to_string()))?
226 .ok_or_else(|| StorageError::KeyNotFound(key.to_string()))?;
227
228 let wrapper: EncryptedData = bincode::deserialize(&value)?;
230
231 let plaintext = self.decrypt(&wrapper.ciphertext, &wrapper.nonce)?;
233
234 let data = bincode::deserialize(&plaintext)?;
236
237 Ok(data)
238 }
239
240 pub async fn store_public<T: Serialize>(
242 &mut self,
243 key: &str,
244 data: &T,
245 _ttl: Duration,
246 ) -> Result<()> {
247 let value = bincode::serialize(data)?;
248
249 let hash = blake3::hash(key.as_bytes());
250 let dht_key = *hash.as_bytes();
251
252 self.dht
253 .store(&DhtKey::from_bytes(dht_key), value)
254 .await
255 .map_err(|e| StorageError::DhtError(e.to_string()))?;
256
257 Ok(())
258 }
259
260 pub async fn get_public<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T> {
262 let hash = blake3::hash(key.as_bytes());
263 let dht_key = *hash.as_bytes();
264 let value = self
265 .dht
266 .retrieve(&DhtKey::from_bytes(dht_key))
267 .await
268 .map_err(|e| StorageError::DhtError(e.to_string()))?
269 .ok_or_else(|| StorageError::KeyNotFound(key.to_string()))?;
270
271 let data = bincode::deserialize(&value)?;
272 Ok(data)
273 }
274
275 pub async fn delete(&mut self, key: &str) -> Result<()> {
277 let hash = blake3::hash(key.as_bytes());
279 let dht_key = *hash.as_bytes();
280 self.dht
281 .store(&DhtKey::from_bytes(dht_key), vec![])
282 .await
283 .map_err(|e| StorageError::DhtError(e.to_string()))?;
284 Ok(())
285 }
286
287 pub async fn list_keys(&self, _prefix: &str) -> Result<Vec<String>> {
289 Ok(vec![])
292 }
293
294 fn encrypt(&self, plaintext: &[u8]) -> Result<(Vec<u8>, [u8; 12])> {
296 let cipher = Aes256Gcm::new(AesKey::<Aes256Gcm>::from_slice(&self.master_key));
297
298 let nonce_bytes = rand::random::<[u8; 12]>();
300 let nonce = Nonce::from_slice(&nonce_bytes);
301
302 let ciphertext = cipher
303 .encrypt(nonce, plaintext)
304 .map_err(|e| StorageError::EncryptionError(e.to_string()))?;
305
306 Ok((ciphertext, nonce_bytes))
307 }
308
309 fn decrypt(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>> {
311 let cipher = Aes256Gcm::new(AesKey::<Aes256Gcm>::from_slice(&self.master_key));
312 let nonce = Nonce::from_slice(nonce);
313
314 let plaintext = cipher
315 .decrypt(nonce, ciphertext)
316 .map_err(|e| StorageError::EncryptionError(e.to_string()))?;
317
318 Ok(plaintext)
319 }
320}
321
322pub struct FileChunker {
400 chunk_size: usize,
401}
402
403impl FileChunker {
404 pub fn new(chunk_size: usize) -> Self {
406 Self { chunk_size }
407 }
408
409 pub fn chunk_file(&self, data: &[u8]) -> Vec<Vec<u8>> {
411 data.chunks(self.chunk_size)
412 .map(|chunk| chunk.to_vec())
413 .collect()
414 }
415
416 pub async fn store_file(
418 &self,
419 storage: &mut StorageManager,
420 file_id: &str,
421 data: &[u8],
422 metadata: FileMetadata,
423 ) -> Result<()> {
424 let chunks = self.chunk_file(data);
425 let total_chunks = chunks.len() as u32;
426
427 let meta_with_chunks = FileMetadata {
429 total_chunks,
430 ..metadata
431 };
432
433 let meta_key = keys::document_meta(file_id);
434 storage
435 .store_encrypted(&meta_key, &meta_with_chunks, ttl::FILE_CHUNK, None)
436 .await?;
437
438 for (i, chunk) in chunks.iter().enumerate() {
440 let chunk_key = keys::file_chunk(file_id, i as u32);
441 storage
442 .store_encrypted(&chunk_key, chunk, ttl::FILE_CHUNK, None)
443 .await?;
444 }
445
446 Ok(())
447 }
448
449 pub async fn get_file(&self, storage: &StorageManager, file_id: &str) -> Result<Vec<u8>> {
451 let meta_key = keys::document_meta(file_id);
453 let metadata: FileMetadata = storage.get_encrypted(&meta_key).await?;
454
455 let mut data = Vec::new();
457 for i in 0..metadata.total_chunks {
458 let chunk_key = keys::file_chunk(file_id, i);
459 let chunk: Vec<u8> = storage.get_encrypted(&chunk_key).await?;
460 data.extend(chunk);
461 }
462
463 Ok(data)
464 }
465}
466
467#[derive(Debug, Clone, Serialize, Deserialize)]
469pub struct FileMetadata {
470 pub file_id: String,
471 pub name: String,
472 pub size: u64,
473 pub mime_type: String,
474 pub hash: Vec<u8>,
475 pub total_chunks: u32,
476 pub created_at: SystemTime,
477 pub created_by: String,
478}