From 2b425b29bc7f924fe2c9049755f9e111e215142a Mon Sep 17 00:00:00 2001
From: Joakim Hulthe <joakim@hulthe.net>
Date: Wed, 27 Mar 2024 14:00:27 +0100
Subject: [PATCH] Add timeout to tester geoip rcp call

---
 test/test-manager/src/tests/relay_ip_overrides.rs |  3 ---
 test/test-rpc/src/net.rs                          | 15 ++++++++++++---
 test/test-runner/src/main.rs                      | 14 +++++++++++---
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/test/test-manager/src/tests/relay_ip_overrides.rs b/test/test-manager/src/tests/relay_ip_overrides.rs
index 1a33639bf227..bc05e8ab5204 100644
--- a/test/test-manager/src/tests/relay_ip_overrides.rs
+++ b/test/test-manager/src/tests/relay_ip_overrides.rs
@@ -72,7 +72,6 @@ pub async fn test_wireguard_ip_override(
     let _remove_nft_rule_on_drop = block_route(guest_ip, relay_ip).await?;
 
     log::info!("checking that the connection does not work while blocked");
-    // FIXME: this fails because of rpc timeouts, which is sort of fine but not ideal
     ensure!(
         helpers::geoip_lookup_with_retries(&rpc).await.is_err(),
         "Assert that relay is blocked by firewall rule"
@@ -129,7 +128,6 @@ pub async fn test_openvpn_ip_override(
     let _remove_nft_rule_on_drop = block_route(guest_ip, relay_ip).await?;
 
     log::info!("checking that the connection does not work while blocked");
-    // FIXME: this fails because of rpc timeouts, which is sort of fine but not ideal
     ensure!(
         helpers::geoip_lookup_with_retries(&rpc).await.is_err(),
         "Assert that relay is blocked by firewall rule"
@@ -212,7 +210,6 @@ pub async fn test_bridge_ip_override(
     let _remove_nft_rule_on_drop = block_route(guest_ip, relay_ip).await?;
 
     log::info!("checking that the connection does not work while blocked");
-    // FIXME: this fails because of rpc timeouts, which is sort of fine but not ideal
     ensure!(
         helpers::geoip_lookup_with_retries(&rpc).await.is_err(),
         "Assert that relay is blocked by firewall rule"
diff --git a/test/test-rpc/src/net.rs b/test/test-rpc/src/net.rs
index 77aa5c938abf..124b4f9ad36f 100644
--- a/test/test-rpc/src/net.rs
+++ b/test/test-rpc/src/net.rs
@@ -2,7 +2,7 @@ use futures::channel::oneshot;
 use hyper::{Client, Uri};
 use once_cell::sync::Lazy;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
-use std::net::SocketAddr;
+use std::{net::SocketAddr, time::Duration};
 use tokio_rustls::rustls::ClientConfig;
 
 use crate::{AmIMullvad, Error};
@@ -72,10 +72,10 @@ impl Drop for SockHandle {
     }
 }
 
-pub async fn geoip_lookup(mullvad_host: String) -> Result<AmIMullvad, Error> {
+pub async fn geoip_lookup(mullvad_host: String, timeout: Duration) -> Result<AmIMullvad, Error> {
     let uri = Uri::try_from(format!("https://ipv4.am.i.{mullvad_host}/json"))
         .map_err(|_| Error::InvalidUrl)?;
-    http_get(uri).await
+    http_get_with_timeout(uri, timeout).await
 }
 
 pub async fn http_get<T: DeserializeOwned>(url: Uri) -> Result<T, Error> {
@@ -106,6 +106,15 @@ pub async fn http_get<T: DeserializeOwned>(url: Uri) -> Result<T, Error> {
     })
 }
 
+pub async fn http_get_with_timeout<T: DeserializeOwned>(
+    url: Uri,
+    timeout: Duration,
+) -> Result<T, Error> {
+    tokio::time::timeout(timeout, http_get(url))
+        .await
+        .map_err(|_| Error::HttpRequest("Request timed out".into()))?
+}
+
 fn read_cert_store() -> tokio_rustls::rustls::RootCertStore {
     let mut cert_store = tokio_rustls::rustls::RootCertStore::empty();
 
diff --git a/test/test-runner/src/main.rs b/test/test-runner/src/main.rs
index d864968bbee5..8399e118d368 100644
--- a/test/test-runner/src/main.rs
+++ b/test/test-runner/src/main.rs
@@ -6,7 +6,7 @@ use std::{
     path::{Path, PathBuf},
     process::Stdio,
     sync::Arc,
-    time::Duration,
+    time::{Duration, SystemTime},
 };
 use util::OnDrop;
 
@@ -171,10 +171,18 @@ impl Service for TestServer {
 
     async fn geoip_lookup(
         self,
-        _: context::Context,
+        ctx: context::Context,
         mullvad_host: String,
     ) -> Result<test_rpc::AmIMullvad, test_rpc::Error> {
-        test_rpc::net::geoip_lookup(mullvad_host).await
+        let timeout = ctx
+            .deadline
+            .duration_since(SystemTime::now())
+            .ok()
+            // account for some time to send the RPC response
+            .and_then(|d| d.checked_sub(Duration::from_millis(500)))
+            .unwrap_or_default();
+
+        test_rpc::net::geoip_lookup(mullvad_host, timeout).await
     }
 
     async fn resolve_hostname(