Revamp PSL updates (#1475)

* build: import Mozilla's Gradle plugin for PSL updates

* autofill-parser: add tests for PublicSuffixListLoader

* autofill-parser: regenerate publicsuffixes list

* github: switch to Gradle plugin for PSL updates

7 files changed, 162 insertions, 4 deletions

diff --git a/.github/workflows/update_publicsuffix_data.yml b/.github/workflows/update_publicsuffix_data.yml
index 6ef3527d..dbe321cc 100644
--- a/.github/workflows/update_publicsuffix_data.yml
+++ b/.github/workflows/update_publicsuffix_data.yml
@@ -11,11 +11,19 @@ jobs:
         uses: actions/checkout@5a4ac9002d0be2fb38bd78e4b4dbde5606d7042f
 
       - name: Download new publicsuffix data
-        run: curl -L https://github.com/mozilla-mobile/android-components/raw/master/components/lib/publicsuffixlist/src/main/assets/publicsuffixes -o autofill-parser/src/main/assets/publicsuffixes
+        uses: burrunan/gradle-cache-action@03c71a8ba93d670980695505f48f49daf43704a6
+        with:
+          arguments: updatePSL
 
       - name: Compare list changes
         run: if [[ $(git diff --binary --stat) != '' ]]; then echo "UPDATED=true" >> $GITHUB_ENV; fi
 
+      - name: Verify update publicsuffixes file
+        uses: burrunan/gradle-cache-action@03c71a8ba93d670980695505f48f49daf43704a6
+        if: env.UPDATED == 'true'
+        with:
+          arguments: :autofill-parser:test
+
       - name: Create Pull Request
         uses: peter-evans/create-pull-request@01f7dd1d28f5131231ba3ede0f1c8cb413584a1d
         if: env.UPDATED == 'true'
diff --git a/autofill-parser/build.gradle.kts b/autofill-parser/build.gradle.kts
index f22081f3..be71acf0 100644
--- a/autofill-parser/build.gradle.kts
+++ b/autofill-parser/build.gradle.kts
@@ -8,9 +8,13 @@ plugins {
   id("com.vanniktech.maven.publish")
   kotlin("android")
   `aps-plugin`
+  `psl-plugin`
 }
 
-android { defaultConfig { consumerProguardFiles("consumer-rules.pro") } }
+android {
+  defaultConfig { consumerProguardFiles("consumer-rules.pro") }
+  sourceSets { getByName("test") { resources.srcDir("src/main/assets") } }
+}
 
 dependencies {
   implementation(libs.androidx.annotation)
@@ -18,4 +22,5 @@ dependencies {
   implementation(libs.kotlin.coroutines.android)
   implementation(libs.kotlin.coroutines.core)
   implementation(libs.thirdparty.timberkt)
+  testImplementation(libs.bundles.testDependencies)
 }
diff --git a/autofill-parser/src/main/assets/publicsuffixes b/autofill-parser/src/main/assets/publicsuffixes
index c12128bd..39d8a034 100644
--- a/autofill-parser/src/main/assets/publicsuffixes
+++ b/autofill-parser/src/main/assets/publicsuffixes
diff --git a/autofill-parser/src/main/java/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoader.kt b/autofill-parser/src/main/java/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoader.kt
index 5f3fc296..8491030d 100644
--- a/autofill-parser/src/main/java/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoader.kt
+++ b/autofill-parser/src/main/java/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoader.kt
@@ -16,8 +16,8 @@ private const val PUBLIC_SUFFIX_LIST_FILE = "publicsuffixes"
 
 internal object PublicSuffixListLoader {
 
-  fun load(context: Context): PublicSuffixListData =
-    context.assets.open(PUBLIC_SUFFIX_LIST_FILE).buffered().use { stream ->
+  fun load(inputStream: BufferedInputStream): PublicSuffixListData =
+    inputStream.use { stream ->
       val publicSuffixSize = stream.readInt()
       val publicSuffixBytes = stream.readFully(publicSuffixSize)
 
@@ -26,6 +26,9 @@ internal object PublicSuffixListLoader {
 
       PublicSuffixListData(publicSuffixBytes, exceptionBytes)
     }
+
+  fun load(context: Context): PublicSuffixListData =
+    load(context.assets.open(PUBLIC_SUFFIX_LIST_FILE).buffered())
 }
 
 @Suppress("MagicNumber")
diff --git a/autofill-parser/src/test/kotlin/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoaderTest.kt b/autofill-parser/src/test/kotlin/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoaderTest.kt
new file mode 100644
index 00000000..a4b7d4eb
--- /dev/null
+++ b/autofill-parser/src/test/kotlin/mozilla/components/lib/publicsuffixlist/PublicSuffixListLoaderTest.kt
@@ -0,0 +1,18 @@
+/*
+ * Copyright © 2014-2021 The Android Password Store Authors. All Rights Reserved.
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+package mozilla.components.lib.publicsuffixlist
+
+import org.junit.Test
+
+internal class PublicSuffixListLoaderTest {
+  @Test
+  fun testLoadingBundledPublicSuffixList() {
+    requireNotNull(javaClass.classLoader).getResourceAsStream("publicsuffixes").buffered().use {
+      stream ->
+      PublicSuffixListLoader.load(stream)
+    }
+  }
+}
diff --git a/buildSrc/build.gradle.kts b/buildSrc/build.gradle.kts
index 2239861b..e9d8587d 100644
--- a/buildSrc/build.gradle.kts
+++ b/buildSrc/build.gradle.kts
@@ -37,6 +37,10 @@ gradlePlugin {
       id = "versioning-plugin"
       implementationClass = "VersioningPlugin"
     }
+    register("psl") {
+      id = "psl-plugin"
+      implementationClass = "PublicSuffixListPlugin"
+    }
   }
 }
 
diff --git a/buildSrc/src/main/java/PublicSuffixListPlugin.kt b/buildSrc/src/main/java/PublicSuffixListPlugin.kt
new file mode 100644
index 00000000..df50f2b4
--- /dev/null
+++ b/buildSrc/src/main/java/PublicSuffixListPlugin.kt
@@ -0,0 +1,120 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+import java.io.File
+import java.util.TreeSet
+import okhttp3.OkHttpClient
+import okhttp3.Request
+import okio.ByteString
+import okio.ByteString.Companion.encodeUtf8
+import okio.buffer
+import okio.sink
+import org.gradle.api.Plugin
+import org.gradle.api.Project
+
+/**
+ * Gradle plugin to update the public suffix list used by the `lib-publicsuffixlist` component.
+ *
+ * Base on PublicSuffixListGenerator from OkHttp:
+ * https://github.com/square/okhttp/blob/master/okhttp/src/test/java/okhttp3/internal/publicsuffix/PublicSuffixListGenerator.java
+ */
+class PublicSuffixListPlugin : Plugin<Project> {
+  override fun apply(project: Project) {
+    project.tasks.register("updatePSL") {
+      doLast {
+        val filename = project.projectDir.absolutePath + "/src/main/assets/publicsuffixes"
+        updatePublicSuffixList(filename)
+      }
+    }
+  }
+
+  private fun updatePublicSuffixList(destination: String) {
+    val list = fetchPublicSuffixList()
+    writeListToDisk(destination, list)
+  }
+
+  private fun writeListToDisk(destination: String, data: PublicSuffixListData) {
+    val fileSink = File(destination).sink()
+
+    fileSink.buffer().use { sink ->
+      sink.writeInt(data.totalRuleBytes)
+
+      for (domain in data.sortedRules) {
+        sink.write(domain).writeByte('\n'.toInt())
+      }
+
+      sink.writeInt(data.totalExceptionRuleBytes)
+
+      for (domain in data.sortedExceptionRules) {
+        sink.write(domain).writeByte('\n'.toInt())
+      }
+    }
+  }
+
+  private fun fetchPublicSuffixList(): PublicSuffixListData {
+    val client = OkHttpClient.Builder().build()
+
+    val request =
+      Request.Builder().url("https://publicsuffix.org/list/public_suffix_list.dat").build()
+
+    client.newCall(request).execute().use { response ->
+      val source = requireNotNull(response.body).source()
+
+      val data = PublicSuffixListData()
+
+      while (!source.exhausted()) {
+        val line = source.readUtf8LineStrict()
+
+        if (line.trim { it <= ' ' }.isEmpty() || line.startsWith("//")) {
+          continue
+        }
+
+        if (line.contains(WILDCARD_CHAR)) {
+          assertWildcardRule(line)
+        }
+
+        var rule = line.encodeUtf8()
+
+        if (rule.startsWith(EXCEPTION_RULE_MARKER)) {
+          rule = rule.substring(1)
+          // We use '\n' for end of value.
+          data.totalExceptionRuleBytes += rule.size + 1
+          data.sortedExceptionRules.add(rule)
+        } else {
+          data.totalRuleBytes += rule.size + 1 // We use '\n' for end of value.
+          data.sortedRules.add(rule)
+        }
+      }
+
+      return data
+    }
+  }
+
+  @Suppress("TooGenericExceptionThrown", "ThrowsCount")
+  private fun assertWildcardRule(rule: String) {
+    if (rule.indexOf(WILDCARD_CHAR) != 0) {
+      throw RuntimeException("Wildcard is not not in leftmost position")
+    }
+
+    if (rule.indexOf(WILDCARD_CHAR, 1) != -1) {
+      throw RuntimeException("Rule contains multiple wildcards")
+    }
+
+    if (rule.length == 1) {
+      throw RuntimeException("Rule wildcards the first level")
+    }
+  }
+
+  companion object {
+    private const val WILDCARD_CHAR = "*"
+    private val EXCEPTION_RULE_MARKER = "!".encodeUtf8()
+  }
+}
+
+data class PublicSuffixListData(
+  var totalRuleBytes: Int = 0,
+  var totalExceptionRuleBytes: Int = 0,
+  val sortedRules: TreeSet<ByteString> = TreeSet(),
+  val sortedExceptionRules: TreeSet<ByteString> = TreeSet()
+)