From 4751e77935c7475a199158e3ebeb105f76b08095 Mon Sep 17 00:00:00 2001
From: Robert Ekl <ekl.rob@gmail.com>
Date: Wed, 25 Mar 2026 11:21:26 -0500
Subject: [PATCH] fix(meshcore): handle disabled export private key

---
 .../MeshCore/Session/MeshCoreSession.swift    | 15 +++--
 ...shCoreSessionCommandCorrelationTests.swift | 58 +++++++++++++++++++
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/MeshCore/Sources/MeshCore/Session/MeshCoreSession.swift b/MeshCore/Sources/MeshCore/Session/MeshCoreSession.swift
index 7d142abd3..8593ecaf1 100644
--- a/MeshCore/Sources/MeshCore/Session/MeshCoreSession.swift
+++ b/MeshCore/Sources/MeshCore/Session/MeshCoreSession.swift
@@ -1331,12 +1331,19 @@ public actor MeshCoreSession: MeshCoreSessionProtocol {
     /// This is a sensitive operation that exposes the device's cryptographic identity.
     /// The exported key can be imported into another device to clone its identity.
     ///
-    /// - Returns: The 32-byte private key, or `nil` if export is disabled.
-    /// - Throws: ``MeshCoreError/timeout`` if the device doesn't respond.
+    /// - Returns: The 32-byte private key.
+    /// - Throws: ``MeshCoreError/featureDisabled`` if private key export is disabled on the device,
+    ///   or ``MeshCoreError/timeout`` if the device doesn't respond.
     public func exportPrivateKey() async throws -> Data {
-        try await sendAndWait(PacketBuilder.exportPrivateKey()) { event in
+        try await sendAndWaitWithError(
+            PacketBuilder.exportPrivateKey()
+        ) { event in
             if case .privateKey(let key) = event { return key }
-            if case .disabled = event { return nil }
+            return nil
+        } errorMatcher: { event in
+            if case .disabled = event {
+                return MeshCoreError.featureDisabled
+            }
             return nil
         }
     }
diff --git a/MeshCore/Tests/MeshCoreTests/Session/MeshCoreSessionCommandCorrelationTests.swift b/MeshCore/Tests/MeshCoreTests/Session/MeshCoreSessionCommandCorrelationTests.swift
index 61aebe77d..0a47eb9b2 100644
--- a/MeshCore/Tests/MeshCoreTests/Session/MeshCoreSessionCommandCorrelationTests.swift
+++ b/MeshCore/Tests/MeshCoreTests/Session/MeshCoreSessionCommandCorrelationTests.swift
@@ -287,6 +287,64 @@ struct MeshCoreSessionCommandCorrelationTests {
         await session.stop()
     }
 
+    @Test("exportPrivateKey throws featureDisabled on disabled response")
+    func exportPrivateKeyThrowsFeatureDisabledOnDisabledResponse() async throws {
+        let transport = MockTransport()
+        let session = MeshCoreSession(
+            transport: transport,
+            configuration: SessionConfiguration(defaultTimeout: 0.2, clientIdentifier: "MCTst")
+        )
+
+        try await startSession(session, transport: transport)
+
+        let exportTask = Task {
+            try await session.exportPrivateKey()
+        }
+
+        try await waitUntil("exportPrivateKey should be sent") {
+            await transport.sentData.count == 2
+        }
+
+        await transport.simulateReceive(Data([ResponseCode.disabled.rawValue]))
+
+        let error = await #expect(throws: MeshCoreError.self) {
+            try await exportTask.value
+        }
+        guard case .featureDisabled? = error else {
+            Issue.record("Expected featureDisabled, got \(String(describing: error))")
+            await session.stop()
+            return
+        }
+
+        await session.stop()
+    }
+
+    @Test("disabled responses do not break unrelated requests")
+    func disabledResponsesDoNotBreakUnrelatedRequests() async throws {
+        let transport = MockTransport()
+        let session = MeshCoreSession(
+            transport: transport,
+            configuration: SessionConfiguration(defaultTimeout: 0.2, clientIdentifier: "MCTst")
+        )
+
+        try await startSession(session, transport: transport)
+
+        let batteryTask = Task {
+            try await session.getBattery()
+        }
+
+        try await waitUntil("getBattery should be sent") {
+            await transport.sentData.count == 2
+        }
+
+        await transport.simulateReceive(Data([ResponseCode.disabled.rawValue]))
+        await transport.simulateReceive(makeBatteryPacket(level: 4018))
+
+        let battery = try await batteryTask.value
+        #expect(battery.level == 4018)
+        await session.stop()
+    }
+
     @Test("requestStatus fails fast on device error before messageSent")
     func requestStatusFailsFastOnDeviceErrorBeforeMessageSent() async throws {
         let transport = MockTransport()