From 22f4c0600e0fcc8aef8bd4b87094c0fb0deaff81 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9rgio=20Santos?= <me@sergiosantos.info>
Date: Mon, 6 Nov 2023 12:32:52 +0000
Subject: [PATCH] fix: Improve connection state checks speed

---
 .../gateway/background/PingRemoteServer.kt    | 20 ++++++++++++++++---
 .../gateway/data/doh/ResolveServiceAddress.kt | 12 ++++++++++-
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/app/src/main/java/tech/relaycorp/gateway/background/PingRemoteServer.kt b/app/src/main/java/tech/relaycorp/gateway/background/PingRemoteServer.kt
index 48216e94..f81a0c56 100644
--- a/app/src/main/java/tech/relaycorp/gateway/background/PingRemoteServer.kt
+++ b/app/src/main/java/tech/relaycorp/gateway/background/PingRemoteServer.kt
@@ -8,11 +8,14 @@ import io.ktor.client.request.head
 import io.ktor.network.selector.ActorSelectorManager
 import io.ktor.network.sockets.aSocket
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.withTimeout
 import tech.relaycorp.gateway.BuildConfig
 import tech.relaycorp.gateway.common.Logging.logger
 import java.io.IOException
 import java.util.logging.Level
 import javax.inject.Inject
+import kotlin.time.Duration.Companion.seconds
 
 class PingRemoteServer
 @Inject constructor() {
@@ -29,7 +32,9 @@ class PingRemoteServer
         try {
             aSocket(ActorSelectorManager(Dispatchers.IO))
                 .tcp()
-                .connect(address, port)
+                .connect(address, port) {
+                    socketTimeout = TIMEOUT.inWholeMilliseconds
+                }
                 .use { true }
         } catch (e: IOException) {
             logger.log(Level.INFO, "Could not ping $address:$port")
@@ -38,11 +43,16 @@ class PingRemoteServer
 
     suspend fun pingURL(url: String) =
         try {
-            ktorClient.head<Unit>(url)
-            true
+            withTimeout(TIMEOUT) {
+                ktorClient.head<Unit>(url)
+                true
+            }
         } catch (e: IOException) {
             logger.log(Level.INFO, "Could not ping $url (${e.message})")
             false
+        } catch (e: TimeoutCancellationException) {
+            logger.log(Level.INFO, "Could not ping $url (${e.message})")
+            false
         } catch (e: ResponseException) {
             logger.log(
                 Level.INFO,
@@ -50,4 +60,8 @@ class PingRemoteServer
             )
             true
         }
+
+    companion object {
+        private val TIMEOUT = 5.seconds
+    }
 }
diff --git a/app/src/main/java/tech/relaycorp/gateway/data/doh/ResolveServiceAddress.kt b/app/src/main/java/tech/relaycorp/gateway/data/doh/ResolveServiceAddress.kt
index f00c4b07..0922ac32 100644
--- a/app/src/main/java/tech/relaycorp/gateway/data/doh/ResolveServiceAddress.kt
+++ b/app/src/main/java/tech/relaycorp/gateway/data/doh/ResolveServiceAddress.kt
@@ -1,9 +1,12 @@
 package tech.relaycorp.gateway.data.doh
 
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.withTimeout
 import tech.relaycorp.doh.DoHClient
 import tech.relaycorp.doh.LookupFailureException
 import tech.relaycorp.gateway.data.model.ServiceAddress
 import javax.inject.Inject
+import kotlin.time.Duration.Companion.seconds
 
 class ResolveServiceAddress
 @Inject constructor(
@@ -13,12 +16,19 @@ class ResolveServiceAddress
     suspend fun resolvePoWeb(address: String): ServiceAddress {
         val srvRecordName = "_awala-gsc._tcp.$address"
         val answer = try {
-            doHClient.lookUp(srvRecordName, "SRV")
+            withTimeout(5.seconds) {
+                doHClient.lookUp(srvRecordName, "SRV")
+            }
         } catch (exc: LookupFailureException) {
             throw InternetAddressResolutionException(
                 "Failed to resolve DNS for PoWeb address",
                 exc
             )
+        } catch (exc: TimeoutCancellationException) {
+            throw InternetAddressResolutionException(
+                "Failed to resolve DNS for PoWeb address",
+                exc
+            )
         }
         val srvRecordData = answer.data.first()
         val srvRecordDataFields = srvRecordData.split(" ")