issue #1378 WKD client timeout test

FlowCrypt/build.gradle

@@ -386,6 +386,7 @@ dependencies {
     androidTestImplementation "com.squareup.okhttp3:okhttp-tls:${rootProject.ext.okhttpVersion}"
     androidTestUtil 'androidx.test:orchestrator:1.4.0'
 
+    testImplementation "com.squareup.okhttp3:mockwebserver:${rootProject.ext.okhttpVersion}"
     testImplementation "junit:junit:${rootProject.ext.junitVersion}"
     testImplementation "androidx.room:room-testing:$roomVersion"
     testImplementation 'org.robolectric:robolectric:4.6.1'

FlowCrypt/src/main/java/com/flowcrypt/email/api/wkd/WkdClient.kt

@@ -7,27 +7,30 @@
 package com.flowcrypt.email.api.wkd
 
 import com.flowcrypt.email.extensions.kotlin.isValidEmail
+import com.flowcrypt.email.extensions.kotlin.isValidLocalhostEmail
 import com.flowcrypt.email.util.BetterInternetAddress
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import org.apache.commons.codec.binary.ZBase32
 import org.apache.commons.codec.digest.DigestUtils
 import org.bouncycastle.openpgp.PGPPublicKeyRingCollection
 import org.bouncycastle.openpgp.jcajce.JcaPGPPublicKeyRingCollection
+import java.io.InterruptedIOException
 import java.net.URLEncoder
 import java.net.UnknownHostException
 import java.util.Locale
 import java.util.concurrent.TimeUnit
 
 object WkdClient {
-  private const val DEFAULT_REQUEST_TIMEOUT = 4
+  const val DEFAULT_REQUEST_TIMEOUT = 4L
 
   fun lookupEmail(
     email: String,
-    timeout: Int = DEFAULT_REQUEST_TIMEOUT,
-    wkdPort: Int? = null
+    wkdPort: Int? = null,
+    useHttps: Boolean = true,
+    timeout: Long = DEFAULT_REQUEST_TIMEOUT
   ): PGPPublicKeyRingCollection? {
-    val keys = rawLookupEmail(email, timeout, wkdPort) ?: return null
+    val keys = rawLookupEmail(email, wkdPort, useHttps, timeout) ?: return null
     val lowerCaseEmail = email.toLowerCase(Locale.ROOT)
     val matchingKeys = keys.keyRings.asSequence().filter {
       for (userId in it.publicKey.userIDs) {
@@ -46,19 +49,23 @@ object WkdClient {
   @Suppress("private")
   fun rawLookupEmail(
     email: String,
-    timeout: Int = DEFAULT_REQUEST_TIMEOUT,
-    wkdPort: Int? = null
+    wkdPort: Int? = null,
+    useHttps: Boolean = true,
+    timeout: Long = DEFAULT_REQUEST_TIMEOUT
   ): PGPPublicKeyRingCollection? {
-    if (!email.isValidEmail()) throw IllegalArgumentException("Invalid email address")
+    if (!email.isValidEmail() && !email.isValidLocalhostEmail()) {
+      throw IllegalArgumentException("Invalid email address")
+    }
     val parts = email.split('@')
     val user = parts[0].toLowerCase(Locale.ROOT)
     val hu = ZBase32().encodeAsString(DigestUtils.sha1(user.toByteArray()))
     val directDomain = parts[1].toLowerCase(Locale.ROOT)
     val advancedDomainPrefix = if (directDomain == "localhost") "" else "openpgpkey."
     val directHost = if (wkdPort == null) directDomain else "${directDomain}:${wkdPort}"
     val advancedHost = "$advancedDomainPrefix$directHost"
-    val advancedUrl = "https://${advancedHost}/.well-known/openpgpkey/${directDomain}"
-    val directUrl = "https://${directHost}/.well-known/openpgpkey"
+    val protocol = if (useHttps) "https" else "http"
+    val advancedUrl = "$protocol://${advancedHost}/.well-known/openpgpkey/${directDomain}"
+    val directUrl = "$protocol://${directHost}/.well-known/openpgpkey"
     val userPart = "hu/$hu?l=${URLEncoder.encode(user, "UTF-8")}"
     try {
       val result = urlLookup(advancedUrl, userPart, timeout)
@@ -71,6 +78,8 @@ object WkdClient {
       urlLookup(directUrl, userPart, timeout).keys
     } catch (ex: UnknownHostException) {
       null
+    } catch (ex: InterruptedIOException) {
+      if (ex.message == "timeout") null else throw ex
     }
   }
 
@@ -79,8 +88,8 @@ object WkdClient {
     val keys: PGPPublicKeyRingCollection? = null
   )
 
-  private fun urlLookup(methodUrlBase: String, userPart: String, timeout: Int): UrlLookupResult {
-    val httpClient = OkHttpClient.Builder().callTimeout(timeout.toLong(), TimeUnit.SECONDS).build()
+  private fun urlLookup(methodUrlBase: String, userPart: String, timeout: Long): UrlLookupResult {
+    val httpClient = OkHttpClient.Builder().callTimeout(timeout, TimeUnit.SECONDS).build()
     val policyRequest = Request.Builder().url("$methodUrlBase/policy").build()
     httpClient.newCall(policyRequest).execute().use { policyResponse ->
       if (policyResponse.code != 200) return UrlLookupResult()

FlowCrypt/src/main/java/com/flowcrypt/email/extensions/kotlin/StringExt.kt

@@ -162,3 +162,7 @@ fun String.stripTrailing(ch: Char): String {
 fun String.isValidEmail(): Boolean {
   return BetterInternetAddress.isValidEmail(this)
 }
+
+fun String.isValidLocalhostEmail(): Boolean {
+  return BetterInternetAddress.isValidLocalhostEmail(this)
+}

FlowCrypt/src/main/java/com/flowcrypt/email/util/BetterEmailAddress.kt

@@ -29,11 +29,16 @@ class BetterInternetAddress(str: String, verifySpecialCharacters: Boolean = true
     private val containsSpecialCharacterRegex = ".*[()<>\\[\\]:;@\\\\,.\"].*".toRegex()
     // double quotes at ends only
     private val doubleQuotedTextRegex = "\"[^\"]*\"".toRegex()
+    private val validLocalhostEmailRegex = Regex("([a-zA-z])([a-zA-z0-9])+@localhost")
 
     fun isValidEmail(email: String): Boolean {
       return validEmailRegex.matchEntire(email) != null
     }
 
+    fun isValidLocalhostEmail(email: String): Boolean {
+      return validLocalhostEmailRegex.matchEntire(email) != null
+    }
+
     fun areValidEmails(emails: Iterable<String>): Boolean {
       return emails.all { it.isValidEmail() }
     }

FlowCrypt/src/test/java/com/flowcrypt/email/api/email/WkdClientTest.kt

@@ -7,6 +7,10 @@
 package com.flowcrypt.email.api.email
 
 import com.flowcrypt.email.api.wkd.WkdClient
+import okhttp3.mockwebserver.Dispatcher
+import okhttp3.mockwebserver.MockResponse
+import okhttp3.mockwebserver.MockWebServer
+import okhttp3.mockwebserver.RecordedRequest
 import org.junit.Assert.assertTrue
 import org.junit.Test
 
@@ -35,4 +39,23 @@ class WkdClientTest {
     val keys = WkdClient.lookupEmail("doesnotexist@thisdomaindoesnotexist.test")
     assertTrue("Key found for non-existing email", keys == null)
   }
+
+  @Test
+  fun requestTimeoutTest() {
+    val mockWebServer = MockWebServer()
+    mockWebServer.dispatcher = object: Dispatcher() {
+      private val sleepTimeout = (WkdClient.DEFAULT_REQUEST_TIMEOUT + 2) * 1000
+      override fun dispatch(request: RecordedRequest): MockResponse {
+        Thread.sleep(sleepTimeout)
+        return MockResponse().setResponseCode(200)
+      }
+    }
+    mockWebServer.start()
+    val port = mockWebServer.port
+    mockWebServer.use {
+      val keys = WkdClient.lookupEmail(email = "user@localhost", wkdPort = port, useHttps = false)
+      assertTrue("Key found for non-existing email", keys == null)
+    }
+  }
 }
+