shadow_dht/

store.rs

//! DHT key-value store implementation

use shadow_core::error::{Result, ShadowError};
use std::collections::HashMap;
use std::time::{Instant, Duration};
use bytes::Bytes;
use serde::{Serialize, Deserialize};

/// Value stored in DHT
#[derive(Debug, Clone)]
pub struct StoredValue {
    /// The actual data
    pub data: Bytes,
    /// Publisher peer ID (as bytes)
    pub publisher: [u8; 32],
    /// Publication timestamp
    pub timestamp: Instant,
    /// Time-to-live in seconds
    pub ttl: u64,
    /// Signature (optional)
    pub signature: Option<Vec<u8>>,
}

impl StoredValue {
    /// Create new stored value
    pub fn new(data: Bytes, publisher: [u8; 32], ttl: u64) -> Self {
        Self {
            data,
            publisher,
            timestamp: Instant::now(),
            ttl,
            signature: None,
        }
    }

    /// Check if value has expired
    pub fn is_expired(&self) -> bool {
        self.timestamp.elapsed() > Duration::from_secs(self.ttl)
    }

    /// Time until expiration
    pub fn time_until_expiry(&self) -> Option<Duration> {
        let elapsed = self.timestamp.elapsed();
        let ttl_duration = Duration::from_secs(self.ttl);
        
        if elapsed < ttl_duration {
            Some(ttl_duration - elapsed)
        } else {
            None
        }
    }
}

/// DHT storage backend
pub struct DHTStore {
    /// Key-value storage
    store: HashMap<[u8; 32], StoredValue>,
    /// Maximum storage size (bytes)
    max_size: usize,
    /// Current storage size
    current_size: usize,
}

impl DHTStore {
    /// Create new DHT store
    pub fn new(max_size: usize) -> Self {
        Self {
            store: HashMap::new(),
            max_size,
            current_size: 0,
        }
    }

    /// Store a value
    pub fn put(&mut self, key: [u8; 32], value: StoredValue) -> Result<()> {
        // Check if we need to make space
        if !self.store.contains_key(&key) {
            let value_size = value.data.len();
            
            // Remove expired entries first
            self.cleanup_expired();
            
            // Check if we have space
            if self.current_size + value_size > self.max_size {
                // Try to evict oldest entries
                if !self.make_space(value_size) {
                    return Err(ShadowError::Storage("Storage full".into()));
                }
            }
            
            self.current_size += value_size;
        } else {
            // Replacing existing value, adjust size
            if let Some(old_value) = self.store.get(&key) {
                self.current_size -= old_value.data.len();
            }
            self.current_size += value.data.len();
        }

        self.store.insert(key, value);
        Ok(())
    }

    /// Get a value
    pub fn get(&self, key: &[u8; 32]) -> Option<&StoredValue> {
        self.store.get(key).filter(|v| !v.is_expired())
    }

    /// Remove a value
    pub fn remove(&mut self, key: &[u8; 32]) -> Option<StoredValue> {
        if let Some(value) = self.store.remove(key) {
            self.current_size -= value.data.len();
            Some(value)
        } else {
            None
        }
    }

    /// Check if key exists (and not expired)
    pub fn contains(&self, key: &[u8; 32]) -> bool {
        self.store.get(key).map_or(false, |v| !v.is_expired())
    }

    /// Remove expired entries
    pub fn cleanup_expired(&mut self) -> usize {
        let expired_keys: Vec<[u8; 32]> = self.store
            .iter()
            .filter(|(_, v)| v.is_expired())
            .map(|(k, _)| *k)
            .collect();

        let count = expired_keys.len();
        for key in expired_keys {
            self.remove(&key);
        }

        count
    }

    /// Make space by evicting oldest entries
    fn make_space(&mut self, needed: usize) -> bool {
        // Collect entries sorted by age (oldest first)
        let mut entries: Vec<_> = self.store.iter().collect();
        entries.sort_by_key(|(_, v)| v.timestamp);

        let mut freed = 0;
        let mut to_remove = Vec::new();

        for (key, value) in entries {
            if freed >= needed {
                break;
            }
            to_remove.push(*key);
            freed += value.data.len();
        }

        // Remove entries
        for key in to_remove {
            self.remove(&key);
        }

        freed >= needed
    }

    /// Get all keys
    pub fn keys(&self) -> Vec<[u8; 32]> {
        self.store.keys().copied().collect()
    }

    /// Get number of stored items
    pub fn len(&self) -> usize {
        self.store.len()
    }

    /// Check if store is empty
    pub fn is_empty(&self) -> bool {
        self.store.is_empty()
    }

    /// Get current storage size
    pub fn size(&self) -> usize {
        self.current_size
    }

    /// Get values that need republishing (approaching expiry)
    pub fn get_expiring_soon(&self, threshold: Duration) -> Vec<([u8; 32], StoredValue)> {
        self.store
            .iter()
            .filter(|(_, v)| {
                if let Some(time_left) = v.time_until_expiry() {
                    time_left < threshold
                } else {
                    false
                }
            })
            .map(|(k, v)| (*k, v.clone()))
            .collect()
    }

    /// Clear all entries
    pub fn clear(&mut self) {
        self.store.clear();
        self.current_size = 0;
    }
}

impl Default for DHTStore {
    fn default() -> Self {
        Self::new(100 * 1024 * 1024) // 100 MB default
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_store_put_get() {
        let mut store = DHTStore::new(1024 * 1024);

        let key = [1u8; 32];
        let value = StoredValue::new(
            Bytes::from("test data"),
            [2u8; 32],
            3600,
        );

        store.put(key, value.clone()).unwrap();

        let retrieved = store.get(&key).unwrap();
        assert_eq!(retrieved.data, value.data);
    }

    #[test]
    fn test_store_remove() {
        let mut store = DHTStore::new(1024 * 1024);

        let key = [1u8; 32];
        let value = StoredValue::new(Bytes::from("test"), [2u8; 32], 3600);

        store.put(key, value).unwrap();
        assert!(store.contains(&key));

        store.remove(&key);
        assert!(!store.contains(&key));
    }

    #[test]
    fn test_store_expiry() {
        let mut store = DHTStore::new(1024 * 1024);

        let key = [1u8; 32];
        let mut value = StoredValue::new(Bytes::from("test"), [2u8; 32], 0);
        value.timestamp = Instant::now() - Duration::from_secs(10);

        store.put(key, value).unwrap();

        // Should not retrieve expired value
        assert!(store.get(&key).is_none());
    }

    #[test]
    fn test_store_size_limit() {
        let mut store = DHTStore::new(100); // Very small store

        let key1 = [1u8; 32];
        let key2 = [2u8; 32];

        let value1 = StoredValue::new(Bytes::from(vec![0u8; 60]), [0u8; 32], 3600);
        let value2 = StoredValue::new(Bytes::from(vec![0u8; 60]), [0u8; 32], 3600);

        store.put(key1, value1).unwrap();
        
        // Should evict old entry to make space
        let result = store.put(key2, value2);
        assert!(result.is_ok());
    }

    #[test]
    fn test_cleanup_expired() {
        let mut store = DHTStore::new(1024 * 1024);

        // Add expired value
        let key = [1u8; 32];
        let mut value = StoredValue::new(Bytes::from("test"), [2u8; 32], 0);
        value.timestamp = Instant::now() - Duration::from_secs(10);
        store.put(key, value).unwrap();

        assert_eq!(store.len(), 1);
        let cleaned = store.cleanup_expired();
        assert_eq!(cleaned, 1);
        assert_eq!(store.len(), 0);
    }
}