From 1059f35fdcf4942fd253748121c883ea38b427a7 Mon Sep 17 00:00:00 2001
From: "Aode (lion)" <asonix@asonix.dog>
Date: Thu, 27 Jan 2022 10:19:28 -0600
Subject: [PATCH] use pre-constructed client for well-known requests also

---
 Cargo.lock              |  3 +--
 Cargo.toml              |  2 +-
 src/database/globals.rs | 30 ++++++++++++++++++------------
 src/server_server.rs    | 20 +-------------------
 4 files changed, 21 insertions(+), 34 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 794a0257..21c27700 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1950,8 +1950,7 @@ dependencies = [
 [[package]]
 name = "reqwest"
 version = "0.11.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "87f242f1488a539a79bac6dbe7c8609ae43b7914b7736210f239a37cccb32525"
+source = "git+https://github.com/niuhuan/reqwest?branch=dns-resolver-fn#57b7cf4feb921573dfafad7d34b9ac6e44ead0bd"
 dependencies = [
  "base64 0.13.0",
  "bytes",
diff --git a/Cargo.toml b/Cargo.toml
index 9ba1ac05..974b4ce8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -48,7 +48,7 @@ rand = "0.8.4"
 # Used to hash passwords
 rust-argon2 = "0.8.3"
 # Used to send requests
-reqwest = { version = "0.11.4", default-features = false, features = ["rustls-tls", "socks"] }
+reqwest = { version = "0.11.4", default-features = false, features = ["rustls-tls", "socks"], git = "https://github.com/niuhuan/reqwest", branch = "dns-resolver-fn" }
 # Used for conduit::Error type
 thiserror = "1.0.28"
 # Used to generate thumbnails for images
diff --git a/src/database/globals.rs b/src/database/globals.rs
index da91c1fb..3278b7f6 100644
--- a/src/database/globals.rs
+++ b/src/database/globals.rs
@@ -10,7 +10,7 @@ use std::{
     collections::{BTreeMap, HashMap},
     fs,
     future::Future,
-    net::IpAddr,
+    net::{IpAddr, SocketAddr},
     path::PathBuf,
     sync::{Arc, Mutex, RwLock},
     time::{Duration, Instant},
@@ -39,6 +39,7 @@ pub struct Globals {
     keypair: Arc<ruma::signatures::Ed25519KeyPair>,
     dns_resolver: TokioAsyncResolver,
     jwt_decoding_key: Option<jsonwebtoken::DecodingKey<'static>>,
+    well_known_client: reqwest::Client,
     basic_client: reqwest::Client,
     pub(super) server_signingkeys: Arc<dyn Tree>,
     pub bad_event_ratelimiter: Arc<RwLock<HashMap<Box<EventId>, RateLimitState>>>,
@@ -133,7 +134,16 @@ impl Globals {
             .as_ref()
             .map(|secret| jsonwebtoken::DecodingKey::from_secret(secret.as_bytes()).into_static());
 
-        let basic_client = reqwest_client_builder(&config, None)?.build()?;
+        let basic_client = reqwest_client_builder(&config)?.build()?;
+        let name_override = Arc::clone(&tls_name_override);
+        let well_known_client = reqwest_client_builder(&config)?
+            .resolve_fn(move |domain| {
+                let read_guard = name_override.read().unwrap();
+                let (override_name, port) = read_guard.get(&domain)?;
+                let first_name = override_name.get(0)?;
+                Some(SocketAddr::new(*first_name, *port))
+            })
+            .build()?;
 
         let s = Self {
             globals,
@@ -144,6 +154,7 @@ impl Globals {
             })?,
             actual_destination_cache: Arc::new(RwLock::new(WellKnownMap::new())),
             tls_name_override,
+            well_known_client,
             basic_client,
             server_signingkeys,
             jwt_decoding_key,
@@ -173,9 +184,10 @@ impl Globals {
         self.basic_client.clone()
     }
 
-    /// Returns a reqwest client builder which can be customized and used to send requests.
-    pub fn reqwest_client_builder(&self) -> Result<reqwest::ClientBuilder> {
-        reqwest_client_builder(&self.config, Some(1))
+    /// Returns a client used for resolving .well-knowns
+    pub fn well_known_client(&self) -> reqwest::Client {
+        // can't return &Client or else we'll hold a lock around the DB across an await
+        self.well_known_client.clone()
     }
 
     #[tracing::instrument(skip(self))]
@@ -343,17 +355,11 @@ impl Globals {
     }
 }
 
-fn reqwest_client_builder(
-    config: &Config,
-    max_idle: Option<usize>,
-) -> Result<reqwest::ClientBuilder> {
+fn reqwest_client_builder(config: &Config) -> Result<reqwest::ClientBuilder> {
     let mut reqwest_client_builder = reqwest::Client::builder()
         .connect_timeout(Duration::from_secs(30))
         .timeout(Duration::from_secs(60 * 3));
 
-    if let Some(max_idle) = max_idle {
-        reqwest_client_builder = reqwest_client_builder.pool_max_idle_per_host(max_idle);
-    }
     if let Some(proxy) = config.proxy.to_proxy()? {
         reqwest_client_builder = reqwest_client_builder.proxy(proxy);
     }
diff --git a/src/server_server.rs b/src/server_server.rs
index 205355f9..978eb67f 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -237,25 +237,7 @@ where
 
     let url = reqwest_request.url().clone();
 
-    let client = if let Some((override_name, port)) = globals
-        .tls_name_override
-        .read()
-        .unwrap()
-        .get(&actual_destination.hostname())
-    {
-        globals
-            .reqwest_client_builder()?
-            .resolve(
-                &actual_destination.hostname(),
-                SocketAddr::new(override_name[0], *port),
-            )
-            .build()?
-        // port will be ignored
-    } else {
-        globals.reqwest_client()
-    };
-
-    let response = client.execute(reqwest_request).await;
+    let response = globals.well_known_client().execute(reqwest_request).await;
 
     match response {
         Ok(mut response) => {