Merge branch 'add-e2e-test-for-shadowsocks-droid-1564'

10 files changed, 167 insertions, 83 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/ObfuscationModeCell.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/ObfuscationModeCell.kt
index 90b5da8891..37e582bb01 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/ObfuscationModeCell.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/ObfuscationModeCell.kt
@@ -21,6 +21,8 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.res.stringResource
+import androidx.compose.ui.semantics.selected
+import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.tooling.preview.PreviewParameter
 import net.mullvad.mullvadvpn.R
@@ -66,6 +68,7 @@ fun ObfuscationModeCell(
                 .fillMaxWidth()
                 .background(MaterialTheme.colorScheme.surfaceContainerLow)
                 .let { if (testTag != null) it.testTag(testTag) else it }
+                .semantics { selected = isSelected }
     ) {
         TwoRowCell(
             modifier = Modifier.weight(1f),
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/SelectableCell.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/SelectableCell.kt
index 9d3688d352..d390e92f6a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/SelectableCell.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/cell/SelectableCell.kt
@@ -11,6 +11,8 @@ import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.alpha
 import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.semantics.selected
+import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.text.TextStyle
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.Dp
@@ -58,7 +60,7 @@ fun SelectableCell(
     testTag: String = "",
 ) {
     BaseCell(
-        modifier = modifier,
+        modifier = modifier.semantics { selected = isSelected },
         onCellClicked = onCellClicked,
         isRowEnabled = isEnabled,
         headlineContent = {
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/extension/UiAutomatorExtensions.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/extension/UiAutomatorExtensions.kt
index ff17a9f497..cb186463cd 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/extension/UiAutomatorExtensions.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/extension/UiAutomatorExtensions.kt
@@ -3,11 +3,15 @@ package net.mullvad.mullvadvpn.test.common.extension
 import android.os.Build
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.BySelector
+import androidx.test.uiautomator.StaleObjectException
 import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.UiObject2
+import androidx.test.uiautomator.UiObject2Condition
 import androidx.test.uiautomator.Until
+import co.touchlab.kermit.Logger
 import java.util.regex.Pattern
 import net.mullvad.mullvadvpn.test.common.constant.DEFAULT_TIMEOUT
+import net.mullvad.mullvadvpn.test.common.constant.LONG_TIMEOUT
 
 fun UiDevice.findObjectByCaseInsensitiveText(text: String): UiObject2 {
     return findObjectWithTimeout(By.text(Pattern.compile(text, Pattern.CASE_INSENSITIVE)))
@@ -29,11 +33,45 @@ fun UiDevice.findObjectWithTimeout(
 
     val foundObject = findObject(selector)
 
-    require(foundObject != null) { "No matches for selector within timeout ($timeout): $selector" }
+    require(foundObject != null) {
+        "No matches for selector within timeout ($timeout ms): $selector"
+    }
 
     return foundObject
 }
 
+fun UiDevice.clickObjectAwaitCondition(
+    selector: BySelector,
+    condition: UiObject2Condition<Boolean>,
+    timeout: Long = LONG_TIMEOUT,
+) {
+    var foundObject = findObjectWithTimeout(selector, timeout)
+    foundObject.click()
+
+    val retryCount = 3
+    repeat(retryCount) {
+        try {
+            val wasChecked = foundObject.wait(condition, timeout)
+            require(wasChecked) {
+                "UiObject2 did not become match condition within timeout $timeout ms"
+            }
+            return
+        } catch (_: StaleObjectException) {
+            Logger.e("Caught StaleObjectException - retrying")
+            foundObject = findObjectWithTimeout(selector, timeout)
+        }
+    }
+    error("Exceeded maximum StaleObjectException count ($retryCount)")
+}
+
+fun UiDevice.clickObjectAwaitIsChecked(selector: BySelector, timeout: Long = LONG_TIMEOUT) {
+    clickObjectAwaitCondition(
+        selector = selector,
+        condition = Until.checked(true),
+        timeout = timeout,
+    )
+}
+
 fun UiDevice.clickAgreeOnPrivacyDisclaimer() {
     findObjectWithTimeout(By.text("Agree and continue")).click()
 }
@@ -63,11 +101,16 @@ fun UiDevice.clickAllowOnNotificationPermissionPromptIfApiLevel33AndAbove(
         findObjectWithTimeout(selector).click()
     } catch (e: IllegalArgumentException) {
         throw IllegalArgumentException(
-            "Failed to allow notification permission within timeout ($timeout)"
+            "Failed to allow notification permission within timeout ($timeout ms)"
         )
     }
 }
 
+fun UiDevice.pressBackTwice() {
+    pressBack()
+    pressBack()
+}
+
 fun UiObject2.findObjectWithTimeout(
     selector: BySelector,
     timeout: Long = DEFAULT_TIMEOUT,
@@ -79,7 +122,7 @@ fun UiObject2.findObjectWithTimeout(
         findObject(selector)
     } catch (e: NullPointerException) {
         throw IllegalArgumentException(
-            "No matches for selector within timeout ($timeout): $selector"
+            "No matches for selector within timeout ($timeout ms): $selector"
         )
     }
 }
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/interactor/AppInteractor.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/interactor/AppInteractor.kt
index c940f9bfba..a75747c76b 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/interactor/AppInteractor.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/interactor/AppInteractor.kt
@@ -16,10 +16,6 @@ import net.mullvad.mullvadvpn.test.common.extension.clickAgreeOnPrivacyDisclaime
 import net.mullvad.mullvadvpn.test.common.extension.clickAllowOnNotificationPermissionPromptIfApiLevel33AndAbove
 import net.mullvad.mullvadvpn.test.common.extension.dismissChangelogDialogIfShown
 import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
-import net.mullvad.mullvadvpn.test.common.page.ConnectPage
-import net.mullvad.mullvadvpn.test.common.page.SettingsPage
-import net.mullvad.mullvadvpn.test.common.page.VpnSettingsPage
-import net.mullvad.mullvadvpn.test.common.page.on
 
 class AppInteractor(
     private val device: UiDevice,
@@ -57,17 +53,6 @@ class AppInteractor(
         ensureLoggedIn()
     }
 
-    fun enableLocalNetworkSharing() {
-        on<ConnectPage> { clickSettings() }
-
-        on<SettingsPage> { clickVpnSettings() }
-
-        on<VpnSettingsPage> { clickLocalNetworkSharingSwitch() }
-
-        device.pressBack()
-        device.pressBack()
-    }
-
     fun attemptLogin(accountNumber: String) {
         val loginObject =
             device.findObjectWithTimeout(By.clazz("android.widget.EditText")).apply {
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Page.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Page.kt
index 3caac5bd95..9a30112557 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Page.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Page.kt
@@ -4,7 +4,7 @@ import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
 
 sealed class Page {
-    protected val uiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation())
+    val uiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation())
 
     abstract fun assertIsDisplayed()
 }
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Story.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Story.kt
new file mode 100644
index 0000000000..7f919ebaf2
--- /dev/null
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/Story.kt
@@ -0,0 +1,50 @@
+package net.mullvad.mullvadvpn.test.common.page
+
+import net.mullvad.mullvadvpn.test.common.extension.pressBackTwice
+
+// This file defines extension methods on Page objects that involve multiple actions
+// that navigate multiple pages.
+
+fun ConnectPage.disableObfuscationStory() {
+    clickSettings()
+    on<SettingsPage> { clickVpnSettings() }
+    on<VpnSettingsPage> {
+        scrollUntilWireGuardObfuscationOffCell()
+        clickWireGuardObfuscationOffCell()
+    }
+    uiDevice.pressBackTwice()
+}
+
+fun ConnectPage.disablePostQuantumStory() {
+    clickSettings()
+    on<SettingsPage> { clickVpnSettings() }
+    on<VpnSettingsPage> {
+        scrollUntilPostQuantumOffCell()
+        clickPostQuantumOffCell()
+    }
+    uiDevice.pressBackTwice()
+}
+
+fun ConnectPage.enableShadowsocksStory() {
+    clickSettings()
+    on<SettingsPage> { clickVpnSettings() }
+    on<VpnSettingsPage> {
+        scrollUntilWireGuardObfuscationShadowsocksCell()
+        clickWireGuardObfuscationShadowsocksCell()
+    }
+    uiDevice.pressBackTwice()
+}
+
+fun ConnectPage.enableDAITAStory() {
+    clickSettings()
+    on<SettingsPage> { clickDaita() }
+    on<DaitaSettingsPage> { clickEnableSwitch() }
+    uiDevice.pressBackTwice()
+}
+
+fun ConnectPage.enableLocalNetworkSharingStory() {
+    clickSettings()
+    on<SettingsPage> { clickVpnSettings() }
+    on<VpnSettingsPage> { clickLocalNetworkSharingSwitch() }
+    uiDevice.pressBackTwice()
+}
diff --git a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt
index 9301768100..6d5dbee492 100644
--- a/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt
+++ b/android/test/common/src/main/kotlin/net/mullvad/mullvadvpn/test/common/page/VpnSettingsPage.kt
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.test.common.page
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
 import androidx.test.uiautomator.Until
+import net.mullvad.mullvadvpn.test.common.extension.clickObjectAwaitIsChecked
 import net.mullvad.mullvadvpn.test.common.extension.findObjectWithTimeout
 
 class VpnSettingsPage internal constructor() : Page() {
@@ -39,23 +40,19 @@ class VpnSettingsPage internal constructor() : Page() {
     }
 
     fun clickWireguardObfuscationUdpOverTcpCell() {
-        uiDevice
-            .findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG))
-            .click()
+        uiDevice.clickObjectAwaitIsChecked(By.res(WIREGUARD_OBFUSCATION_UDP_OVER_TCP_CELL_TEST_TAG))
     }
 
     fun clickWireGuardObfuscationOffCell() {
-        uiDevice.findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG)).click()
+        uiDevice.clickObjectAwaitIsChecked(By.res(WIREGUARD_OBFUSCATION_OFF_CELL_TEST_TAG))
     }
 
     fun clickPostQuantumOffCell() {
-        uiDevice.findObjectWithTimeout(By.res(QUANTUM_RESISTANCE_OFF_CELL_TEST_TAG)).click()
+        uiDevice.clickObjectAwaitIsChecked(By.res(QUANTUM_RESISTANCE_OFF_CELL_TEST_TAG))
     }
 
     fun clickWireGuardObfuscationShadowsocksCell() {
-        uiDevice
-            .findObjectWithTimeout(By.res(WIREGUARD_OBFUSCATION_SHADOWSOCKS_CELL_TEST_TAG))
-            .click()
+        uiDevice.clickObjectAwaitIsChecked(By.res(WIREGUARD_OBFUSCATION_SHADOWSOCKS_CELL_TEST_TAG))
     }
 
     private fun scrollUntilCell(testTag: String) {
diff --git a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
index 0ed78f6db3..688b83c875 100644
--- a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
+++ b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/ConnectionTest.kt
@@ -10,6 +10,9 @@ import net.mullvad.mullvadvpn.test.common.page.SelectLocationPage
 import net.mullvad.mullvadvpn.test.common.page.SettingsPage
 import net.mullvad.mullvadvpn.test.common.page.SystemVpnConfigurationAlert
 import net.mullvad.mullvadvpn.test.common.page.VpnSettingsPage
+import net.mullvad.mullvadvpn.test.common.page.disableObfuscationStory
+import net.mullvad.mullvadvpn.test.common.page.enableLocalNetworkSharingStory
+import net.mullvad.mullvadvpn.test.common.page.enableShadowsocksStory
 import net.mullvad.mullvadvpn.test.common.page.on
 import net.mullvad.mullvadvpn.test.common.rule.ForgetAllVpnAppsInSettingsTestRule
 import net.mullvad.mullvadvpn.test.e2e.annotations.HasDependencyOnLocalAPI
@@ -71,7 +74,7 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
     @ClearFirewallRules
     fun testWireGuardObfuscationAutomatic() = runBlocking {
         app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
-        app.enableLocalNetworkSharing()
+        on<ConnectPage> { enableLocalNetworkSharingStory() }
 
         on<ConnectPage> { clickSelectLocation() }
 
@@ -108,7 +111,7 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
     @ClearFirewallRules
     fun testWireGuardObfuscationOff() = runBlocking {
         app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
-        app.enableLocalNetworkSharing()
+        on<ConnectPage> { enableLocalNetworkSharingStory() }
 
         on<ConnectPage> { clickSelectLocation() }
 
@@ -161,7 +164,7 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
     fun testUDPOverTCP() =
         runBlocking<Unit> {
             app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
-            app.enableLocalNetworkSharing()
+            on<ConnectPage> { enableLocalNetworkSharingStory() }
 
             on<ConnectPage> { clickSelectLocation() }
 
@@ -205,8 +208,44 @@ class ConnectionTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
             }
         }
 
+    @Test
+    @HasDependencyOnLocalAPI
+    @ClearFirewallRules
+    fun testShadowsocks() =
+        runBlocking<Unit> {
+            app.launchAndEnsureLoggedIn(accountTestRule.validAccountNumber)
+            on<ConnectPage> { enableLocalNetworkSharingStory() }
+
+            on<ConnectPage> { disableObfuscationStory() }
+
+            // Block all WireGuard traffic
+            val firewallRule = DropRule.blockWireGuardTrafficRule(ANY_IP_ADDRESS)
+            firewallClient.createRule(firewallRule)
+
+            on<ConnectPage> { clickConnect() }
+
+            on<SystemVpnConfigurationAlert> { clickOk() }
+
+            // Ensure it is not possible to connect to relay
+            on<ConnectPage> {
+                delay(UNSUCCESSFUL_CONNECTION_TIMEOUT.milliseconds)
+                waitForConnectingLabel()
+                clickCancel()
+            }
+
+            on<ConnectPage> { enableShadowsocksStory() }
+
+            // Ensure we can now connect with Shadowsocks enabled
+            on<ConnectPage> {
+                clickConnect()
+                waitForConnectedLabel(timeout = EXTREMELY_LONG_TIMEOUT)
+                clickDisconnect()
+            }
+        }
+
     companion object {
         const val VERY_FORGIVING_WIREGUARD_OFF_CONNECTION_TIMEOUT = 60000L
         const val UNSUCCESSFUL_CONNECTION_TIMEOUT = 60000L
+        const val ANY_IP_ADDRESS = "0.0.0.0/0"
     }
 }
diff --git a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/LeakTest.kt b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/LeakTest.kt
index f109ff3b74..79471a2363 100644
--- a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/LeakTest.kt
+++ b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/LeakTest.kt
@@ -5,12 +5,15 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import net.mullvad.mullvadvpn.test.common.misc.Attachment
 import net.mullvad.mullvadvpn.test.common.page.ConnectPage
-import net.mullvad.mullvadvpn.test.common.page.DaitaSettingsPage
 import net.mullvad.mullvadvpn.test.common.page.SelectLocationPage
 import net.mullvad.mullvadvpn.test.common.page.SettingsPage
 import net.mullvad.mullvadvpn.test.common.page.SystemVpnConfigurationAlert
 import net.mullvad.mullvadvpn.test.common.page.VpnSettingsPage
 import net.mullvad.mullvadvpn.test.common.page.WireGuardCustomPortDialog
+import net.mullvad.mullvadvpn.test.common.page.disableObfuscationStory
+import net.mullvad.mullvadvpn.test.common.page.disablePostQuantumStory
+import net.mullvad.mullvadvpn.test.common.page.enableDAITAStory
+import net.mullvad.mullvadvpn.test.common.page.enableShadowsocksStory
 import net.mullvad.mullvadvpn.test.common.page.on
 import net.mullvad.mullvadvpn.test.common.rule.ForgetAllVpnAppsInSettingsTestRule
 import net.mullvad.mullvadvpn.test.e2e.annotations.HasDependencyOnLocalAPI
@@ -181,9 +184,8 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
         runBlocking<Unit> {
             app.launch()
             // Obfuscation and Post-Quantum are by default set to automatic. Explicitly set to off.
-            disableObfuscation()
-            disablePostQuantum()
-
+            on<ConnectPage> { disableObfuscationStory() }
+            on<ConnectPage> { disablePostQuantumStory() }
             on<ConnectPage> { clickSelectLocation() }
 
             on<SelectLocationPage> {
@@ -207,9 +209,8 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
                         ) // Give it some time for generating traffic in tunnel before changing
                         // settings
 
-                        enableDAITA()
-                        enableShadowsocks()
-
+                        on<ConnectPage> { enableDAITAStory() }
+                        on<ConnectPage> { enableShadowsocksStory() }
                         on<ConnectPage> { waitForConnectedLabel() }
 
                         delay(
@@ -232,48 +233,4 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
                 NoTrafficToHostRule(targetIpAddress),
             )
         }
-
-    private fun disableObfuscation() {
-        on<ConnectPage> { clickSettings() }
-        on<SettingsPage> { clickVpnSettings() }
-        on<VpnSettingsPage> {
-            scrollUntilWireGuardObfuscationUdpOverTcpCell()
-            clickWireGuardObfuscationOffCell()
-        }
-
-        device.pressBack()
-        device.pressBack()
-    }
-
-    private fun disablePostQuantum() {
-        on<ConnectPage> { clickSettings() }
-        on<SettingsPage> { clickVpnSettings() }
-        on<VpnSettingsPage> {
-            scrollUntilPostQuantumOffCell()
-            clickPostQuantumOffCell()
-        }
-
-        device.pressBack()
-        device.pressBack()
-    }
-
-    private fun enableShadowsocks() {
-        on<ConnectPage> { clickSettings() }
-        on<SettingsPage> { clickVpnSettings() }
-        on<VpnSettingsPage> {
-            scrollUntilWireGuardObfuscationShadowsocksCell()
-            clickWireGuardObfuscationShadowsocksCell()
-        }
-
-        device.pressBack()
-        device.pressBack()
-    }
-
-    private fun enableDAITA() {
-        on<ConnectPage> { clickSettings() }
-        on<SettingsPage> { clickDaita() }
-        on<DaitaSettingsPage> { clickEnableSwitch() }
-        device.pressBack()
-        device.pressBack()
-    }
 }
diff --git a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/router/firewall/DropRule.kt b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/router/firewall/DropRule.kt
index 627ecdc0fe..1431d9bf1b 100644
--- a/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/router/firewall/DropRule.kt
+++ b/android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/router/firewall/DropRule.kt
@@ -14,13 +14,21 @@ import net.mullvad.mullvadvpn.test.e2e.router.NetworkingProtocol
 data class DropRule(
     @SerialName("src") val source: String,
     @SerialName("dst") val destination: String,
+    @SerialName("block_wireguard") val blockWireGuard: Boolean = false,
     val protocols: List<NetworkingProtocol>,
     @EncodeDefault val label: String = "urn:uuid:${SessionIdentifier.fromDeviceIdentifier()}",
 ) {
     companion object {
         fun blockUDPTrafficRule(to: String): DropRule {
             val testDeviceIpAddress = Networking.getDeviceIpv4Address()
-            return DropRule(testDeviceIpAddress, to, listOf(NetworkingProtocol.UDP))
+            return DropRule(
+                source = testDeviceIpAddress,
+                destination = to,
+                protocols = listOf(NetworkingProtocol.UDP),
+            )
         }
+
+        fun blockWireGuardTrafficRule(to: String): DropRule =
+            blockUDPTrafficRule(to).copy(blockWireGuard = true)
     }
 }