Fix NFC signature result in UI

Modules/Test/CommonsTestShared/Sources/CommonsTestShared/Concurrency/AwaitableCondition.swift

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2017 - 2026 Riigi Infosüsteemi Amet
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ *
+ */
+
+import Foundation
+
+public actor AwaitableCondition {
+    private var waiter: CheckedContinuation<Void, Never>?
+    private var isFulfilled = false
+
+    public init() {}
+
+    public func fulfill() {
+        isFulfilled = true
+        let waiter = self.waiter
+        self.waiter = nil
+        waiter?.resume()
+    }
+
+    public func wait() async {
+        if isFulfilled { return }
+        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
+            waiter = continuation
+        }
+    }
+}

RIADigiDoc/Domain/NFC/NFCOperationBase.swift

@@ -35,6 +35,8 @@ public class NFCOperationBase: NSObject, Loggable, @MainActor NFCTagReaderSessio
 
     let connection = NFCConnection()
 
+    public var onStepChange: (@MainActor (Int) -> Void)?
+
     func updateAlertMessage(step: Int) {
         let stepMessages = [
             strings?.initialMessage ?? "",
@@ -50,6 +52,7 @@ public class NFCOperationBase: NSObject, Loggable, @MainActor NFCTagReaderSessio
         Self.logger().info("NFC: Updating alert message to: \(message)")
         message += "\n\n\(progressBar.generate())"
         session?.alertMessage = message
+        onStepChange?(step)
     }
 
     func success() {

RIADigiDoc/Domain/NFC/OperationReadCertAndSign.swift

@@ -36,6 +36,9 @@ public class OperationReadCertAndSign: NFCOperationBase, OperationReadCertAndSig
     private var userAgent: String = ""
     private var returnData: SignedContainerProtocol?
 
+    private var isOperationRunning = false
+    private var pendingCancellation: Error?
+
     private var continuation: CheckedContinuation<SignedContainerProtocol, Error>?
 
     // swiftlint:disable:next function_parameter_count
@@ -57,11 +60,18 @@ public class OperationReadCertAndSign: NFCOperationBase, OperationReadCertAndSig
         self.userAgent = userAgent
         self.strings = strings
 
+        returnData = nil
+        operationError = nil
+        didCompleteSuccessfully = false
+        nfcError = ""
+        isOperationRunning = false
+        pendingCancellation = nil
+
         return try await withCheckedThrowingContinuation { continuation in
             self.continuation = continuation
 
             guard NFCTagReaderSession.readingAvailable else {
-                continuation.resume(throwing: IdCardInternalError.nfcNotSupported)
+                resume(with: .failure(IdCardInternalError.nfcNotSupported))
                 return
             }
 
@@ -71,13 +81,40 @@ public class OperationReadCertAndSign: NFCOperationBase, OperationReadCertAndSig
         }
     }
 
+    private func finishOperation() {
+        if let returnData {
+            resume(with: .success(returnData))
+            return
+        }
+
+        resume(with: .failure(pendingCancellation ?? operationError ?? IdCardInternalError.sessionInvalidated))
+    }
+
+    private static func userCancellation(from error: Error) -> Error? {
+        guard let nfcError = error as? NFCReaderError,
+              nfcError.code == .readerSessionInvalidationErrorUserCanceled else {
+            return nil
+        }
+        return IdCardInternalError.cancelledByUser
+    }
+
+    private func resume(with result: Result<SignedContainerProtocol, Error>) {
+        guard let continuation else { return }
+        self.continuation = nil
+        continuation.resume(with: result)
+    }
+
     // MARK: - NFCTagReaderSessionDelegate
 
     // swiftlint:disable:next cyclomatic_complexity
     public override func tagReaderSession(_ session: NFCTagReaderSession, didDetect tags: [NFCTag]) {
         Task { @MainActor in
+            isOperationRunning = true
+
             defer {
                 self.session = nil
+                isOperationRunning = false
+                finishOperation()
             }
 
             guard let signedContainer else {
@@ -193,30 +230,22 @@ public class OperationReadCertAndSign: NFCOperationBase, OperationReadCertAndSig
         Self.logger().info("NFC: Reader session finished with error: \(error)")
         self.session = nil
 
-        guard let continuationToResume = self.continuation else { return }
-        self.continuation = nil
-
-        if let returnData, didCompleteSuccessfully {
-            continuationToResume.resume(with: .success(returnData))
+        guard !isOperationRunning else {
+            pendingCancellation = Self.userCancellation(from: error)
             return
         }
 
         if let storedError = self.operationError {
-            continuationToResume.resume(throwing: storedError)
+            resume(with: .failure(storedError))
             return
         }
 
-        if let nfcError = error as? NFCReaderError {
-            switch nfcError.code {
-            case .readerSessionInvalidationErrorUserCanceled:
-                continuationToResume.resume(throwing: IdCardInternalError.cancelledByUser)
-                return
-
-            default:
-                break
-            }
+        if let nfcError = error as? NFCReaderError,
+           nfcError.code == .readerSessionInvalidationErrorUserCanceled {
+            resume(with: .failure(IdCardInternalError.cancelledByUser))
+            return
         }
 
-        continuationToResume.resume(throwing: error)
+        resume(with: .failure(error))
     }
 }

RIADigiDoc/Domain/NFC/Protocols/OperationDecryptProtocol.swift

@@ -24,7 +24,9 @@ import nfclib
 
 /// @mockable
 @MainActor
-public protocol OperationDecryptProtocol {
+public protocol OperationDecryptProtocol: AnyObject {
+    var onStepChange: (@MainActor (Int) -> Void)? { get set }
+
     func processDecrypt(
         canNumber: String,
         pin1Number: SecureData,

RIADigiDoc/Domain/NFC/Protocols/OperationReadCertAndSignProtocol.swift

@@ -24,7 +24,9 @@ import LibdigidocLibSwift
 
 /// @mockable
 @MainActor
-public protocol OperationReadCertAndSignProtocol {
+public protocol OperationReadCertAndSignProtocol: AnyObject {
+    var onStepChange: (@MainActor (Int) -> Void)? { get set }
+
     // swiftlint:disable:next function_parameter_count
     func startOperation(
         canNumber: String,

RIADigiDoc/UI/Component/Container/Signing/NFC/NFCView.swift

@@ -137,11 +137,11 @@ struct NFCView: View {
     ) {
         _viewModel = State(wrappedValue: Container.shared.nfcViewModel())
         _webEidViewModel = State(wrappedValue: webEidViewModel)
-        self.actionType = actionType
-        self.pinType = pinType
-        self._isWebEidAuthenticating = isWebEidAuthenticating
-        self.rememberMe = rememberMe
-        self.actionMethods = actionMethods
+        _actionType = State(wrappedValue: actionType)
+        _actionMethods = State(wrappedValue: actionMethods)
+        _pinType = State(wrappedValue: pinType)
+        _rememberMe = State(wrappedValue: rememberMe)
+        _isWebEidAuthenticating = isWebEidAuthenticating
         self.cryptoContainer = cryptoContainer
         self.signedContainer = signedContainer
         self.onSuccess = onSuccess
@@ -376,6 +376,9 @@ struct NFCView: View {
 
             viewModel.resetErrors()
         }
+        .onChange(of: viewModel.actionMessageKey) { _, newMessageKey in
+            nfcActionMessage = newMessageKey
+        }
         .onChange(of: viewModel.certMismatch) { _, mismatch in
             if mismatch {
                 canNumber = ""

RIADigiDoc/UI/Component/Container/Signing/SigningView.swift

@@ -414,13 +414,14 @@ struct SigningView: View {
                             }
                         }
 
-                        containerLoadingTask = Task {
-                            await viewModel.loadContainerData(
-                                signedContainer: viewModel.signedContainer
-                            )
-
-                            await updateSignAndEncryptButtonVisibility()
+                        loadContainer(signedContainer: viewModel.signedContainer)
+                    }
+                    .onChange(of: viewModel.currentContainerID) { _, newContainerID in
+                        guard newContainerID != nil else { return }
+                        if viewModel.isSignatureAdded() {
+                            selectedTab = .signatures
                         }
+                        loadContainer(signedContainer: nil)
                     }
                     .onDisappear {
                         containerLoadingTask?.cancel()
@@ -565,6 +566,15 @@ struct SigningView: View {
         }
     }
 
+    private func loadContainer(signedContainer: SignedContainerProtocol?) {
+        let previousLoad = containerLoadingTask
+        containerLoadingTask = Task {
+            _ = await previousLoad?.value
+            await viewModel.loadContainerData(signedContainer: signedContainer)
+            await updateSignAndEncryptButtonVisibility()
+        }
+    }
+
     private func updateSignAndEncryptButtonVisibility() async {
         let shouldShowSignButton = await viewModel
             .isSignButtonShown(

RIADigiDoc/ViewModel/Signing/NFC/NFCViewModel.swift

@@ -46,6 +46,8 @@ class NFCViewModel: NFCViewModelProtocol, Loggable {
     var nfcAlertMessageUrl: String?
 
     var signatureExtensionFailed = false
+
+    private(set) var actionMessageKey: String = "NFC hold card"
     var certMismatch: Bool = false
 
     private let nfcCANKeyFilename = Constants.File.nfcCANKey
@@ -328,6 +330,11 @@ class NFCViewModel: NFCViewModelProtocol, Loggable {
         let recipients = await cryptoContainer?.getRecipients() ?? []
         let pinSecureData = SecureData(Array(pin1.utf8))
         await clearTempCAN()
+        actionMessageKey = "NFC hold card"
+        operationDecrypt.onStepChange = { [weak self] step in
+            self?.actionMessageKey = step >= 4 ? "Decrypting in progress" : "NFC hold card"
+        }
+
         do {
             NFCViewModel.logger().info("NFC: Starting decryption operation")
             let container = try await operationDecrypt.processDecrypt(
@@ -395,6 +402,11 @@ class NFCViewModel: NFCViewModelProtocol, Loggable {
         let appInfo = userAgentUtil.appInfo(diagnostics: .nfc, language: appLanguage)
         await clearTempCAN()
 
+        actionMessageKey = "NFC hold card"
+        operationReadCertAndSign.onStepChange = { [weak self] step in
+            self?.actionMessageKey = step >= 4 ? "Signing in progress" : "NFC hold card"
+        }
+
         do {
             NFCViewModel.logger().info("NFC: Starting signing operation")
             let result = try await operationReadCertAndSign.startOperation(

RIADigiDoc/ViewModel/SigningViewModel.swift

@@ -91,6 +91,10 @@ class SigningViewModel: SigningViewModelProtocol, Loggable {
         self.containerUtil = containerUtil
     }
 
+    var currentContainerID: ObjectIdentifier? {
+        sharedContainerViewModel.currentContainer().map { ObjectIdentifier($0) }
+    }
+
     func loadContainerData(signedContainer: SignedContainerProtocol?) async {
         SigningViewModel.logger().info("Loading signed container data")
         sharedContainerViewModel.setIsSignatureAdded(false)

RIADigiDocTests/ViewModel/EncryptRecipientViewModelTests.swift

@@ -23,7 +23,6 @@ import CryptoSwift
 import Foundation
 import Testing
 
-
 @MainActor
 struct EncryptRecipientViewModelTests {
 

RIADigiDocTests/ViewModel/Signing/MobileId/MobileIdViewModelTests.swift

@@ -26,6 +26,7 @@ import MobileIdLibMocks
 import CommonsLib
 import LibdigidocLibSwift
 import LibdigidocLibSwiftMocks
+import CommonsTestShared
 import UtilsLibMocks
 
 @MainActor
@@ -935,3 +936,53 @@ struct MobileIdViewModelTests {
         )
     }
 }
+
+extension MobileIdViewModelTests {
+
+    @Test
+    func sign_writesSignatureAfterTaskCancelledDuringAddSignature() async {
+        mockMobileIdSignService.getCertificateRequestHandler = { _, _, _, _, _, _, _, _ in
+            await mockMobileIdCertificateResponse()
+        }
+        mockMobileIdSignService.getVerificationCodeHandler = { _ in "1234" }
+        mockMobileIdSignService.getSignatureRequestHandler = { _, _, _, _, _, _, _, _, _, _, _, _, _ in
+            await mockSuccessSignature()
+        }
+        mockMobileIdSignService.getSessionRequestHandler = { _, _, _, _, _, _ in
+            await mockSuccessSession()
+        }
+        mockProxyUtil.getProxyInfoHandler = { ProxyInfo() }
+
+        let writeStarted = AwaitableCondition()
+        let writeMayFinish = AwaitableCondition()
+
+        let updatedContainer = SignedContainerProtocolMock()
+        let container = SignedContainerProtocolMock()
+        container.getRawContainerFileHandler = { URL(fileURLWithPath: "/tmp/test.asice") }
+        container.prepareSignatureHandler = { _, _, _, _ in Data([0x01]) }
+        container.addSignatureHandler = { _, _ in
+            await writeStarted.fulfill()
+            await writeMayFinish.wait()
+            return updatedContainer
+        }
+
+        let signingTask = Task { @MainActor in
+            await viewModel.sign(
+                phoneNumber: "37251234567",
+                personalCode: "60001019906",
+                roleData: roleData,
+                signedContainer: container
+            )
+        }
+
+        await writeStarted.wait()
+
+        signingTask.cancel()
+        await writeMayFinish.fulfill()
+
+        let result = await signingTask.value
+
+        #expect(result === updatedContainer, "the container written after cancellation is returned")
+        #expect(container.addSignatureCallCount == 1)
+    }
+}