Skip to content

Commit

Permalink
resolver async awake lookup
Browse files Browse the repository at this point in the history
  • Loading branch information
@valesteban
valesteban committed Oct 17, 2023
1 parent bae1a62 commit 7bdf209
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 43 deletions.
17 changes: 15 additions & 2 deletions src/resolver/async_resolver.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ impl AsyncResolver {

//Async query
let response = LookupIpFutureStub::lookup(domain_name_struct, self.cache.clone(),conn).await;

println!("[LOOKUP IP RESPONSE => {:?}]",response);
let ip_addr = match response {
Ok(val) => {
Expand Down Expand Up @@ -66,6 +66,7 @@ mod async_resolver_test {
// use tokio::runtime::Runtime;
use crate::resolver::config::ResolverConfig;
use super::AsyncResolver;



// #[test]
Expand All @@ -87,12 +88,24 @@ mod async_resolver_test {
#[tokio::test]
#[ignore]
async fn lookupip_example() {
println!("[TEST INIT]");

let resolver = AsyncResolver::new(ResolverConfig::default());

let response = resolver.lookup_ip("example.com", "UDP").await.unwrap();

println!("[TEST => {}]",response);
println!("[TEST FINISH=> {}]",response);



}

// #[test]
// fn test_spawn() {

// let resolver = AsyncResolver::new(ResolverConfig::default());

// let handle = tokio::spawn(resolver.lookup_ip("example.com", "UDP"));
// unimplemented!();
// }
}
116 changes: 75 additions & 41 deletions src/resolver/lookup.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,26 @@ use crate::message::type_qtype::Qtype;
use crate::message::question::Question;
use crate::client::client_error::ClientError;

use futures_util::FutureExt;
use std::sync::Arc;
use futures_util::{FutureExt,task::Waker};
use tokio::io::AsyncWriteExt;
use std::pin::Pin;
use std::task::{Poll,Context};
//TODO: Eliminar librerias
use std::net::{IpAddr,Ipv4Addr};
use futures_util::{future::Future,future};
use super::resolver_error::ResolverError;
use std::time::Duration;
use std::sync:: Mutex;
use crate::client::client_connection::ClientConnectionType;

//Future returned fron AsyncResolver when performing a lookup with rtype A
pub struct LookupIpFutureStub {
name: DomainName, // cache: DnsCache,
query: Pin< Box< dyn Future <Output = Result<DnsMessage, ResolverError >> > >,
query:Pin< Box< dyn Future<Output = Result<DnsMessage, ResolverError>> >>,
cache: DnsCache,
conn: ClientConnectionType,
waker: Option<Waker>,
}

impl Future for LookupIpFutureStub{
Expand All @@ -34,63 +38,61 @@ impl Future for LookupIpFutureStub{
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
println!("[POLL FUTURE]");

loop {
let query = self.query.as_mut().poll(cx);
println!("[POLL] query ");

match query {
Poll::Pending => {
println!(" [Pending]");
return Poll::Pending;
},
Poll::Ready(Err(_)) => {
println!(" [ready err]");
self.query = Box::pin(lookup_stub(self.name.clone(), self.cache.clone(),self.conn.clone()));
},
Poll::Ready(Ok(ip_addr)) => {
println!(" [Ready]");
return Poll::Ready(Ok(ip_addr));
}
let query = self.query.as_mut().poll(cx);
println!("[POLL query {:?}",query);

match query {
Poll::Pending => {
println!(" [return pending]");
return Poll::Pending;
},
Poll::Ready(Err(_)) => {
println!(" [ready empty]");
self.waker = Some(cx.waker().clone());

tokio::spawn(
lookup_stub(self.name.clone(),self.cache.clone(),self.conn.clone(),self.waker.clone()));
println!(" [return pending]");
return Poll::Pending;
},
Poll::Ready(Ok(ip_addr)) => {
println!(" [return ready]");
return Poll::Ready(Ok(ip_addr));
}
}
}

}



impl LookupIpFutureStub{
pub fn lookup(
name: DomainName,
cache:DnsCache,
conn: ClientConnectionType,
) -> Self {
println!("[LOOKUP FUTURE]");
println!("[LOOKUP CREATE FUTURE]");;

Self {
name: name,
query: future::err(ResolverError::Message("Empty")).boxed(), //FIXME: cambiar a otro tipo el error/inicio
cache: cache,
conn: conn,
waker: None,
}

}
}

}
pub async fn lookup_stub( //FIXME: podemos ponerle de nombre lookup_strategy y que se le pase ahi un parametro strategy que diga si son los pasos o si funciona como stub
name: DomainName,
mut cache: DnsCache,
conn: ClientConnectionType,
waker: Option<Waker>,
// query:Pin< Box< dyn Future <Output = Result<DnsMessage, ResolverError >> > >,
) -> Result<DnsMessage,ResolverError> {
println!("[LOOKUP STUB]");
// //FIXME: change values
// let server:IpAddr = IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1));
// let timeout:Duration = Duration::new(2, 0);

// //Connection type
// let conn = ClientUDPConnection::new(server, timeout);
//Create query
let mut query = DnsMessage::new_query_message(

let mut new_query = DnsMessage::new_query_message(
DomainName::new_from_string("example.com".to_string()),
Qtype::A,
Qclass::IN,
Expand All @@ -101,7 +103,7 @@ pub async fn lookup_stub( //FIXME: podemos ponerle de nombre lookup_strategy y

let mut question = Question::new();
question.set_qclass(Qclass::IN);
query.set_question(question);
new_query.set_question(question);



Expand All @@ -114,25 +116,33 @@ pub async fn lookup_stub( //FIXME: podemos ponerle de nombre lookup_strategy y
.iter()
.map(|rr_cache_value| rr_cache_value.get_resource_record())
.collect::<Vec<ResourceRecord>>();
query.set_answer(answer);
return Ok(query);
new_query.set_answer(answer);
return Ok(new_query);
}

//FIXME:
let responseResult: Result<DnsMessage, ResolverError> = match conn {
ClientConnectionType::TCP(client_conn) => {
match client_conn.send(query) {
match client_conn.send(new_query) {
Err(_) => Err(ResolverError::Message("Error: Receiving DNS message")),
Ok(val) => Ok(val),
Ok(val) => {
Ok(val)
},
}
}
ClientConnectionType::UDP(client_conn) => {
match client_conn.send(query) {
match client_conn.send(new_query) {
Err(_) => Err(ResolverError::Message("Error: Receiving DNS message")),
Ok(val) => Ok(val),
Ok(val) => {
Ok(val)},
}
}
};
};
//para que en siguient eciclo de tokio despierte esta task
if let Some(waker) = waker {
println!("[LOOKUP STUB] wake");
waker.wake();
}

println!("[LOOKUP STUB] return");
responseResult
Expand All @@ -141,8 +151,32 @@ pub async fn lookup_stub( //FIXME: podemos ponerle de nombre lookup_strategy y



#[cfg(test)]
mod async_resolver_test {
// use tokio::runtime::Runtime;
use crate::{ domain_name::DomainName, dns_cache::DnsCache};
use super::lookup_stub;
use tokio::time::Duration;
use std::net::{IpAddr, Ipv4Addr};
use super::*;

#[tokio::test]
async fn lookup_stub_test() {
let name = DomainName::new_from_string("example.com".to_string());
let cache = DnsCache::new();
let waker = None;

let google_server: IpAddr = IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8));
let timeout: Duration = Duration::from_secs(20);

let client_udp = ClientUDPConnection::new(google_server, timeout);
let conn = ClientConnectionType::UDP(client_udp);

let result = lookup_stub(name, cache, conn, waker).await;
println!("[Test Result ] {:?}", result);
}





}

0 comments on commit 7bdf209

Please sign in to comment.