From 1a3aefc1c70197721453ca4d6806e6d5f85b892b Mon Sep 17 00:00:00 2001
From: Ertugrul Aypek <ertugrul.aypek@toolsforhumanity.com>
Date: Sun, 12 Jan 2025 01:08:55 +0100
Subject: [PATCH] use async resolver

---
 .../stage/smpcv2-1-stage/values-iris-mpc.yaml |  2 +-
 iris-mpc/src/bin/server.rs                    | 36 ++++++++++++++-----
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml b/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml
index 95bd124ea..efab13ddf 100644
--- a/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml
+++ b/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml
@@ -1,6 +1,6 @@
 env:
   - name: RUST_LOG
-    value: "info"
+    value: "debug"
 
   - name: RUST_BACKTRACE
     value: "full"
diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
index 3ed1b4c0f..1cd6f322a 100644
--- a/iris-mpc/src/bin/server.rs
+++ b/iris-mpc/src/bin/server.rs
@@ -5,14 +5,14 @@ use aws_sdk_s3::{config::Builder as S3ConfigBuilder, Client as S3Client};
 use aws_sdk_sns::{types::MessageAttributeValue, Client as SNSClient};
 use aws_sdk_sqs::{config::Region, Client};
 use aws_smithy_experimental::hyper_1_0::{CryptoMode, HyperClientBuilder};
-use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns};
+use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns, ResolveDnsError};
 use axum::{response::IntoResponse, routing::get, Router};
 use clap::Parser;
 use eyre::{eyre, Context};
 use futures::{stream::select_all, StreamExt, TryStreamExt};
 use hickory_resolver::{
     config::{ResolverConfig, ResolverOpts},
-    Resolver,
+    TokioAsyncResolver,
 };
 use iris_mpc_common::{
     config::{Config, Opt},
@@ -56,7 +56,9 @@ use std::{
     backtrace::Backtrace,
     collections::{HashMap, HashSet},
     fmt::{Debug, Formatter},
-    mem, panic,
+    mem,
+    net::IpAddr,
+    panic,
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc, LazyLock, Mutex,
@@ -670,13 +672,13 @@ async fn main() -> eyre::Result<()> {
 }
 
 struct StaticResolver {
-    resolver: Arc<Resolver>,
+    resolver: Arc<TokioAsyncResolver>,
 }
 
 impl StaticResolver {
     fn new() -> Self {
-        let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default())
-            .expect("Failed to create resolver");
+        let resolver =
+            TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
         StaticResolver {
             resolver: Arc::new(resolver),
         }
@@ -701,9 +703,25 @@ impl Clone for StaticResolver {
 
 impl ResolveDns for StaticResolver {
     fn resolve_dns<'a>(&'a self, _name: &'a str) -> DnsFuture<'a> {
-        let lookup_result = self.resolver.lookup_ip(_name).unwrap();
-        let ip_addresses = lookup_result.iter().collect();
-        DnsFuture::ready(Ok(ip_addresses))
+        let resolver = Arc::clone(&self.resolver);
+        let hostname = _name.to_string();
+
+        // Create the async block that performs DNS resolution
+        let future = async move {
+            match resolver.lookup_ip(&hostname).await {
+                Ok(lookup_result) => {
+                    let ips: Vec<IpAddr> = lookup_result.iter().collect();
+                    Ok(ips)
+                }
+                Err(e) => Err(ResolveDnsError::new(format!(
+                    "Failed to resolve {}: {}",
+                    hostname, e
+                ))),
+            }
+        };
+
+        // Wrap the future into DnsFuture
+        DnsFuture::new(Box::pin(future))
     }
 }