diff --git a/src/client.rs b/src/client.rs index 60044d6..2b2195f 100644 --- a/src/client.rs +++ b/src/client.rs @@ -34,6 +34,21 @@ impl Client { Self::extract_answers(resp) } + pub async fn query_with_retry( + &self, + questions: Vec>, + retries: usize, + ) -> Result>>, String> { + let mut last_res = Err("Dummy".to_string()); + for i in (0..retries) { + last_res = self.query(questions.clone()).await; + if last_res.is_ok() { + break; + } + } + return last_res; + } + // Select an upstream randomly fn select_upstream(&self) -> String { let idx = unsafe { Math::random() } * self.options.upstream_urls.len() as f64; @@ -51,7 +66,7 @@ impl Client { header.set_qr(false); // For queries, QR = false header.set_opcode(Opcode::Query); header.set_rd(true); // Ask for recursive queries - // Set up the questions + // Set up the questions for q in questions { builder .push(q) diff --git a/src/server.rs b/src/server.rs index f8959d4..6929d1a 100644 --- a/src/server.rs +++ b/src/server.rs @@ -44,6 +44,7 @@ enum DnsResponseFormat { #[derive(Deserialize)] pub struct ServerOptions { upstream_urls: Vec, + retries: usize, } pub struct Server { @@ -75,7 +76,11 @@ impl Server { let body = err_response!(Self::parse_dns_body(&req).await); let query_id = body.header().id(); // random ID that needs to be preserved in response let questions = err_response!(Self::extract_questions(body)); - let records = err_response!(self.client.query(questions).await); + let records = err_response!( + self.client + .query_with_retry(questions, self.options.retries) + .await + ); let resp_format = Self::get_response_format(&req); let mut resp_body = err_response!(match &resp_format {