fix(firmware): Polish ESP32 OTA connection hardening

Author: jeremiah-k

core/data/src/commonMain/kotlin/org/meshtastic/core/data/manager/FromRadioPacketHandlerImpl.kt

@@ -22,6 +22,7 @@ import kotlinx.coroutines.SupervisorJob
 import org.koin.core.annotation.Single
 import org.meshtastic.core.common.util.handledLaunch
 import org.meshtastic.core.common.util.ioDispatcher
+import org.meshtastic.core.model.util.isOtaStatusNotification
 import org.meshtastic.core.repository.FromRadioPacketHandler
 import org.meshtastic.core.repository.LockdownCoordinator
 import org.meshtastic.core.repository.MeshConfigFlowManager
@@ -176,14 +177,3 @@ class FromRadioPacketHandlerImpl(
         }
     }
 }
-
-private const val OTA_KEYWORD = "OTA"
-private const val OTA_CONFIRM_PREFIX = "Rebooting to"
-private val OTA_REJECTION_PREFIXES = listOf("Cannot start OTA", "OTA Loader", "Unable to switch to the OTA partition")
-
-internal fun ClientNotification.isOtaStatusNotification(): Boolean {
-    val message = message.trim()
-    if (message.isBlank() || !message.contains(OTA_KEYWORD)) return false
-
-    return message.startsWith(OTA_CONFIRM_PREFIX) || OTA_REJECTION_PREFIXES.any { message.startsWith(it) }
-}

core/data/src/commonTest/kotlin/org/meshtastic/core/data/manager/FromRadioPacketHandlerImplTest.kt

@@ -21,6 +21,7 @@ import dev.mokkery.answering.returns
 import dev.mokkery.every
 import dev.mokkery.mock
 import dev.mokkery.verify
+import org.meshtastic.core.model.util.isOtaStatusNotification
 import org.meshtastic.core.repository.MeshConfigFlowManager
 import org.meshtastic.core.repository.MeshConfigHandler
 import org.meshtastic.core.repository.MqttManager

core/model/src/commonMain/kotlin/org/meshtastic/core/model/util/Extensions.kt

@@ -19,6 +19,7 @@
 package org.meshtastic.core.model.util
 
 import org.meshtastic.proto.Channel
+import org.meshtastic.proto.ClientNotification
 import org.meshtastic.proto.Config
 import org.meshtastic.proto.MeshPacket
 import org.meshtastic.proto.ModuleConfig
@@ -139,3 +140,12 @@ fun getInitials(fullName: String): String {
 }
 
 fun String.withoutEmojis(): String = filterNot { char -> char.isSurrogate() }
+
+private const val OTA_KEYWORD = "OTA"
+
+/**
+ * Broad OTA-status detector matching the firmware update preflight gate's criterion: any ClientNotification whose
+ * message contains "OTA" (e.g. "Rebooting to BLE OTA", "Cannot start OTA: ...") is consumed by the firmware update flow
+ * and suppressed as a generic alert. Shared so the UI and data layers do not duplicate the matcher.
+ */
+fun ClientNotification.isOtaStatusNotification(): Boolean = message.isNotBlank() && message.contains(OTA_KEYWORD)

core/ui/src/commonMain/kotlin/org/meshtastic/core/ui/viewmodel/UIViewModel.kt

@@ -47,6 +47,7 @@ import org.meshtastic.core.model.TracerouteMapAvailability
 import org.meshtastic.core.model.evaluateTracerouteMapAvailability
 import org.meshtastic.core.model.service.TracerouteResponse
 import org.meshtastic.core.model.util.dispatchMeshtasticUri
+import org.meshtastic.core.model.util.isOtaStatusNotification
 import org.meshtastic.core.navigation.DeepLinkRouter
 import org.meshtastic.core.repository.EventFirmwareRepository
 import org.meshtastic.core.repository.FirmwareReleaseRepository
