From dc18f89c0bd851aa861b59cdf78342176b31f426 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Thu, 4 Jul 2024 11:46:30 +0000
Subject: [PATCH] don't cache server name lookups indefinitely (#436)

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/service/globals/resolver.rs |  1 +
 src/service/sending/resolve.rs  | 30 ++++++++++++++++++++++++++++--
 src/service/sending/send.rs     |  1 +
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/service/globals/resolver.rs b/src/service/globals/resolver.rs
index 0cfaa1b9..3082f2fd 100644
--- a/src/service/globals/resolver.rs
+++ b/src/service/globals/resolver.rs
@@ -122,6 +122,7 @@ impl Resolve for Hooked {
 			.read()
 			.expect("locked for reading")
 			.get(name.as_str())
+			.filter(|cached| cached.valid())
 			.cloned();
 
 		if let Some(cached) = cached {
diff --git a/src/service/sending/resolve.rs b/src/service/sending/resolve.rs
index 01943cc0..77311006 100644
--- a/src/service/sending/resolve.rs
+++ b/src/service/sending/resolve.rs
@@ -2,6 +2,7 @@ use std::{
 	fmt,
 	fmt::Debug,
 	net::{IpAddr, SocketAddr},
+	time::SystemTime,
 };
 
 use hickory_resolver::{error::ResolveError, lookup::SrvLookup};
@@ -9,7 +10,7 @@ use ipaddress::IPAddress;
 use ruma::{OwnedServerName, ServerName};
 use tracing::{debug, error, trace};
 
-use crate::{debug_error, debug_info, debug_warn, services, Error, Result};
+use crate::{debug_error, debug_info, debug_warn, services, utils::rand, Error, Result};
 
 /// Wraps either an literal IP address plus port, or a hostname plus complement
 /// (colon-plus-port if it was specified).
@@ -47,12 +48,14 @@ pub(crate) struct ActualDest {
 pub struct CachedDest {
 	pub dest: FedDest,
 	pub host: String,
+	pub expire: SystemTime,
 }
 
 #[derive(Clone, Debug)]
 pub struct CachedOverride {
 	pub ips: Vec<IpAddr>,
 	pub port: u16,
+	pub expire: SystemTime,
 }
 
 #[tracing::instrument(skip_all, name = "resolve")]
@@ -125,6 +128,7 @@ pub async fn resolve_actual_dest(dest: &ServerName, cache: bool) -> Result<Cache
 	Ok(CachedDest {
 		dest: actual_dest,
 		host: host.into_uri_string(),
+		expire: CachedDest::default_expire(),
 	})
 }
 
@@ -286,6 +290,7 @@ async fn query_and_cache_override(overname: &'_ str, hostname: &'_ str, port: u1
 				CachedOverride {
 					ips: override_ip.iter().collect(),
 					port,
+					expire: CachedOverride::default_expire(),
 				},
 			);
 
@@ -416,6 +421,7 @@ impl crate::globals::resolver::Resolver {
 			.read()
 			.expect("locked for reading")
 			.get(name)
+			.filter(|cached| cached.valid())
 			.cloned()
 	}
 
@@ -431,10 +437,30 @@ impl crate::globals::resolver::Resolver {
 		self.overrides
 			.read()
 			.expect("locked for reading")
-			.contains_key(name)
+			.get(name)
+			.filter(|cached| cached.valid())
+			.is_some()
 	}
 }
 
+impl CachedDest {
+	#[inline]
+	#[must_use]
+	pub fn valid(&self) -> bool { self.expire > SystemTime::now() }
+
+	#[must_use]
+	pub(crate) fn default_expire() -> SystemTime { rand::timepoint_secs(60 * 60 * 18..60 * 60 * 36) }
+}
+
+impl CachedOverride {
+	#[inline]
+	#[must_use]
+	pub fn valid(&self) -> bool { self.expire > SystemTime::now() }
+
+	#[must_use]
+	pub(crate) fn default_expire() -> SystemTime { rand::timepoint_secs(60 * 60 * 6..60 * 60 * 12) }
+}
+
 impl FedDest {
 	fn into_https_string(self) -> String {
 		match self {
diff --git a/src/service/sending/send.rs b/src/service/sending/send.rs
index 1b977a73..a89ea2f8 100644
--- a/src/service/sending/send.rs
+++ b/src/service/sending/send.rs
@@ -111,6 +111,7 @@ where
 			CachedDest {
 				dest: actual.dest.clone(),
 				host: actual.host.clone(),
+				expire: CachedDest::default_expire(),
 			},
 		);
 	}