2 files changed, 51 insertions, 20 deletions
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/CustomDnsAdapter.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/CustomDnsAdapter.kt
index a8dd24528a..d7d9a5813e 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/CustomDnsAdapter.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/CustomDnsAdapter.kt
@@ -124,35 +124,66 @@ class CustomDnsAdapter(val customDns: CustomDns) : Adapter<CustomDnsItemHolder>(
         }
     }
 
-    fun addDnsServer(address: String) {
-        jobTracker.newUiJob("addDnsServer $address") {
-            var added = false
+    fun saveDnsServer(address: String) {
+        jobTracker.newUiJob("saveDnsServer $address") {
+            editingPosition?.let { position ->
+                if (position >= cachedCustomDnsServers.size) {
+                    addDnsServer(address)
+                } else {
+                    replaceDnsServer(address, position)
+                }
+            }
+        }
+    }
+
+    fun removeDnsServer(address: InetAddress) {
+        jobTracker.newBackgroundJob("removeDnsServer $address") {
+            customDns.removeDnsServer(address)
+        }
+    }
+
+    private suspend fun addDnsServer(address: String) {
+        var added = false
 
-            jobTracker.runOnBackground {
-                if (inetAddressValidator.isValid(address)) {
-                    val address = InetAddress.getByName(address)
+        jobTracker.runOnBackground {
+            if (inetAddressValidator.isValid(address)) {
+                val address = InetAddress.getByName(address)
 
-                    if (customDns.addDnsServer(address)) {
-                        cachedCustomDnsServers.add(address)
-                        added = true
-                    }
+                if (customDns.addDnsServer(address)) {
+                    cachedCustomDnsServers.add(address)
+                    added = true
                 }
             }
+        }
 
-            if (added) {
-                editingPosition = null
+        if (added) {
+            editingPosition = null
 
-                val count = getItemCount()
+            val count = getItemCount()
 
-                notifyItemChanged(count - 3)
-                notifyItemInserted(count - 2)
-            }
+            notifyItemChanged(count - 3)
+            notifyItemInserted(count - 2)
         }
     }
 
-    fun removeDnsServer(address: InetAddress) {
-        jobTracker.newBackgroundJob("removeDnsServer $address") {
-            customDns.removeDnsServer(address)
+    private suspend fun replaceDnsServer(address: String, position: Int) {
+        var replaced = false
+
+        jobTracker.runOnBackground {
+            if (inetAddressValidator.isValid(address)) {
+                val newAddress = InetAddress.getByName(address)
+                val oldAddress = cachedCustomDnsServers[position]
+
+                if (customDns.replaceDnsServer(oldAddress, newAddress)) {
+                    cachedCustomDnsServers[position] = newAddress
+                    replaced = true
+                }
+            }
+        }
+
+        if (replaced) {
+            editingPosition = null
+            notifyItemChanged(position)
         }
     }
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/EditCustomDnsServerHolder.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/EditCustomDnsServerHolder.kt
index c0ca781628..d2decf9241 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/EditCustomDnsServerHolder.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/customdns/EditCustomDnsServerHolder.kt
@@ -9,7 +9,7 @@ class EditCustomDnsServerHolder(view: View, adapter: CustomDnsAdapter) : CustomD
 
     init {
         view.findViewById<View>(R.id.save).setOnClickListener {
-            adapter.addDnsServer(input.text.toString())
+            adapter.saveDnsServer(input.text.toString())
         }
     }
 }