use std::sync::Arc;
use async_trait::async_trait;
use dig_dht::{ContentId, DhtService, ProviderRecord};
use crate::error::DownloadError;
#[async_trait]
pub trait ProviderLocator: Send + Sync {
async fn find_providers(
&self,
content: &ContentId,
) -> Result<Vec<ProviderRecord>, DownloadError>;
}
#[derive(Clone)]
pub struct DhtProviderLocator {
dht: Arc<DhtService>,
}
impl DhtProviderLocator {
pub fn new(dht: Arc<DhtService>) -> Self {
DhtProviderLocator { dht }
}
pub fn service(&self) -> &Arc<DhtService> {
&self.dht
}
}
#[async_trait]
impl ProviderLocator for DhtProviderLocator {
async fn find_providers(
&self,
content: &ContentId,
) -> Result<Vec<ProviderRecord>, DownloadError> {
self.dht
.find_providers(content)
.await
.map_err(|e| DownloadError::transport("dht", e))
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use dig_dht::transport::DhtTransport;
use dig_dht::{
CandidateAddr, Contact, DhtConfig, DhtError, DhtRequest, DhtResponse, Key, PeerId,
ProviderRecord,
};
struct CannedTransport {
providers: Vec<ProviderRecord>,
}
#[async_trait]
impl DhtTransport for CannedTransport {
async fn rpc(
&self,
_from: &Contact,
_peer: &Contact,
request: &DhtRequest,
) -> Result<DhtResponse, DhtError> {
match request {
DhtRequest::FindProviders { .. } => Ok(DhtResponse::Providers {
providers: self.providers.clone(),
closer: vec![],
}),
DhtRequest::Ping { nonce } => Ok(DhtResponse::Pong { nonce: *nonce }),
_ => Ok(DhtResponse::Nodes { nodes: vec![] }),
}
}
}
#[tokio::test]
async fn dht_locator_returns_found_providers() {
let content = ContentId::resource([1; 32], [2; 32], [3; 32]);
let rec = ProviderRecord::new(
&content.to_key(),
&PeerId::from_bytes([7; 32]),
vec![CandidateAddr::direct("203.0.113.7", 9444)],
u64::MAX,
);
let transport = Arc::new(CannedTransport {
providers: vec![rec.clone()],
});
let dht = DhtService::new(
PeerId::from_bytes([0; 32]),
vec![CandidateAddr::direct("127.0.0.1", 1)],
DhtConfig::default(),
transport,
);
dht.bootstrap(&[dig_dht::BootstrapPeer::direct(
PeerId::from_bytes([5; 32]),
"127.0.0.1",
2,
)])
.await
.unwrap();
let locator = DhtProviderLocator::new(Arc::new(dht));
let got = locator.find_providers(&content).await.unwrap();
assert_eq!(got.len(), 1);
assert_eq!(got[0].provider_peer_id, rec.provider_peer_id);
assert!(locator.service().local_id() == &PeerId::from_bytes([0; 32]));
let _ = Key::from_bytes([0; 32]);
}
}