Try every resolved address when connecting a TCP socket

common/src/main/java/com/pedro/common/socket/java/TcpStreamSocketJava.kt

@@ -2,9 +2,9 @@ package com.pedro.common.socket.java
 
 import android.os.Build
 import java.io.IOException
+import java.net.InetAddress
 import java.net.InetSocketAddress
 import java.net.Socket
-import java.net.SocketAddress
 import java.security.GeneralSecurityException
 import java.security.SecureRandom
 import javax.net.ssl.SSLContext
@@ -18,8 +18,51 @@ class TcpStreamSocketJava(
     private val certificates: TrustManager? = null
 ): TcpStreamSocketJavaBase() {
 
+    /**
+     * Resolves the host to every address it has. Replaceable in tests.
+     */
+    internal var resolver: (String) -> Array<InetAddress> = { InetAddress.getAllByName(it) }
+
+    /**
+     * Tries every address the host resolves to, in the order the system returns them, instead of
+     * only the first one. A host with both IPv6 and IPv4 addresses usually resolves IPv6 first; on a
+     * network where IPv6 is assigned but broken that address refuses or times out, and the stream
+     * used to fail even though the same server was reachable over IPv4.
+     * The whole attempt keeps the original timeout: every address gets an equal share of the time
+     * that is left, so a slow address can't use up the time of the ones after it.
+     */
     override fun onConnectSocket(timeout: Long): Socket {
-        val socket = if (secured) {
+        return configure(connectAny(timeout), timeout)
+    }
+
+    /**
+     * Only the TCP connect moves on to the next address. A TLS handshake failure is not about the
+     * address, so it is reported as it is instead of being replaced by the error of the last address.
+     */
+    private fun connectAny(timeout: Long): Socket {
+        val addresses = resolver(host)
+        val deadline = System.currentTimeMillis() + timeout
+        var lastError: IOException? = null
+        for ((index, address) in addresses.withIndex()) {
+            val remaining = deadline - System.currentTimeMillis()
+            if (remaining <= 0) break
+            val attemptTimeout = remaining / (addresses.size - index)
+            val socket = createSocket()
+            try {
+                // Keeping the host name in the address keeps it for SNI and certificate checks.
+                val endpoint = InetSocketAddress(InetAddress.getByAddress(host, address.address), port)
+                socket.connect(endpoint, attemptTimeout.toInt().coerceAtLeast(1))
+                return socket
+            } catch (e: IOException) {
+                runCatching { socket.close() }
+                lastError = e
+            }
+        }
+        throw lastError ?: IOException("Unable to connect to $host")
+    }
+
+    private fun createSocket(): Socket {
+        return if (secured) {
             try {
                 val context = SSLContext.getInstance("TLS")
                 context.init(null, certificates?.let { arrayOf(it) }, SecureRandom())
@@ -28,8 +71,9 @@ class TcpStreamSocketJava(
                 throw IOException("Create SSL socket failed: ${e.message}")
             }
         } else Socket()
-        val socketAddress: SocketAddress = InetSocketAddress(host, port)
-        socket.connect(socketAddress, timeout.toInt())
+    }
+
+    private fun configure(socket: Socket, timeout: Long): Socket {
         socket.soTimeout = timeout.toInt()
         if (hostVerification && socket is SSLSocket && Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
             socket.sslParameters = socket.sslParameters.apply {

common/src/test/java/com/pedro/common/socket/TcpStreamSocketJavaTest.kt

@@ -0,0 +1,89 @@
+/*
+ * Copyright (C) 2024 pedroSG94.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.pedro.common.socket.java
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Test
+import java.io.IOException
+import java.net.InetAddress
+import java.net.ServerSocket
+
+class TcpStreamSocketJavaTest {
+
+  private val ipv4 = InetAddress.getByName("127.0.0.1")
+  // Nothing listens on ::1 in these tests (the server binds 127.0.0.1 only), so a connect there fails
+  // at once, the same way a refused or unusable IPv6 address does.
+  private val ipv6 = InetAddress.getByName("::1")
+
+  @Test
+  fun `falls back to the next address when the first one fails`() {
+    ServerSocket(0, 1, ipv4).use { server ->
+      val socket = TcpStreamSocketJava("example.test", server.localPort, secured = false)
+      socket.resolver = { arrayOf(ipv6, ipv4) }
+      val connected = socket.onConnectSocket(3000)
+      try {
+        assertTrue(connected.isConnected)
+        assertEquals(ipv4, connected.inetAddress)
+      } finally {
+        connected.close()
+      }
+    }
+  }
+
+  @Test
+  fun `uses the first address when it works`() {
+    ServerSocket(0, 1, ipv4).use { server ->
+      val socket = TcpStreamSocketJava("example.test", server.localPort, secured = false)
+      socket.resolver = { arrayOf(ipv4, ipv6) }
+      val connected = socket.onConnectSocket(3000)
+      try {
+        assertEquals(ipv4, connected.inetAddress)
+      } finally {
+        connected.close()
+      }
+    }
+  }
+
+  @Test
+  fun `keeps the host name for the connected address`() {
+    ServerSocket(0, 1, ipv4).use { server ->
+      val socket = TcpStreamSocketJava("example.test", server.localPort, secured = false)
+      socket.resolver = { arrayOf(ipv4) }
+      val connected = socket.onConnectSocket(3000)
+      try {
+        assertEquals("example.test", connected.inetAddress.hostName)
+      } finally {
+        connected.close()
+      }
+    }
+  }
+
+  @Test
+  fun `throws when no address connects`() {
+    val closedPort = ServerSocket(0, 1, ipv4).use { it.localPort }
+    val socket = TcpStreamSocketJava("example.test", closedPort, secured = false)
+    socket.resolver = { arrayOf(ipv6, ipv4) }
+    try {
+      socket.onConnectSocket(2000)
+      fail("connect should fail when every address fails")
+    } catch (e: IOException) {
+      // expected
+    }
+  }
+}