1use crate::dht::{DHT, Key};
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 dht_key = Key::new(key.as_bytes());
206
207 self.dht
208 .put(dht_key, wrapper_bytes)
209 .await
210 .map_err(|e| StorageError::DhtError(e.to_string()))?;
211
212 Ok(())
213 }
214
215 pub async fn get_encrypted<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T> {
217 let dht_key = Key::new(key.as_bytes());
219 let record = self
220 .dht
221 .get(&dht_key)
222 .await
223 .ok_or_else(|| StorageError::KeyNotFound(key.to_string()))?;
224
225 let wrapper: EncryptedData = bincode::deserialize(&record.value)?;
227
228 let plaintext = self.decrypt(&wrapper.ciphertext, &wrapper.nonce)?;
230
231 let data = bincode::deserialize(&plaintext)?;
233
234 Ok(data)
235 }
236
237 pub async fn store_public<T: Serialize>(
239 &mut self,
240 key: &str,
241 data: &T,
242 _ttl: Duration,
243 ) -> Result<()> {
244 let value = bincode::serialize(data)?;
245
246 let dht_key = Key::new(key.as_bytes());
247
248 self.dht
249 .put(dht_key, value)
250 .await
251 .map_err(|e| StorageError::DhtError(e.to_string()))?;
252
253 Ok(())
254 }
255
256 pub async fn get_public<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T> {
258 let dht_key = Key::new(key.as_bytes());
259 let record = self
260 .dht
261 .get(&dht_key)
262 .await
263 .ok_or_else(|| StorageError::KeyNotFound(key.to_string()))?;
264
265 let data = bincode::deserialize(&record.value)?;
266 Ok(data)
267 }
268
269 pub async fn delete(&mut self, key: &str) -> Result<()> {
271 let dht_key = Key::new(key.as_bytes());
273 self.dht
274 .put(dht_key, vec![])
275 .await
276 .map_err(|e| StorageError::DhtError(e.to_string()))?;
277 Ok(())
278 }
279
280 pub async fn list_keys(&self, _prefix: &str) -> Result<Vec<String>> {
282 Ok(vec![])
285 }
286
287 fn encrypt(&self, plaintext: &[u8]) -> Result<(Vec<u8>, [u8; 12])> {
289 let cipher = Aes256Gcm::new(AesKey::<Aes256Gcm>::from_slice(&self.master_key));
290
291 let nonce_bytes = rand::random::<[u8; 12]>();
293 let nonce = Nonce::from_slice(&nonce_bytes);
294
295 let ciphertext = cipher
296 .encrypt(nonce, plaintext)
297 .map_err(|e| StorageError::EncryptionError(e.to_string()))?;
298
299 Ok((ciphertext, nonce_bytes))
300 }
301
302 fn decrypt(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>> {
304 let cipher = Aes256Gcm::new(AesKey::<Aes256Gcm>::from_slice(&self.master_key));
305 let nonce = Nonce::from_slice(nonce);
306
307 let plaintext = cipher
308 .decrypt(nonce, ciphertext)
309 .map_err(|e| StorageError::EncryptionError(e.to_string()))?;
310
311 Ok(plaintext)
312 }
313}
314
315pub struct FileChunker {
393 chunk_size: usize,
394}
395
396impl FileChunker {
397 pub fn new(chunk_size: usize) -> Self {
399 Self { chunk_size }
400 }
401
402 pub fn chunk_file(&self, data: &[u8]) -> Vec<Vec<u8>> {
404 data.chunks(self.chunk_size)
405 .map(|chunk| chunk.to_vec())
406 .collect()
407 }
408
409 pub async fn store_file(
411 &self,
412 storage: &mut StorageManager,
413 file_id: &str,
414 data: &[u8],
415 metadata: FileMetadata,
416 ) -> Result<()> {
417 let chunks = self.chunk_file(data);
418 let total_chunks = chunks.len() as u32;
419
420 let meta_with_chunks = FileMetadata {
422 total_chunks,
423 ..metadata
424 };
425
426 let meta_key = keys::document_meta(file_id);
427 storage
428 .store_encrypted(&meta_key, &meta_with_chunks, ttl::FILE_CHUNK, None)
429 .await?;
430
431 for (i, chunk) in chunks.iter().enumerate() {
433 let chunk_key = keys::file_chunk(file_id, i as u32);
434 storage
435 .store_encrypted(&chunk_key, chunk, ttl::FILE_CHUNK, None)
436 .await?;
437 }
438
439 Ok(())
440 }
441
442 pub async fn get_file(&self, storage: &StorageManager, file_id: &str) -> Result<Vec<u8>> {
444 let meta_key = keys::document_meta(file_id);
446 let metadata: FileMetadata = storage.get_encrypted(&meta_key).await?;
447
448 let mut data = Vec::new();
450 for i in 0..metadata.total_chunks {
451 let chunk_key = keys::file_chunk(file_id, i);
452 let chunk: Vec<u8> = storage.get_encrypted(&chunk_key).await?;
453 data.extend(chunk);
454 }
455
456 Ok(data)
457 }
458}
459
460#[derive(Debug, Clone, Serialize, Deserialize)]
462pub struct FileMetadata {
463 pub file_id: String,
464 pub name: String,
465 pub size: u64,
466 pub mime_type: String,
467 pub hash: Vec<u8>,
468 pub total_chunks: u32,
469 pub created_at: SystemTime,
470 pub created_by: String,
471}