server: implement query retries

src/client.rs

@@ -34,6 +34,21 @@ impl Client {
         Self::extract_answers(resp)
     }
 
+    pub async fn query_with_retry(
+        &self,
+        questions: Vec<Question<ParsedDname>>,
+        retries: usize,
+    ) -> Result<Vec<Record<ParsedDname, AllRecordData<ParsedDname>>>, 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)

src/server.rs

@@ -44,6 +44,7 @@ enum DnsResponseFormat {
 #[derive(Deserialize)]
 pub struct ServerOptions {
     upstream_urls: Vec<String>,
+    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 {