use cid::Cid;
use libp2p::PeerId;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum DhtProviderError {
#[error("DHT provider not found: {0}")]
ProviderNotFound(String),
#[error("Operation not supported: {0}")]
OperationNotSupported(String),
#[error("Configuration error: {0}")]
ConfigurationError(String),
#[error("Query failed: {0}")]
QueryFailed(String),
#[error("Internal error: {0}")]
InternalError(String),
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DhtCapabilities {
pub supports_content_routing: bool,
pub supports_peer_routing: bool,
pub supports_kv_storage: bool,
pub supports_range_queries: bool,
pub supports_semantic_queries: bool,
pub max_query_hops: Option<usize>,
pub supports_custom_routing: bool,
}
impl DhtCapabilities {
pub fn basic() -> Self {
Self {
supports_content_routing: true,
supports_peer_routing: true,
supports_kv_storage: false,
supports_range_queries: false,
supports_semantic_queries: false,
max_query_hops: Some(20),
supports_custom_routing: false,
}
}
pub fn advanced() -> Self {
Self {
supports_content_routing: true,
supports_peer_routing: true,
supports_kv_storage: true,
supports_range_queries: true,
supports_semantic_queries: true,
max_query_hops: Some(20),
supports_custom_routing: true,
}
}
}
#[derive(Debug, Clone)]
pub struct DhtQueryResult {
pub providers: Vec<PeerId>,
pub hops: usize,
pub duration_ms: u64,
pub success: bool,
}
#[derive(Debug, Clone)]
pub struct DhtPeerInfo {
pub peer_id: PeerId,
pub addresses: Vec<String>,
pub distance: Option<u64>,
}
pub trait DhtProvider: Send + Sync {
fn name(&self) -> &str;
fn version(&self) -> &str;
fn capabilities(&self) -> DhtCapabilities;
fn bootstrap(&self, peers: Vec<PeerId>) -> Result<(), DhtProviderError>;
fn provide(&self, cid: &Cid) -> Result<(), DhtProviderError>;
fn find_providers(&self, cid: &Cid) -> Result<DhtQueryResult, DhtProviderError>;
fn find_peer(&self, peer_id: &PeerId) -> Result<DhtPeerInfo, DhtProviderError>;
fn get_closest_peers(
&self,
key: &[u8],
count: usize,
) -> Result<Vec<DhtPeerInfo>, DhtProviderError>;
fn put_value(&self, key: &[u8], value: &[u8]) -> Result<(), DhtProviderError> {
let _ = (key, value);
Err(DhtProviderError::OperationNotSupported(
"Key-value storage not supported".to_string(),
))
}
fn get_value(&self, key: &[u8]) -> Result<Vec<u8>, DhtProviderError> {
let _ = key;
Err(DhtProviderError::OperationNotSupported(
"Key-value storage not supported".to_string(),
))
}
fn stats(&self) -> DhtProviderStats;
fn is_healthy(&self) -> bool {
let stats = self.stats();
stats.routing_table_size > 0 && stats.success_rate > 0.5
}
}
#[derive(Debug, Clone, Default)]
pub struct DhtProviderStats {
pub routing_table_size: usize,
pub total_queries: u64,
pub successful_queries: u64,
pub failed_queries: u64,
pub avg_query_duration_ms: f64,
pub success_rate: f64,
}
impl DhtProviderStats {
pub fn calculate_success_rate(&mut self) {
if self.total_queries > 0 {
self.success_rate = self.successful_queries as f64 / self.total_queries as f64;
} else {
self.success_rate = 0.0;
}
}
}
pub struct DhtProviderRegistry {
providers: HashMap<String, Arc<dyn DhtProvider>>,
active_provider: Option<String>,
}
impl DhtProviderRegistry {
pub fn new() -> Self {
Self {
providers: HashMap::new(),
active_provider: None,
}
}
pub fn register(&mut self, name: impl Into<String>, provider: Arc<dyn DhtProvider>) {
let name = name.into();
self.providers.insert(name.clone(), provider);
if self.active_provider.is_none() {
self.active_provider = Some(name);
}
}
pub fn get(&self, name: &str) -> Option<Arc<dyn DhtProvider>> {
self.providers.get(name).cloned()
}
pub fn get_active(&self) -> Option<Arc<dyn DhtProvider>> {
self.active_provider
.as_ref()
.and_then(|name| self.get(name))
}
pub fn set_active(&mut self, name: impl Into<String>) -> Result<(), DhtProviderError> {
let name = name.into();
if self.providers.contains_key(&name) {
self.active_provider = Some(name);
Ok(())
} else {
Err(DhtProviderError::ProviderNotFound(name))
}
}
pub fn list_providers(&self) -> Vec<String> {
self.providers.keys().cloned().collect()
}
pub fn unregister(&mut self, name: &str) -> Option<Arc<dyn DhtProvider>> {
let provider = self.providers.remove(name);
if self.active_provider.as_deref() == Some(name) {
self.active_provider = None;
}
provider
}
pub fn count(&self) -> usize {
self.providers.len()
}
pub fn has_provider(&self, name: &str) -> bool {
self.providers.contains_key(name)
}
}
impl Default for DhtProviderRegistry {
fn default() -> Self {
Self::new()
}
}
pub mod kademlia {
use super::*;
pub struct KademliaDhtProvider {
stats: parking_lot::RwLock<DhtProviderStats>,
config: KademliaConfig,
}
#[derive(Debug, Clone)]
pub struct KademliaConfig {
pub k_bucket_size: usize,
pub alpha: usize,
pub replication_factor: usize,
pub query_timeout_secs: u64,
}
impl Default for KademliaConfig {
fn default() -> Self {
Self {
k_bucket_size: 20,
alpha: 3,
replication_factor: 20,
query_timeout_secs: 60,
}
}
}
impl KademliaDhtProvider {
pub fn new() -> Self {
Self::with_config(KademliaConfig::default())
}
pub fn with_config(config: KademliaConfig) -> Self {
Self {
stats: parking_lot::RwLock::new(DhtProviderStats::default()),
config,
}
}
#[allow(dead_code)]
pub fn config(&self) -> &KademliaConfig {
&self.config
}
}
impl Default for KademliaDhtProvider {
fn default() -> Self {
Self::new()
}
}
impl DhtProvider for KademliaDhtProvider {
fn name(&self) -> &str {
"kademlia"
}
fn version(&self) -> &str {
"1.0.0"
}
fn capabilities(&self) -> DhtCapabilities {
DhtCapabilities {
supports_content_routing: true,
supports_peer_routing: true,
supports_kv_storage: true,
supports_range_queries: false,
supports_semantic_queries: false,
max_query_hops: Some(20),
supports_custom_routing: false,
}
}
fn bootstrap(&self, peers: Vec<PeerId>) -> Result<(), DhtProviderError> {
let mut stats = self.stats.write();
stats.routing_table_size = peers.len();
Ok(())
}
fn provide(&self, _cid: &Cid) -> Result<(), DhtProviderError> {
Ok(())
}
fn find_providers(&self, _cid: &Cid) -> Result<DhtQueryResult, DhtProviderError> {
let mut stats = self.stats.write();
stats.total_queries += 1;
stats.successful_queries += 1;
stats.calculate_success_rate();
Ok(DhtQueryResult {
providers: vec![],
hops: 0,
duration_ms: 0,
success: true,
})
}
fn find_peer(&self, peer_id: &PeerId) -> Result<DhtPeerInfo, DhtProviderError> {
Ok(DhtPeerInfo {
peer_id: *peer_id,
addresses: vec![],
distance: None,
})
}
fn get_closest_peers(
&self,
_key: &[u8],
count: usize,
) -> Result<Vec<DhtPeerInfo>, DhtProviderError> {
let _ = count;
Ok(vec![])
}
fn put_value(&self, _key: &[u8], _value: &[u8]) -> Result<(), DhtProviderError> {
Ok(())
}
fn get_value(&self, _key: &[u8]) -> Result<Vec<u8>, DhtProviderError> {
Err(DhtProviderError::QueryFailed("Not found".to_string()))
}
fn stats(&self) -> DhtProviderStats {
self.stats.read().clone()
}
}
}
#[cfg(test)]
mod tests {
use super::kademlia::*;
use super::*;
#[test]
fn test_dht_capabilities() {
let basic = DhtCapabilities::basic();
assert!(basic.supports_content_routing);
assert!(basic.supports_peer_routing);
assert!(!basic.supports_kv_storage);
let advanced = DhtCapabilities::advanced();
assert!(advanced.supports_content_routing);
assert!(advanced.supports_kv_storage);
assert!(advanced.supports_semantic_queries);
}
#[test]
fn test_kademlia_provider() {
let provider = KademliaDhtProvider::new();
assert_eq!(provider.name(), "kademlia");
assert_eq!(provider.version(), "1.0.0");
let caps = provider.capabilities();
assert!(caps.supports_content_routing);
assert!(caps.supports_peer_routing);
assert!(caps.supports_kv_storage);
}
#[test]
fn test_provider_registry() {
let mut registry = DhtProviderRegistry::new();
assert_eq!(registry.count(), 0);
let provider = Arc::new(KademliaDhtProvider::new());
registry.register("kademlia", provider);
assert_eq!(registry.count(), 1);
assert!(registry.has_provider("kademlia"));
}
#[test]
fn test_registry_active_provider() {
let mut registry = DhtProviderRegistry::new();
let provider = Arc::new(KademliaDhtProvider::new());
registry.register("kademlia", provider);
let active = registry.get_active();
assert!(active.is_some());
assert_eq!(
active
.expect("test: active provider should be set after registration")
.name(),
"kademlia"
);
}
#[test]
fn test_registry_set_active() {
let mut registry = DhtProviderRegistry::new();
let provider1 = Arc::new(KademliaDhtProvider::new());
registry.register("kademlia", provider1);
registry
.set_active("kademlia")
.expect("test: set_active should succeed for registered provider");
assert_eq!(
registry
.get_active()
.expect("test: active provider should be kademlia after set_active")
.name(),
"kademlia"
);
}
#[test]
fn test_registry_unregister() {
let mut registry = DhtProviderRegistry::new();
let provider = Arc::new(KademliaDhtProvider::new());
registry.register("kademlia", provider);
assert_eq!(registry.count(), 1);
registry.unregister("kademlia");
assert_eq!(registry.count(), 0);
}
#[test]
fn test_provider_bootstrap() {
let provider = KademliaDhtProvider::new();
let peers = vec![PeerId::random(), PeerId::random()];
let result = provider.bootstrap(peers);
assert!(result.is_ok());
}
#[test]
fn test_provider_stats() {
let provider = KademliaDhtProvider::new();
let stats = provider.stats();
assert_eq!(stats.total_queries, 0);
assert_eq!(stats.successful_queries, 0);
}
#[test]
fn test_provider_health() {
let provider = KademliaDhtProvider::new();
assert!(!provider.is_healthy());
provider
.bootstrap(vec![PeerId::random()])
.expect("test: bootstrap with one peer should succeed");
let cid = Cid::default();
provider
.find_providers(&cid)
.expect("test: find_providers should succeed on bootstrapped provider");
assert!(provider.is_healthy());
}
#[test]
fn test_list_providers() {
let mut registry = DhtProviderRegistry::new();
let provider = Arc::new(KademliaDhtProvider::new());
registry.register("kademlia", provider);
let providers = registry.list_providers();
assert_eq!(providers.len(), 1);
assert!(providers.contains(&"kademlia".to_string()));
}
#[test]
fn test_provider_not_found() {
let mut registry = DhtProviderRegistry::new();
let result = registry.set_active("nonexistent");
assert!(result.is_err());
assert!(matches!(result, Err(DhtProviderError::ProviderNotFound(_))));
}
#[test]
fn test_stats_success_rate() {
let mut stats = DhtProviderStats {
total_queries: 10,
successful_queries: 7,
..Default::default()
};
stats.calculate_success_rate();
assert!((stats.success_rate - 0.7).abs() < 0.01);
}
#[test]
fn test_kademlia_config_default() {
let config = KademliaConfig::default();
assert_eq!(config.k_bucket_size, 20);
assert_eq!(config.alpha, 3);
assert_eq!(config.replication_factor, 20);
assert_eq!(config.query_timeout_secs, 60);
}
#[test]
fn test_unsupported_operation() {
let provider = KademliaDhtProvider::new();
let result = provider.put_value(b"key", b"value");
assert!(result.is_ok());
}
}