Move DoH public/private host rules into the DnsQuery

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

@@ -23,6 +23,7 @@ import okhttp3.MediaType
 import okhttp3.MediaType.Companion.toMediaType
 import okhttp3.OkHttpClient
 import okhttp3.dnsoverhttps.internal.DnsOverHttpsQuery
+import okhttp3.internal.concurrent.TaskRunner
 import okhttp3.internal.dns.StateMachineDnsCall
 import okhttp3.internal.dns.execute
 import okhttp3.internal.publicsuffix.PublicSuffixDatabase
@@ -38,6 +39,7 @@ import okhttp3.internal.publicsuffix.PublicSuffixDatabase
  * [doh_spec]: https://tools.ietf.org/html/draft-ietf-doh-dns-over-https-13
  */
 class DnsOverHttps internal constructor(
+  private val taskRunner: TaskRunner,
   @get:JvmName("client") val client: OkHttpClient,
   @get:JvmName("url") val url: HttpUrl,
   @get:JvmName("includeIPv6") val includeIPv6: Boolean,
@@ -48,43 +50,28 @@ class DnsOverHttps internal constructor(
 ) : Dns {
   private val queryFactory =
     DnsOverHttpsQuery.Factory(
+      taskRunner = taskRunner,
+      resolvePrivateAddresses = resolvePrivateAddresses,
+      resolvePublicAddresses = resolvePublicAddresses,
       client = client,
       dnsUrl = url,
       post = post,
     )
 
   override fun newCall(request: Dns.Request): Dns.Call =
     StateMachineDnsCall(
+      taskRunner = taskRunner,
       request = request,
       queryFactory = queryFactory,
-      canceledException = validate(request.hostname),
       includeIPv6 = includeIPv6,
       includeServiceMetadata = includeServiceMetadata,
     )
 
-  /**
-   * Returns an exception if [hostname] should not be resolved.
-   *
-   * We **return** this exception rather than throwing it because in the [Dns.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)
   override fun lookup(hostname: String): List<InetAddress> {
     val withoutServiceMetadata =
       DnsOverHttps(
+        taskRunner = taskRunner,
         client = client,
         url = url,
         includeIPv6 = includeIPv6,
@@ -101,6 +88,7 @@ class DnsOverHttps internal constructor(
   }
 
   class Builder {
+    internal val taskRunner = TaskRunner.INSTANCE
     internal var client: OkHttpClient? = null
     internal var url: HttpUrl? = null
     internal var includeIPv6 = true
@@ -114,13 +102,14 @@ class DnsOverHttps internal constructor(
     fun build(): DnsOverHttps {
       val client = this.client ?: throw NullPointerException("client not set")
       return DnsOverHttps(
-        client.newBuilder().dns(buildBootstrapClient(this)).build(),
-        checkNotNull(url) { "url not set" },
-        includeIPv6,
-        includeServiceMetadata,
-        post,
-        resolvePrivateAddresses,
-        resolvePublicAddresses,
+        taskRunner = taskRunner,
+        client = client.newBuilder().dns(buildBootstrapClient(this)).build(),
+        url = checkNotNull(url) { "url not set" },
+        includeIPv6 = includeIPv6,
+        includeServiceMetadata = includeServiceMetadata,
+        post = post,
+        resolvePrivateAddresses = resolvePrivateAddresses,
+        resolvePublicAddresses = resolvePublicAddresses,
       )
     }
 

okhttp-dnsoverhttps/src/main/kotlin/okhttp3/dnsoverhttps/internal/-DnsOverHttpsQuery.kt

@@ -19,6 +19,7 @@ package okhttp3.dnsoverhttps.internal
 
 import java.io.IOException
 import java.net.ProtocolException
+import java.net.UnknownHostException
 import okhttp3.Call
 import okhttp3.Callback
 import okhttp3.HttpUrl
@@ -29,7 +30,9 @@ import okhttp3.RequestBody
 import okhttp3.Response
 import okhttp3.dnsoverhttps.DnsOverHttps.Companion.DNS_MESSAGE
 import okhttp3.dnsoverhttps.DnsOverHttps.Companion.MAX_RESPONSE_SIZE
+import okhttp3.dnsoverhttps.DnsOverHttps.Companion.isPrivateHost
 import okhttp3.internal.OkHttpInternalApi
+import okhttp3.internal.concurrent.TaskRunner
 import okhttp3.internal.dns.DnsMessage
 import okhttp3.internal.dns.DnsMessageReader
 import okhttp3.internal.dns.DnsMessageWriter
@@ -41,7 +44,7 @@ import okio.BufferedSink
 
 @OkHttpInternalApi
 internal class DnsOverHttpsQuery(
-  val call: Call,
+  private val call: Call,
 ) : DnsQuery {
   override fun enqueue(callback: DnsQuery.Callback) {
     call.enqueue(
@@ -75,11 +78,19 @@ internal class DnsOverHttpsQuery(
   }
 
   class Factory(
+    private val taskRunner: TaskRunner,
+    private val resolvePrivateAddresses: Boolean,
+    private val resolvePublicAddresses: Boolean,
     private val client: OkHttpClient,
     private val dnsUrl: HttpUrl,
     private val post: Boolean,
   ) : DnsQuery.Factory {
     override fun newQuery(question: Question): DnsQuery {
+      val overrideQuery = overrideQuery(question)
+      if (overrideQuery != null) {
+        return overrideQuery
+      }
+
       val dnsMessage = DnsMessage.query(question)
       return DnsOverHttpsQuery(
         call =
@@ -111,6 +122,40 @@ internal class DnsOverHttpsQuery(
           ),
       )
     }
+
+    /** Returns an alternate query for [question] if the [DnsOverHttpsQuery] should not be used. */
+    private fun overrideQuery(question: Question): DnsQuery? {
+      // Don't load the public suffix list unless necessary.
+      if (resolvePrivateAddresses && resolvePublicAddresses) return null
+
+      val privateHost = isPrivateHost(question.name)
+
+      val resolveUnsupportedMessage =
+        when {
+          privateHost && !resolvePrivateAddresses -> "private hosts not resolved"
+          !privateHost && !resolvePublicAddresses -> "public hosts not resolved"
+          else -> return null
+        }
+
+      return ResolveUnsupportedQuery(
+        hostname = question.name,
+        message = resolveUnsupportedMessage,
+      )
+    }
+
+    private inner class ResolveUnsupportedQuery(
+      private val hostname: String,
+      private val message: String,
+    ) : DnsQuery {
+      override fun enqueue(callback: DnsQuery.Callback) {
+        taskRunner.newQueue().execute("$hostname dns") {
+          callback.onFailure(UnknownHostException(message))
+        }
+      }
+
+      override fun cancel() {
+      }
+    }
   }
 }
 

