From 2fa03e3fa0bab743e65a8e5964ed510758b4a908 Mon Sep 17 00:00:00 2001
From: Fabian Henneke <FabianHenneke@users.noreply.github.com>
Date: Tue, 9 Jun 2020 13:45:23 +0200
Subject: [PATCH] Allow custom public suffixes for Autofill (#841)

Adds a preference that allows the user to specify domains that are then
treated as additional public suffixes for the purposes of Autofill.
---
 .../java/com/zeapo/pwdstore/UserPreference.kt | 10 ++++-
 .../autofill/oreo/PublicSuffixListCache.kt    | 41 ++++++++++++++++++-
 app/src/main/res/values/strings.xml           |  3 ++
 app/src/main/res/xml/preference.xml           |  4 ++
 4 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/app/src/main/java/com/zeapo/pwdstore/UserPreference.kt b/app/src/main/java/com/zeapo/pwdstore/UserPreference.kt
index 41404716..63c9ce7c 100644
--- a/app/src/main/java/com/zeapo/pwdstore/UserPreference.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/UserPreference.kt
@@ -116,6 +116,7 @@ class UserPreference : AppCompatActivity() {
             autoFillEnablePreference = findPreference("autofill_enable")
             val oreoAutofillDirectoryStructurePreference = findPreference<ListPreference>("oreo_autofill_directory_structure")
             val oreoAutofillDefaultUsername = findPreference<EditTextPreference>("oreo_autofill_default_username")
+            val oreoAutofillCustomPublixSuffixes = findPreference<EditTextPreference>("oreo_autofill_custom_public_suffixes")
             val autoFillAppsPreference = findPreference<Preference>("autofill_apps")
             val autoFillDefaultPreference = findPreference<CheckBoxPreference>("autofill_default")
             val autoFillAlwaysShowDialogPreference = findPreference<CheckBoxPreference>("autofill_always")
@@ -128,8 +129,15 @@ class UserPreference : AppCompatActivity() {
             )
             oreoAutofillDependencies = listOfNotNull(
                 oreoAutofillDirectoryStructurePreference,
-                oreoAutofillDefaultUsername
+                oreoAutofillDefaultUsername,
+                oreoAutofillCustomPublixSuffixes
             )
+            oreoAutofillCustomPublixSuffixes?.apply {
+                setOnBindEditTextListener {
+                    it.isSingleLine = false
+                    it.setHint(R.string.preference_custom_public_suffixes_hint)
+                }
+            }
 
             // Misc preferences
             val appVersionPreference = findPreference<Preference>("app_version")
diff --git a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/PublicSuffixListCache.kt b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/PublicSuffixListCache.kt
index 12d9a8c4..f1ce6bce 100644
--- a/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/PublicSuffixListCache.kt
+++ b/app/src/main/java/com/zeapo/pwdstore/autofill/oreo/PublicSuffixListCache.kt
@@ -6,6 +6,7 @@ package com.zeapo.pwdstore.autofill.oreo
 
 import android.content.Context
 import android.util.Patterns
+import androidx.preference.PreferenceManager
 import kotlinx.coroutines.runBlocking
 import mozilla.components.lib.publicsuffixlist.PublicSuffixList
 
@@ -44,7 +45,43 @@ fun getPublicSuffixPlusOne(context: Context, domain: String) = runBlocking {
     ) {
         domain
     } else {
-        PublicSuffixListCache.getOrCachePublicSuffixList(context).getPublicSuffixPlusOne(domain)
-            .await() ?: domain
+        getCanonicalSuffix(context, domain)
     }
 }
+
+/**
+ * Returns:
+ * - [domain], if [domain] equals [suffix];
+ * - null, if [domain] does not have [suffix] as a domain suffix or only with an empty prefix;
+ * - the direct subdomain of [suffix] of which [domain] is a subdomain.
+ */
+fun getSuffixPlusUpToOne(domain: String, suffix: String): String? {
+    if (domain == suffix)
+        return domain
+    val prefix = domain.removeSuffix(".$suffix")
+    if (prefix == domain || prefix.isEmpty())
+        return null
+    val lastPrefixPart = prefix.takeLastWhile { it != '.' }
+    return "$lastPrefixPart.$suffix"
+}
+
+fun getCustomSuffixes(context: Context): Sequence<String> {
+    val prefs = PreferenceManager.getDefaultSharedPreferences(context)
+    return prefs.getString("oreo_autofill_custom_public_suffixes", "")!!
+        .splitToSequence('\n')
+        .filter { it.isNotBlank() && it.first() != '.' && it.last() != '.' }
+}
+
+suspend fun getCanonicalSuffix(context: Context, domain: String): String {
+    val publicSuffixList = PublicSuffixListCache.getOrCachePublicSuffixList(context)
+    val publicSuffixPlusOne = publicSuffixList.getPublicSuffixPlusOne(domain).await()
+        ?: return domain
+    var longestSuffix = publicSuffixPlusOne
+    for (customSuffix in getCustomSuffixes(context)) {
+        val suffixPlusUpToOne = getSuffixPlusUpToOne(domain, customSuffix) ?: continue
+        // A shorter suffix is automatically a substring.
+        if (suffixPlusUpToOne.length > longestSuffix.length)
+            longestSuffix = suffixPlusUpToOne
+    }
+    return longestSuffix
+}
diff --git a/app/src/main/res/values/strings.xml b/app/src/main/res/values/strings.xml
index ffd10068..2e37b0eb 100644
--- a/app/src/main/res/values/strings.xml
+++ b/app/src/main/res/values/strings.xml
@@ -373,4 +373,7 @@
     <string name="preference_default_username_title">Default username</string>
     <string name="git_operation_remember_password">Remember password</string>
     <string name="git_operation_hint_password">Password</string>
+    <string name="preference_custom_public_suffixes_title">Custom domains</string>
+    <string name="preference_custom_public_suffixes_summary">Autofill will distinguish subdomains of these domains</string>
+    <string name="preference_custom_public_suffixes_hint">company.com\npersonal.com</string>
 </resources>
diff --git a/app/src/main/res/xml/preference.xml b/app/src/main/res/xml/preference.xml
index 9e56ce8f..0d71d6cc 100644
--- a/app/src/main/res/xml/preference.xml
+++ b/app/src/main/res/xml/preference.xml
@@ -21,6 +21,10 @@
             app:key="oreo_autofill_default_username"
             app:summary="@string/preference_default_username_summary"
             app:title="@string/preference_default_username_title" />
+        <EditTextPreference
+            app:key="oreo_autofill_custom_public_suffixes"
+            app:summary="@string/preference_custom_public_suffixes_summary"
+            app:title="@string/preference_custom_public_suffixes_title" />
         <Preference
             app:key="autofill_apps"
             app:title="@string/pref_autofill_apps_title" />