83 lines
2.6 KiB
Rust
83 lines
2.6 KiB
Rust
use std::{future::Future, pin::Pin, task::{Context, Poll}};
|
|
use std::sync::Arc;
|
|
use futures_util::{ready, Stream, StreamExt};
|
|
|
|
use std::ops::{Deref, DerefMut};
|
|
|
|
use tokio::{net::TcpStream as TokioTcpStream, task};
|
|
use trust_dns_client::{client::AsyncClient, error::ClientError, op::DnsResponse, tcp::TcpClientStream};
|
|
use trust_dns_client::rr::dnssec::tsig::TSigner;
|
|
use trust_dns_proto::error::{ProtoError, ProtoErrorKind};
|
|
use trust_dns_proto::iocompat::AsyncIoTokioAsStd;
|
|
|
|
use crate::config::DnsClientConfig;
|
|
|
|
|
|
pub struct DnsClient(AsyncClient);
|
|
|
|
impl Deref for DnsClient {
|
|
type Target = AsyncClient;
|
|
|
|
fn deref(&self) -> &Self::Target {
|
|
&self.0
|
|
}
|
|
}
|
|
|
|
impl DerefMut for DnsClient {
|
|
fn deref_mut(&mut self) -> &mut Self::Target {
|
|
&mut self.0
|
|
}
|
|
}
|
|
|
|
impl DnsClient {
|
|
pub async fn from_config(dns_config: &DnsClientConfig) -> Result<Self, ProtoError> {
|
|
info!("Creating DNS client for {}", dns_config.server);
|
|
let (stream, handle) = TcpClientStream::<AsyncIoTokioAsStd<TokioTcpStream>>::new(dns_config.server);
|
|
let signer = if let Some(tsig_config) = dns_config.tsig.as_ref() {
|
|
info!("Client configured with TSIG authentication");
|
|
Some(Arc::new(TSigner::new(
|
|
tsig_config.key.clone(),
|
|
tsig_config.algorithm.clone(),
|
|
tsig_config.name.clone().into_inner(),
|
|
60,
|
|
)?.into()))
|
|
} else {
|
|
info!("Client configured without authentication");
|
|
None
|
|
};
|
|
|
|
let client = AsyncClient::with_timeout(
|
|
stream,
|
|
handle,
|
|
dns_config.timeout,
|
|
signer,
|
|
);
|
|
let (client, bg) = client.await?;
|
|
task::spawn(bg);
|
|
return Ok(DnsClient(client))
|
|
}
|
|
}
|
|
|
|
// Reimplement this type here as ClientReponse in trust-dns crate have private fields
|
|
// https://github.com/bluejekyll/trust-dns/blob/v0.21.2/crates/client/src/client/async_client.rs#L621-L641
|
|
pub struct ClientResponse<R>(pub(crate) R)
|
|
where
|
|
R: Stream<Item = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static;
|
|
|
|
impl<R> Future for ClientResponse<R>
|
|
where
|
|
R: Stream<Item = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static,
|
|
{
|
|
type Output = Result<DnsResponse, ClientError>;
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
Poll::Ready(
|
|
match ready!(self.0.poll_next_unpin(cx)) {
|
|
Some(r) => r,
|
|
None => Err(ProtoError::from(ProtoErrorKind::Timeout)),
|
|
}
|
|
.map_err(ClientError::from),
|
|
)
|
|
}
|
|
}
|