okhttp/api/android/okhttp.api

@@ -538,7 +538,6 @@ public final class okhttp3/Dns$Request {
 }
 
 public final class okhttp3/DnsCache {
-	public final field -delegate Lokhttp3/internal/dns/RealDnsCache;
 	public synthetic fun <init> (JJJJIILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public synthetic fun <init> (JJJJILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun evictAll ()V

okhttp/api/jvm/okhttp.api

@@ -538,7 +538,6 @@ public final class okhttp3/Dns$Request {
 }
 
 public final class okhttp3/DnsCache {
-	public final field -delegate Lokhttp3/internal/dns/RealDnsCache;
 	public synthetic fun <init> (JJJJIILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public synthetic fun <init> (JJJJILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun evictAll ()V

okhttp/src/androidMain/kotlin/okhttp3/android/AndroidDns.kt

@@ -87,9 +87,9 @@ class AndroidDns(
     includeServiceMetadata: Boolean,
   ): Dns.Call =
     StateMachineDnsCall(
+      taskRunner = taskRunner,
       request = request,
       queryFactory = queryFactory,
-      canceledException = null,
       // A single `A` query stands in for both families: the system resolver returns IPv4 and
       // IPv6 addresses together, so there's no separate `AAAA` query.
       includeIPv6 = false,

okhttp/src/commonJvmAndroid/kotlin/okhttp3/DnsCache.kt

@@ -76,7 +76,7 @@ import okhttp3.internal.dns.RealDnsCache
 class DnsCache internal constructor(
   @OkHttpInternalApi
   @JvmField
-  val `-delegate`: RealDnsCache,
+  internal val `-delegate`: RealDnsCache,
 ) {
   @get:JvmName("minimumTimeToLive")
   val minimumTimeToLive: Duration

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/dns/-StateMachineDnsCall.kt

@@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicReference
 import okhttp3.Dns
 import okhttp3.Protocol
 import okhttp3.internal.OkHttpInternalApi
+import okhttp3.internal.concurrent.TaskRunner
 
 /**
  * An application-layer [Dns.Call] that performs multiple transport-layer [DnsQuery]s in parallel.
@@ -53,9 +54,9 @@ import okhttp3.internal.OkHttpInternalApi
  */
 @OkHttpInternalApi
 class StateMachineDnsCall(
+  private val taskRunner: TaskRunner,
   override val request: Dns.Request,
   private val queryFactory: DnsQuery.Factory,
-  private val canceledException: IOException?,
   private val includeIPv6: Boolean,
   private val includeServiceMetadata: Boolean,
 ) : Dns.Call {
@@ -75,30 +76,39 @@ class StateMachineDnsCall(
         add(Question(request.hostname, TYPE_A))
       }
 
-    val queries =
-      questions.map { question ->
-        queryFactory.newQuery(question)
-      }
-
     while (true) {
       val previous =
         state.get() as? State.Idle
           ?: error("already enqueued")
 
+      // If it's canceled before it is enqueued, jump straight to Complete.
+      if (previous.canceled) {
+        val next = State.Complete(canceled = true)
+
+        if (!state.compareAndSet(previous, next)) continue // Lost a race, retry.
+
+        taskRunner.newQueue().execute("${request.hostname} dns") {
+          callback.onFailure(this, IOException("canceled"))
+        }
+
+        return
+      }
+
+      val queries =
+        questions.map { question ->
+          queryFactory.newQuery(question)
+        }
+
       val next =
         State.Running(
-          canceled = previous.canceled,
+          canceled = false,
           callback = callback,
           runningQueries = queries,
         )
 
       if (!state.compareAndSet(previous, next)) continue // Lost a race, retry.
 
       for (query in queries) {
-        if (previous.canceled || canceledException != null) {
-          query.cancel()
-        }
-
         query.enqueue(
           callback =
             object : DnsQuery.Callback {
@@ -210,7 +220,6 @@ class StateMachineDnsCall(
 
       val allExceptions =
         when {
-          canceledException != null -> listOf(canceledException)
           newException != null -> previous.pendingExceptions + newException
           else -> previous.pendingExceptions
         }

okhttp/src/commonTest/kotlin/okhttp3/internal/dns/StateMachineDnsCallTest.kt

@@ -841,59 +841,22 @@ class StateMachineDnsCallTest {
       }
     }
 
-  /**
-   * The implementation still enqueues canceled queries, because that's an easy way to jump to a
-   * dispatcher thread to post the failures back to the callback.
-   */
-  @Test
-  fun `cancel before enqueue`() =
-    testStateMachineDnsCall {
-      val call =
-        newCall(
-          request = Dns.Request(hostname = "lysine.dev"),
-          caching = false,
-        )
-      call.cancel()
-      call.enqueue()
-
-      queryFactory.takeCancel("lysine.dev", TYPE_HTTPS)
-      val query0 = queryFactory.takeQuery("lysine.dev", TYPE_HTTPS)
-      queryFactory.takeCancel("lysine.dev", TYPE_AAAA)
-      val query1 = queryFactory.takeQuery("lysine.dev", TYPE_AAAA)
-      queryFactory.takeCancel("lysine.dev", TYPE_A)
-      val query2 = queryFactory.takeQuery("lysine.dev", TYPE_A)
-
-      query0.respondFailure("canceled")
-      query1.respondFailure("canceled")
-      query2.respondFailure("canceled")
-
-      call.takeOnFailure("canceled")
-    }
-
   @Test
-  fun `cancel before enqueue with caching`() =
+  fun `cancel before enqueue`(caching: Boolean) =
     testStateMachineDnsCall {
       val call =
         newCall(
           request = Dns.Request(hostname = "lysine.dev"),
-          caching = true,
+          caching = caching,
         )
       call.cancel()
       call.enqueue()
 
-      val query0 = queryFactory.takeQuery("lysine.dev", TYPE_HTTPS)
-      val query1 = queryFactory.takeQuery("lysine.dev", TYPE_AAAA)
-      val query2 = queryFactory.takeQuery("lysine.dev", TYPE_A)
-
-      query0.respondFailure("canceled")
-      query1.respondFailure("canceled")
-      query2.respondFailure("canceled")
-
       call.takeOnFailure("canceled")
 
-      assertThat(cache.size).isEqualTo(3)
-      assertThat(cache.requestCount).isEqualTo(3)
-      assertThat(cache.networkCount).isEqualTo(3)
+      assertThat(cache.size).isEqualTo(0)
+      assertThat(cache.requestCount).isEqualTo(0)
+      assertThat(cache.networkCount).isEqualTo(0)
       assertThat(cache.hitCount).isEqualTo(0)
     }
 

okhttp/src/commonTest/kotlin/okhttp3/internal/dns/StateMachineDnsCallTester.kt

@@ -146,13 +146,13 @@ class StateMachineDnsCallTester internal constructor() {
 
     val call =
       StateMachineDnsCall(
+        taskRunner = taskFaker.taskRunner,
         request = request,
         queryFactory =
           when {
             caching -> cache.wrap(queryFactory)
             else -> queryFactory
           },
-        canceledException = null,
         includeIPv6 = includeIPv6,
         includeServiceMetadata = includeServiceMetadata,
       )

okhttp/src/jvmMain/java9/module-info.java

@@ -5,7 +5,7 @@
   requires java.logging;
   exports okhttp3;
   exports okhttp3.internal to okhttp3.logging, okhttp3.sse, okhttp3.java.net.cookiejar, okhttp3.dnsoverhttps, mockwebserver3, okhttp3.mockwebserver, mockwebserver3.junit5, okhttp3.coroutines, okhttp3.tls;
-  exports okhttp3.internal.concurrent to mockwebserver3, okhttp3.mockwebserver;
+  exports okhttp3.internal.concurrent to mockwebserver3, okhttp3.dnsoverhttps, okhttp3.mockwebserver;
   exports okhttp3.internal.connection to mockwebserver3, okhttp3.mockwebserver, okhttp3.logging;
   exports okhttp3.internal.dns to okhttp3.dnsoverhttps;
   exports okhttp3.internal.http to okhttp3.logging, okhttp3.brotli, mockwebserver3;