@@ -258,8 +259,7 @@ class UIViewModel(
             .onEach { notification ->
                 // OTA status notifications (e.g. "Rebooting to WiFi OTA") are consumed by the firmware update
                 // preflight gate — don't show a popup dialog for them.
-                val msg = notification.message.trim()
-                if (msg.contains("OTA") && (msg.startsWith("Rebooting to") || msg.startsWith("Cannot start OTA"))) {
+                if (notification.isOtaStatusNotification()) {
                     return@onEach
                 }
                 val isCompromised = notification.low_entropy_key != null || notification.duplicated_public_key != null

core/ui/src/commonTest/kotlin/org/meshtastic/core/ui/viewmodel/UIViewModelImportSummaryTest.kt

@@ -17,6 +17,8 @@
 package org.meshtastic.core.ui.viewmodel
 
 import org.meshtastic.core.common.util.CommonUri
+import org.meshtastic.core.model.util.isOtaStatusNotification
+import org.meshtastic.proto.ClientNotification
 import kotlin.test.Test
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
@@ -46,4 +48,21 @@ class UIViewModelImportSummaryTest {
 
         assertTrue(summary.contains("hasFragment=false"))
     }
+
+    @Test
+    fun ota_status_notifications_are_suppressed() {
+        assertTrue(ClientNotification(message = "Rebooting to WiFi OTA").isOtaStatusNotification())
+        assertTrue(ClientNotification(message = "Rebooting to BLE OTA").isOtaStatusNotification())
+        assertTrue(
+            ClientNotification(message = "Cannot start OTA: OTA Loader partition not found.").isOtaStatusNotification(),
+        )
+        assertTrue(ClientNotification(message = "OTA Loader does not support WiFi").isOtaStatusNotification())
+        assertTrue(ClientNotification(message = "Unable to switch to the OTA partition.").isOtaStatusNotification())
+    }
+
+    @Test
+    fun non_ota_status_notifications_are_not_suppressed() {
+        assertFalse(ClientNotification(message = "Low battery").isOtaStatusNotification())
+        assertFalse(ClientNotification(message = "Key verification requested").isOtaStatusNotification())
+    }
 }

feature/firmware/src/commonMain/kotlin/org/meshtastic/feature/firmware/ota/BleOtaTransport.kt

@@ -43,6 +43,21 @@ import kotlin.time.Duration
 import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
 
+private const val MAC_OCTET_COUNT = 6
+private const val MAC_OUI_OCTETS = 3
+private const val MASKED_MAC_MIDDLE = "**:**"
+private const val INVALID_MAC_LABEL = "<invalid-mac>"
+
+/**
+ * Masks a BLE MAC address for logging/UI-facing exception text: preserves the OUI (first 3 octets) and last octet for
+ * diagnostics while hiding the unique middle portion. Privacy rule: never log raw PII.
+ */
+internal fun maskMac(address: String): String {
+    val octets = address.split(":")
+    if (octets.size != MAC_OCTET_COUNT || octets.any { it.length != 2 }) return INVALID_MAC_LABEL
+    return "${octets.take(MAC_OUI_OCTETS).joinToString(":")}:$MASKED_MAC_MIDDLE:${octets.last()}"
+}
+
 /** BLE transport implementation for ESP32 Unified OTA protocol using Kable. */
 class BleOtaTransport(
     private val scanner: BleScanner,
@@ -85,7 +100,7 @@ class BleOtaTransport(
             val device =
                 scanForOtaDevice()
                     ?: throw OtaProtocolException.ConnectionFailed(
-                        "Device not found at address $address. " +
+                        "Device not found at address ${maskMac(address)}. " +
                             "Ensure the device has rebooted into OTA mode and is advertising.",
                     )
 
@@ -101,12 +116,14 @@ class BleOtaTransport(
                 if (finalState is BleConnectionState.Disconnected) {
                     Logger.w { "BLE OTA: Failed to connect to ${maskMac(device.address)} (state=$finalState)" }
                     throw OtaProtocolException.ConnectionFailed(
-                        "Failed to connect to device at address ${device.address}",
+                        "Failed to connect to device at address ${maskMac(device.address)}",
                     )
                 }
             } catch (@Suppress("SwallowedException") e: kotlinx.coroutines.TimeoutCancellationException) {
                 Logger.w { "BLE OTA: Timed out waiting to connect to ${maskMac(device.address)}. Error: ${e.message}" }
-                throw OtaProtocolException.Timeout("Timed out connecting to device at address ${device.address}")
+                throw OtaProtocolException.Timeout(
+                    "Timed out connecting to device at address ${maskMac(device.address)}",
+                )
             }
 
             Logger.i { "BLE OTA: Connected to ${maskMac(device.address)}, discovering services..." }
@@ -295,13 +312,6 @@ class BleOtaTransport(
             OtaProtocolException.TransferFailed("Transfer failed: ${error.message}")
         }
 
-    /**
-     * Masks a BLE MAC address for logging: preserves the OUI (first 3 octets) and last octet for diagnostics while
-     * hiding the unique middle portion. Privacy rule: never log raw PII.
-     */
-    private fun maskMac(address: String): String =
-        address.takeIf { it.length >= 17 }?.let { "${it.substring(0, 8)}:**:**:**:${it.substring(15)}" } ?: address
-
     override suspend fun close() {
         try {
             bleConnection.disconnect()

feature/firmware/src/commonMain/kotlin/org/meshtastic/feature/firmware/ota/Esp32OtaUpdateHandler.kt

@@ -87,6 +87,11 @@ private const val OTA_KEYWORD = "OTA"
 // — keeps an IPv6 WiFi target (which also has colons) from being misrouted to the BLE path.
 private val MAC_ADDRESS_REGEX = Regex("([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}")
 
+private data class TransportFactoryResolution(
+    val factory: () -> UnifiedOtaProtocol,
+    val readinessAlreadyWaited: Boolean,
+)
+
 internal fun isBleMacAddress(target: String): Boolean = MAC_ADDRESS_REGEX.matches(target)
 
 /**
@@ -115,6 +120,7 @@ class Esp32OtaUpdateHandler(
     internal var delayFn: suspend (Long) -> Unit = { delay(it) }
     internal var bleTransportFactoryOverride: ((String) -> UnifiedOtaProtocol)? = null
     internal var wifiTransportFactoryOverride: ((String) -> UnifiedOtaProtocol)? = null
+    internal var wifiOtaDiscoveryOverride: (suspend () -> String?)? = null
 
     /** Entry point for FirmwareUpdateHandler interface. Routes to BLE (target is a MAC) or WiFi (anything else). */
     override suspend fun startUpdate(
@@ -320,49 +326,71 @@ class Esp32OtaUpdateHandler(
         }
     }
 
-    @Suppress("ThrowsCount") // CancellationException rethrow + two post-confirm failure paths
     private suspend fun connectToDevice(
         transportFactory: () -> UnifiedOtaProtocol,
         attempts: Int,
         rebootMode: Int,
         postConfirmReadinessDelayMs: Long,
         updateState: (FirmwareUpdateState) -> Unit,
     ): UnifiedOtaProtocol {
-        // Show "waiting for reboot" state before first connection attempt
         updateState(
             FirmwareUpdateState.Processing(ProgressState(UiText.Resource(Res.string.firmware_update_waiting_reboot))),
         )
 
-        if (postConfirmReadinessDelayMs > 0L) {
-            Logger.i {
-                "ESP32 OTA: Waiting ${postConfirmReadinessDelayMs}ms for ${otaModeName(rebootMode)} OTA service"
-            }
-            delayFn(postConfirmReadinessDelayMs)
-        }
+        val resolution = resolvePostConfirmTransportFactory(rebootMode, transportFactory, updateState)
+        waitForPostConfirmReadiness(rebootMode, postConfirmReadinessDelayMs, resolution.readinessAlreadyWaited)
+
+        return runTransportConnectRetries(resolution.factory, attempts, rebootMode, updateState)
+    }
 
+    private suspend fun resolvePostConfirmTransportFactory(
+        rebootMode: Int,
+        transportFactory: () -> UnifiedOtaProtocol,
+        updateState: (FirmwareUpdateState) -> Unit,
+    ): TransportFactoryResolution {
         // In production WiFi mode the device may have picked up a different DHCP lease after rebooting into the OTA
         // loader. Listen for the loader's UDP discovery broadcast and, if one arrives, redirect the TCP transport at
-        // the discovered IP. Skipped when a transport override is installed (tests inject fake transports).
-        val effectiveTransportFactory: () -> UnifiedOtaProtocol =
-            if (rebootMode == REBOOT_MODE_WIFI && wifiTransportFactoryOverride == null) {
-                updateState(
-                    FirmwareUpdateState.Processing(
-                        ProgressState(UiText.Resource(Res.string.firmware_update_searching_device)),
-                    ),
-                )
-                val discoveredIp = WifiOtaDiscovery.discoverOtaDevice()
-                if (discoveredIp != null) {
-                    Logger.i { "ESP32 OTA: Using UDP-discovered OTA device IP $discoveredIp for TCP transport" }
-                    val factory: () -> UnifiedOtaProtocol = { WifiOtaTransport(discoveredIp) }
-                    factory
-                } else {
-                    Logger.i { "ESP32 OTA: No UDP discovery broadcast received; falling back to configured device IP" }
-                    transportFactory
+        // the discovered IP. Skipped for transport-only test overrides unless a discovery override is installed.
+        if (
+            rebootMode != REBOOT_MODE_WIFI || (wifiTransportFactoryOverride != null && wifiOtaDiscoveryOverride == null)
+        ) {
+            return TransportFactoryResolution(factory = transportFactory, readinessAlreadyWaited = false)
+        }
+
+        updateState(
+            FirmwareUpdateState.Processing(ProgressState(UiText.Resource(Res.string.firmware_update_searching_device))),
+        )
+        val discoverOtaDevice = wifiOtaDiscoveryOverride ?: { WifiOtaDiscovery.discoverOtaDevice() }
+        val discoveredIp = discoverOtaDevice()
+        val factory: () -> UnifiedOtaProtocol =
+            if (discoveredIp != null) {
+                val discoveredFactory: () -> UnifiedOtaProtocol = {
+                    wifiTransportFactoryOverride?.invoke(discoveredIp) ?: WifiOtaTransport(discoveredIp)
                 }
+                Logger.i { "ESP32 OTA: Using UDP-discovered OTA device for TCP transport" }
+                discoveredFactory
             } else {
+                Logger.i { "ESP32 OTA: No UDP discovery broadcast received; falling back to configured device IP" }
                 transportFactory
             }
+        // Only treat the readiness window as already-spent when discovery actually resolved a device — its 15 s
+        // listening window covers the loader's reboot+DHCP+TCP-server bring-up. On a null result (timeout / bind
+        // failure) the device has not yet been heard from, so the caller's 8 s readiness margin still applies.
+        return TransportFactoryResolution(factory = factory, readinessAlreadyWaited = discoveredIp != null)
+    }
+
+    private suspend fun waitForPostConfirmReadiness(rebootMode: Int, delayMs: Long, readinessAlreadyWaited: Boolean) {
+        if (delayMs <= 0L || readinessAlreadyWaited) return
+        Logger.i { "ESP32 OTA: Waiting ${delayMs}ms for ${otaModeName(rebootMode)} OTA service" }
+        delayFn(delayMs)
+    }
 
+    private suspend fun runTransportConnectRetries(
+        transportFactory: () -> UnifiedOtaProtocol,
+        attempts: Int,
+        rebootMode: Int,
+        updateState: (FirmwareUpdateState) -> Unit,
+    ): UnifiedOtaProtocol {
         var connectedTransport: UnifiedOtaProtocol? = null
         val result =
             retryWithDelay(
@@ -376,40 +404,38 @@ class Esp32OtaUpdateHandler(
                     )
                 },
             ) {
-                val transport = effectiveTransportFactory()
-                val connectResult =
-                    try {
-                        transport.connect()
-                    } catch (e: CancellationException) {
-                        closeFailedTransport(transport)
-                        throw e
-                    } catch (@Suppress("TooGenericExceptionCaught") e: Throwable) {
-                        Result.failure(e)
-                    }
-
-                if (connectResult.isSuccess) {
-                    connectedTransport = transport
-                    connectResult
-                } else {
-                    val error = connectResult.exceptionOrNull()
-                    if (error is CancellationException) {
-                        closeFailedTransport(transport)
-                        throw error
-                    }
-                    Logger.w(error) {
-                        "ESP32 OTA: ${otaModeName(
-                            rebootMode,
-                        )} connection attempt failed; closing transport before retry"
-                    }
-                    closeFailedTransport(transport)
-                    Result.failure(error ?: OtaProtocolException.ConnectionFailed("Connection attempt failed"))
-                }
+                val transport = transportFactory()
+                val connectResult = connectTransportAttempt(transport, rebootMode)
+                connectResult.onSuccess { connectedTransport = transport }
             }
 
         result.getOrElse { cause -> throw postConfirmConnectionFailed(rebootMode, attempts, cause) }
         return connectedTransport ?: throw postConfirmConnectionFailed(rebootMode, attempts, null)
     }
 
+    private suspend fun connectTransportAttempt(transport: UnifiedOtaProtocol, rebootMode: Int): Result<Unit> {
+        val connectResult =
+            try {
+                transport.connect()
+            } catch (e: CancellationException) {
+                closeFailedTransport(transport)
+                throw e
+            } catch (@Suppress("TooGenericExceptionCaught") e: Throwable) {
+                Result.failure(e)
+            }
+
+        val error = connectResult.exceptionOrNull() ?: return connectResult
+        if (error is CancellationException) {
+            closeFailedTransport(transport)
+            throw error
+        }
+        Logger.w(error) {
+            "ESP32 OTA: ${otaModeName(rebootMode)} connection attempt failed; closing transport before retry"
+        }
+        closeFailedTransport(transport)
+        return Result.failure(error)
+    }
+
     private suspend fun closeFailedTransport(transport: UnifiedOtaProtocol) {
         withContext(NonCancellable) {
             try {