2 files changed, 48 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
+}