Make the packet tunnel enter error state when failing device check

mullvad · Sep 18, 2024 · 08456d4 · 08456d4
1 parent 2c71523
commit 08456d4
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 22 deletions.
diff --git a/ios/MullvadVPNTests/MullvadVPN/PacketTunnel/DeviceCheck/DeviceCheckOperationTests.swift b/ios/MullvadVPNTests/MullvadVPN/PacketTunnel/DeviceCheck/DeviceCheckOperationTests.swift
@@ -241,7 +241,7 @@ class DeviceCheckOperationTests: XCTestCase {
         waitForExpectations(timeout: .UnitTest.timeout)
     }
 
-    func testShouldReportFailedKeyRotataionAttempt() {
+    func testShouldReportFailedKeyRotationAttempt() {
         let expect = expectation(description: "Wait for operation to complete")
 
         let currentKey = PrivateKey()

diff --git a/ios/PacketTunnel/DeviceCheck/DeviceCheckOperation.swift b/ios/PacketTunnel/DeviceCheck/DeviceCheckOperation.swift
@@ -278,14 +278,14 @@ final class DeviceCheckOperation: ResultOperation<DeviceCheck> {
 }
 
 /// An error used internally by `DeviceCheckOperation`.
-private enum DeviceCheckError: LocalizedError, Equatable {
+public enum DeviceCheckError: LocalizedError, Equatable {
     /// Device is no longer logged in.
     case invalidDeviceState
 
     /// Main process has likely performed key rotation at the same time when packet tunnel was doing so.
     case keyRotationRace
 
-    var errorDescription: String? {
+    public var errorDescription: String? {
         switch self {
         case .invalidDeviceState:
             return "Cannot complete device check because device is no longer logged in."

diff --git a/ios/PacketTunnel/PacketTunnelProvider/DeviceChecker.swift b/ios/PacketTunnel/PacketTunnelProvider/DeviceChecker.swift
@@ -35,18 +35,18 @@ final class DeviceChecker {
         key.
      4. Rotate WireGuard key on key mismatch.
      */
-    func start(rotateKeyOnMismatch: Bool) async throws -> DeviceCheck {
+    func start(rotateKeyOnMismatch: Bool) async -> Result<DeviceCheck, Error> {
         let checkOperation = DeviceCheckOperation(
             dispatchQueue: dispatchQueue,
             remoteSevice: DeviceCheckRemoteService(accountsProxy: accountsProxy, devicesProxy: devicesProxy),
             deviceStateAccessor: DeviceStateAccessor(),
             rotateImmediatelyOnKeyMismatch: rotateKeyOnMismatch
         )
 
-        return try await withTaskCancellationHandler {
-            return try await withCheckedThrowingContinuation { continuation in
+        return await withTaskCancellationHandler {
+            return await withCheckedContinuation { continuation in
                 checkOperation.completionHandler = { result in
-                    continuation.resume(with: result)
+                    continuation.resume(with: .success(result))
                 }
                 operationQueue.addOperation(checkOperation)
             }

diff --git a/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
@@ -287,26 +287,40 @@ extension PacketTunnelProvider {
 extension PacketTunnelProvider {
     private func startDeviceCheck(rotateKeyOnMismatch: Bool = false) {
         Task {
-            do {
-                try await startDeviceCheckInner(rotateKeyOnMismatch: rotateKeyOnMismatch)
-            } catch {
-                providerLogger.error(error: error, message: "Failed to perform device check.")
-            }
+            await startDeviceCheckInner(rotateKeyOnMismatch: rotateKeyOnMismatch)
         }
     }
 
-    private func startDeviceCheckInner(rotateKeyOnMismatch: Bool) async throws {
-        let result = try await deviceChecker.start(rotateKeyOnMismatch: rotateKeyOnMismatch)
+    private func startDeviceCheckInner(rotateKeyOnMismatch: Bool) async {
+        let result = await deviceChecker.start(rotateKeyOnMismatch: rotateKeyOnMismatch)
 
-        if let blockedStateReason = result.blockedStateReason {
-            actor.setErrorState(reason: blockedStateReason)
-        }
+        switch result {
+        case let .failure(error):
+            switch error {
+            case is DeviceCheckError:
+                providerLogger.error("\(error.localizedDescription) Forcing a log out")
+                actor.setErrorState(reason: .deviceLoggedOut)
+            default:
+                providerLogger
+                    .error(
+                        "Entering blocked state because device check encountered a generic error: \(error.localizedDescription)"
+                    )
+                actor.setErrorState(reason: .unknown)
+            }
 
-        switch result.keyRotationStatus {
-        case let .attempted(date), let .succeeded(date):
-            actor.notifyKeyRotation(date: date)
-        case .noAction:
-            break
+        case let .success(keyRotationResult):
+            if let blockedStateReason = keyRotationResult.blockedStateReason {
+                providerLogger.error("Entering blocked state after unsuccessful device check: \(blockedStateReason)")
+                actor.setErrorState(reason: blockedStateReason)
+                return
+            }
+
+            switch keyRotationResult.keyRotationStatus {
+            case let .attempted(date), let .succeeded(date):
+                actor.notifyKeyRotation(date: date)
+            case .noAction:
+                break
+            }
         }
     }
 }