DnsOverHttps Fixups: event sequencing, resolvePublicAddresses, IPv6 first

okhttp-dnsoverhttps/api/okhttp-dnsoverhttps.api

@@ -18,6 +18,7 @@ public final class okhttp3/dnsoverhttps/DnsOverHttps$Builder {
 	public final fun bootstrapDnsHosts ([Ljava/net/InetAddress;)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;
 	public final fun build ()Lokhttp3/dnsoverhttps/DnsOverHttps;
 	public final fun client (Lokhttp3/OkHttpClient;)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;
+	public final fun includeHttps (Z)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;
 	public final fun includeIPv6 (Z)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;
 	public final fun post (Z)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;
 	public final fun resolvePrivateAddresses (Z)Lokhttp3/dnsoverhttps/DnsOverHttps$Builder;

okhttp-dnsoverhttps/src/main/kotlin/okhttp3/dnsoverhttps/-DnsMessage.kt

@@ -17,9 +17,16 @@
 
 package okhttp3.dnsoverhttps
 
+import java.io.IOException
 import java.net.InetAddress
+import java.net.ProtocolException
+import java.net.UnknownHostException
+import okhttp3.Protocol
 import okhttp3.RequestBody
+import okhttp3.Response
 import okhttp3.dnsoverhttps.DnsOverHttps.Companion.DNS_MESSAGE
+import okhttp3.dnsoverhttps.DnsOverHttps.Companion.MAX_RESPONSE_SIZE
+import okhttp3.internal.platform.Platform
 import okio.Buffer
 import okio.BufferedSink
 import okio.ByteString
@@ -109,8 +116,8 @@ internal sealed interface ResourceRecord {
   data class Https(
     override val name: String,
     override val timeToLive: Int,
-    val priority: Int,
-    val targetName: String,
+    val priority: Int = 1,
+    val targetName: String = "",
     val alpnIds: List<String>? = null,
     var port: Int = 443,
     val ipAddressHints: List<InetAddress> = listOf(),
@@ -169,3 +176,34 @@ internal class QueryRequestBody(
     sink.emitCompleteSegments()
   }
 }
+
+@Throws(IOException::class)
+internal fun decodeResponse(response: Response): List<ResourceRecord> {
+  if (
+    response.cacheResponse == null &&
+    response.protocol !== Protocol.HTTP_2 &&
+    response.protocol !== Protocol.QUIC
+  ) {
+    Platform.get().log("Unexpected protocol: ${response.protocol}", Platform.WARN)
+  }
+
+  response.use {
+    if (!response.isSuccessful) {
+      throw IOException("response: ${response.code} ${response.message}")
+    }
+
+    val body = response.body
+    if (body.contentLength() > MAX_RESPONSE_SIZE) {
+      throw ProtocolException(
+        "response size exceeds limit ($MAX_RESPONSE_SIZE bytes): ${body.contentLength()} bytes",
+      )
+    }
+
+    val dnsResponse = DnsMessageReader(body.source()).read()
+    when (dnsResponse.responseCode) {
+      RESPONSE_CODE_SUCCESS -> return dnsResponse.answers
+      RESPONSE_CODE_SERVER_FAILURE -> throw UnknownHostException("DNS server failure")
+      else -> throw UnknownHostException()
+    }
+  }
+}

okhttp-dnsoverhttps/src/main/kotlin/okhttp3/dnsoverhttps/-DnsOverHttpsCall.kt

@@ -36,43 +36,24 @@ import okhttp3.internal.testAndSet
  * Implements [Dns2.Call] by making multiple HTTPS calls.
  */
 internal class DnsOverHttpsCall(
-  private val dnsOverHttps: DnsOverHttps,
   override val request: Dns2.Request,
+  private val calls: List<Call>,
+  private val canceledException: IOException?,
 ) : Dns2.Call,
   Callback {
+  @Volatile private var canceled = false
   private val state = AtomicReference<State>(State.Idle)
 
   override fun enqueue(callback: Dns2.Callback) {
-    val calls =
-      buildList {
-        if (dnsOverHttps.includeHttps) {
-          add(dnsOverHttps.createCall(request.hostname, TYPE_HTTPS))
-        }
-
-        add(dnsOverHttps.createCall(request.hostname, TYPE_A))
-
-        if (dnsOverHttps.includeIPv6) {
-          add(dnsOverHttps.createCall(request.hostname, TYPE_AAAA))
-        }
-      }
-
     val running = State.Running(callback, calls)
 
     val previous = state.testAndSet(running) { it is State.Idle }
-    when (previous) {
-      is State.Idle -> {
-        for (call in calls) {
-          call.enqueue(this)
-        }
-      }
-
-      is State.Running, State.Complete -> {
-        throw IllegalStateException("already enqueued")
-      }
+    check(previous is State.Idle) {
+      "already enqueued"
+    }
 
-      State.Canceled -> {
-        callback.onFailure(this, IOException("canceled"))
-      }
+    for (call in calls) {
+      call.enqueue(this)
     }
   }
 
@@ -90,7 +71,11 @@ internal class DnsOverHttpsCall(
           ?: return // Already complete or canceled; nothing to do.
 
       val newRunningCalls = previous.runningCalls - call
-      val allFailures = previous.delayedFailures + e
+      val allFailures =
+        when {
+          canceledException != null -> listOf(canceledException)
+          else -> previous.delayedFailures + e
+        }
       val next =
         when {
           newRunningCalls.isEmpty() -> {
@@ -125,7 +110,7 @@ internal class DnsOverHttpsCall(
   ) {
     val resourceRecords =
       try {
-        dnsOverHttps.decodeResponse(response)
+        decodeResponse(response)
       } catch (e: IOException) {
         return onFailure(call, e)
       }
@@ -165,9 +150,10 @@ internal class DnsOverHttpsCall(
           ?: return // Already complete or canceled; nothing to do.
 
       val newRunningCalls = previous.runningCalls - call
+      val lastRunningCall = newRunningCalls.isEmpty()
       val next =
         when {
-          newRunningCalls.isEmpty() -> {
+          lastRunningCall -> {
             State.Complete
           }
 
@@ -182,22 +168,36 @@ internal class DnsOverHttpsCall(
 
       if (!state.compareAndSet(previous, next)) continue // Lost a race; retry.
 
-      previous.callback.onRecords(
-        call = this,
-        last = newRunningCalls.isEmpty(),
-        records = dns2Records,
-      )
+      val emitDelayedFailures = lastRunningCall && previous.delayedFailures.isNotEmpty()
+      val lastEvent = lastRunningCall && !emitDelayedFailures
+      if (dns2Records.isNotEmpty() || lastEvent) {
+        previous.callback.onRecords(
+          call = this,
+          last = lastEvent,
+          records = dns2Records,
+        )
+      }
+      if (emitDelayedFailures) {
+        previous.callback.onFailure(
+          call = this,
+          exceptions = previous.delayedFailures,
+        )
+      }
 
       return
     }
   }
 
   override fun cancel() {
-    val previous = state.testAndSet(State.Canceled) { it is State.Running || it == State.Idle }
-    (previous as? State.Running)?.callback?.onFailure(this, IOException("canceled"))
+    if (canceled) return // Already canceled.
+
+    canceled = true
+    for (call in calls) {
+      call.cancel()
+    }
   }
 
-  override fun isCanceled() = state.get() is State.Canceled
+  override fun isCanceled() = canceled
 
   private sealed interface State {
     object Idle : State
@@ -209,8 +209,6 @@ internal class DnsOverHttpsCall(
     ) : State
 
     object Complete : State
-
-    object Canceled : State
   }
 }
 

okhttp-dnsoverhttps/src/main/kotlin/okhttp3/dnsoverhttps/DnsOverHttps.kt

@@ -17,7 +17,6 @@ package okhttp3.dnsoverhttps
 
 import java.io.IOException
 import java.net.InetAddress
-import java.net.ProtocolException
 import java.net.UnknownHostException
 import java.util.concurrent.CountDownLatch
 import okhttp3.Call
@@ -28,10 +27,8 @@ import okhttp3.HttpUrl
 import okhttp3.MediaType
 import okhttp3.MediaType.Companion.toMediaType
 import okhttp3.OkHttpClient
-import okhttp3.Protocol
 import okhttp3.Request
 import okhttp3.Response
-import okhttp3.internal.platform.Platform
 import okhttp3.internal.publicsuffix.PublicSuffixDatabase
 
 /**
@@ -54,34 +51,48 @@ class DnsOverHttps internal constructor(
   @get:JvmName("resolvePublicAddresses") val resolvePublicAddresses: Boolean,
 ) : Dns,
   Dns2 {
-  override fun newCall(request: Dns2.Request): Dns2.Call = DnsOverHttpsCall(this, request)
+  override fun newCall(request: Dns2.Request): Dns2.Call {
+    val calls = callsList(request.hostname)
 
-  @Throws(UnknownHostException::class)
-  override fun lookup(hostname: String): List<InetAddress> {
-    if (!resolvePrivateAddresses || !resolvePublicAddresses) {
-      val privateHost = isPrivateHost(hostname)
-
-      if (privateHost && !resolvePrivateAddresses) {
-        throw UnknownHostException("private hosts not resolved")
-      }
-
-      if (!privateHost && !resolvePublicAddresses) {
-        throw UnknownHostException("public hosts not resolved")
+    val canceledException = validate(request.hostname)
+    if (canceledException != null) {
+      for (call in calls) {
+        call.cancel()
       }
     }
 
-    return lookupHttps(hostname)
+    return DnsOverHttpsCall(
+      request = request,
+      calls = calls,
+      canceledException = canceledException,
+    )
+  }
+
+  /**
+   * Returns an exception if [hostname] should not be resolved.
+   *
+   * We **return** this exception rather than throwing it because in the [Dns2.Callback] case we want
+   * `onFailure()` to be called on a dispatcher thread and not synchronously.
+   */
+  private fun validate(hostname: String): UnknownHostException? {
+    // Don't load the public suffix list unless necessary.
+    if (resolvePrivateAddresses && resolvePublicAddresses) return null
+
+    val privateHost = isPrivateHost(hostname)
+
+    return when {
+      privateHost && !resolvePrivateAddresses -> UnknownHostException("private hosts not resolved")
+      !privateHost && !resolvePublicAddresses -> UnknownHostException("public hosts not resolved")
+      else -> null
+    }
   }
 
   @Throws(UnknownHostException::class)
-  private fun lookupHttps(hostname: String): List<InetAddress> {
-    val calls =
-      buildList {
-        add(createCall(hostname, TYPE_A))
-        if (includeIPv6) {
-          add(createCall(hostname, TYPE_AAAA))
-        }
-      }
+  override fun lookup(hostname: String): List<InetAddress> {
+    val validationException = validate(hostname)
+    if (validationException != null) throw validationException
+
+    val calls = callsList(hostname, inetAddressesOnly = true)
 
     val failures = ArrayList<Exception>(3)
     val results = ArrayList<InetAddress>(5)
@@ -175,37 +186,6 @@ class DnsOverHttps internal constructor(
     throw unknownHostException
   }
 
-  @Throws(IOException::class)
-  internal fun decodeResponse(response: Response): List<ResourceRecord> {
-    if (
-      response.cacheResponse == null &&
-      response.protocol !== Protocol.HTTP_2 &&
-      response.protocol !== Protocol.QUIC
-    ) {
-      Platform.get().log("Incorrect protocol: ${response.protocol}", Platform.WARN)
-    }
-
-    response.use {
-      if (!response.isSuccessful) {
-        throw IOException("response: ${response.code} ${response.message}")
-      }
-
-      val body = response.body
-      if (body.contentLength() > MAX_RESPONSE_SIZE) {
-        throw ProtocolException(
-          "response size exceeds limit ($MAX_RESPONSE_SIZE bytes): ${body.contentLength()} bytes",
-        )
-      }
-
-      val dnsResponse = DnsMessageReader(body.source()).read()
-      when (dnsResponse.responseCode) {
-        RESPONSE_CODE_SUCCESS -> return dnsResponse.answers
-        RESPONSE_CODE_SERVER_FAILURE -> throw UnknownHostException("DNS server failure")
-        else -> throw UnknownHostException()
-      }
-    }
-  }
-
   internal fun createCall(
     hostname: String,
     type: Int,
@@ -238,6 +218,22 @@ class DnsOverHttps internal constructor(
           }.build(),
     )
 
+  private fun callsList(
+    hostname: String,
+    inetAddressesOnly: Boolean = false,
+  ): List<Call> =
+    buildList {
+      if (includeHttps && !inetAddressesOnly) {
+        add(createCall(hostname, TYPE_HTTPS))
+      }
+
+      if (includeIPv6) {
+        add(createCall(hostname, TYPE_AAAA))
+      }
+
+      add(createCall(hostname, TYPE_A))
+    }
+
   class Builder {
     internal var client: OkHttpClient? = null
     internal var url: HttpUrl? = null
@@ -278,6 +274,11 @@ class DnsOverHttps internal constructor(
      *
      * This is false by default, but that default is subject to change in 2026.
      */
+    fun includeHttps(includeHttps: Boolean) =
+      apply {
+        this.includeHttps = includeHttps
+      }
+
     fun includeIPv6(includeIPv6: Boolean) =
       apply {
         this.includeIPv6 = includeIPv6