ET-2216 Push notification decryption

Modules/App/Sources/Services/AppConfigService.swift

@@ -50,17 +50,7 @@ final class AppConfigService: Sendable {
 
         return AppConfig(appVersion: appVersion, environment: environment)
 #else
-        let domain = Bundle.main.infoDictionary?["PMApiHost"] as? String ?? "proton.me"
-        let appVersion = "ios-mail@7.0.0" // Read from config once "ios-mail@x.y.z" is supported.
-        let environment = AppConfig.Environment(
-            domain: domain,
-            apiBaseUrl: "https://mail-api.\(domain)",
-            userAgent: "Mozilla/5.0",
-            isSrpProofSkipped: false,
-            isHttpAllowed: false
-        )
-
-        return AppConfig(appVersion: appVersion, environment: environment)
+        return .default
 #endif
     }()
 }

Modules/InboxCore/Sources/AppConfig.swift

@@ -46,6 +46,20 @@ public struct AppConfig: Sendable {
 
 public extension AppConfig {
 
+    static let `default`: Self = {
+        let domain = Bundle.main.infoDictionary?["PMApiHost"] as? String ?? "proton.me"
+        let appVersion = "ios-mail@7.0.0" // Read from config once "ios-mail@x.y.z" is supported.
+        let environment = AppConfig.Environment(
+            domain: domain,
+            apiBaseUrl: "https://mail-api.\(domain)",
+            userAgent: "Mozilla/5.0",
+            isSrpProofSkipped: false,
+            isHttpAllowed: false
+        )
+
+        return .init(appVersion: appVersion, environment: environment)
+    }()
+
     var apiEnvConfig: ApiConfig {
         let environment = self.environment
 

Modules/NotificationService/Sources/NotificationService.swift

@@ -16,6 +16,7 @@
 // along with Proton Mail. If not, see https://www.gnu.org/licenses/.
 
 import InboxCore
+import TestableNotificationService
 import UserNotifications
 
 class NotificationService: UNNotificationServiceExtension {
@@ -23,15 +24,13 @@ class NotificationService: UNNotificationServiceExtension {
         _ request: UNNotificationRequest,
         withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void
     ) {
-        AppLogger.log(message: "Notification received: \(request.content.title)", category: .notifications)
+        Task {
+            let updatedContent = await TestableNotificationService().transform(originalContent: request.content)
 
-        guard let mutableContent = (request.content.mutableCopy() as? UNMutableNotificationContent) else {
-            contentHandler(request.content)
-            return
+            await MainActor.run {
+                contentHandler(updatedContent)
+            }
         }
-
-        mutableContent.body = "overridden by app extension"
-        contentHandler(mutableContent)
     }
 
     override func serviceExtensionTimeWillExpire() {

Modules/TestableNotificationService/Package.swift

@@ -0,0 +1,21 @@
+// swift-tools-version: 5.9
+// The swift-tools-version declares the minimum version of Swift required to build this package.
+
+import PackageDescription
+
+let package = Package(
+    name: "TestableNotificationService",
+    platforms: [.iOS(.v17)],
+    products: [
+        .library(name: "TestableNotificationService", targets: ["TestableNotificationService"]),
+    ],
+    dependencies: [
+        .package(path: "../InboxCore"),
+        .package(path: "../InboxKeychain"),
+        .package(path: "../../ProtonPackages/proton_app_uniffi")
+    ],
+    targets: [
+        .target(name: "TestableNotificationService", dependencies: ["InboxCore", "InboxKeychain", "proton_app_uniffi"]),
+        .testTarget(name: "NotificationServiceTests", dependencies: ["TestableNotificationService"]),
+    ]
+)

Modules/TestableNotificationService/Sources/ProcessWideMailSessionCache.swift

@@ -0,0 +1,49 @@
+// Copyright (c) 2025 Proton Technologies AG
+//
+// This file is part of Proton Mail.
+//
+// Proton Mail is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Proton Mail 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 General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Proton Mail. If not, see https://www.gnu.org/licenses/.
+
+import InboxCore
+import InboxKeychain
+import proton_app_uniffi
+
+/*
+ This type exists because even though `NotificationService` is actually deinitialized every time after it finishes processing a notification,
+ the `MailSession` created on the Rust side remains in memory and needs to be reused.
+
+ Attempting to create a new one will trigger a panic - a SetGlobalDefaultError related to the logger in the SDK.
+ */
+enum ProcessWideMailSessionCache {
+    private static var cachedMailSession: MailSession?
+
+    static func prepareMailSession() throws -> MailSession {
+        if let cachedMailSession = cachedMailSession {
+            return cachedMailSession
+        } else {
+            let params = MailSessionParamsFactory.make(appConfig: .default)
+            let mailSessionResult = createMailSession(params: params, keyChain: KeychainSDKWrapper())
+
+            switch mailSessionResult {
+            case .ok(let mailSession):
+                cachedMailSession = mailSession
+                return mailSession
+            case .error(let userSessionError):
+                throw userSessionError
+            }
+        }
+    }
+}
+
+extension UserSessionError: Error {}

Modules/TestableNotificationService/Sources/TestableNotificationService.swift

@@ -0,0 +1,112 @@
+// Copyright (c) 2025 Proton Technologies AG
+//
+// This file is part of Proton Mail.
+//
+// Proton Mail is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Proton Mail 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 General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Proton Mail. If not, see https://www.gnu.org/licenses/.
+
+import InboxCore
+import proton_app_uniffi
+import UserNotifications
+
+public struct TestableNotificationService {
+    typealias DecryptRemoteNotification = (EncryptedPushNotification) async throws -> DecryptedPushNotification
+
+    private let decryptRemoteNotification: DecryptRemoteNotification
+
+    public init() {
+        self.init {
+            let mailSession = try ProcessWideMailSessionCache.prepareMailSession()
+
+            switch await decryptPushNotification(session: mailSession, encrypted: $0) {
+            case .ok(let value):
+                return value
+            case .error(let error):
+                throw error
+            }
+        }
+    }
+
+    init(decryptRemoteNotification: @escaping DecryptRemoteNotification) {
+        self.decryptRemoteNotification = decryptRemoteNotification
+    }
+
+    public func transform(originalContent: UNNotificationContent) async -> UNNotificationContent {
+        guard let mutableContent = (originalContent.mutableCopy() as? UNMutableNotificationContent) else {
+            AppLogger.log(message: "Notification content cannot be mutated", category: .notifications, isError: true)
+            return originalContent
+        }
+
+        // this is a temporary "marker" body to see if the extension has been launched by the OS, which is known to not be the case sometimes
+        mutableContent.body = "You received a new message!"
+
+        if let encryptedPushNotification = parseDecryptablePayload(from: originalContent.userInfo) {
+            await replaceTitleAndBody(of: mutableContent, byDecrypting: encryptedPushNotification)
+        }
+
+        return mutableContent
+    }
+
+    private func parseDecryptablePayload(from userInfo: [AnyHashable: Any]) -> EncryptedPushNotification? {
+        guard
+            let encryptedMessage = userInfo["encryptedMessage"] as? String,
+            let sessionId = userInfo["UID"] as? String
+        else {
+            AppLogger.log(message: "Missing required fields in the payload", category: .notifications, isError: true)
+            return nil
+        }
+
+        return .init(authId: sessionId, encryptedMessage: encryptedMessage)
+    }
+
+    private func replaceTitleAndBody(
+        of mutableContent: UNMutableNotificationContent,
+        byDecrypting encryptedPushNotification: EncryptedPushNotification
+    ) async {
+        do {
+            let notificationData = try await decryptRemoteNotification(encryptedPushNotification)
+            mutableContent.title = notificationData.sender.displayableName
+            mutableContent.body = notificationData.body
+        } catch {
+            AppLogger.log(error: error, category: .notifications)
+        }
+    }
+}
+
+private extension DecryptedPushNotification {
+    var body: String {
+        switch self {
+        case .email(let payload):
+            payload.subject
+        case .openUrl(let payload):
+            payload.content
+        }
+    }
+
+    var sender: NotificationSender {
+        switch self {
+        case .email(let payload):
+            payload.sender
+        case .openUrl(let payload):
+            payload.sender
+        }
+    }
+}
+
+private extension NotificationSender {
+    var displayableName: String {
+        name.isEmpty ? address : name
+    }
+}
+
+extension ActionError: Error {}

Modules/TestableNotificationService/Tests/NotificationServiceTests.swift

@@ -0,0 +1,102 @@
+// Copyright (c) 2025 Proton Technologies AG
+//
+// This file is part of Proton Mail.
+//
+// Proton Mail is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Proton Mail 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 General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Proton Mail. If not, see https://www.gnu.org/licenses/.
+
+import proton_app_uniffi
+import Testing
+import UserNotifications
+
+@testable import TestableNotificationService
+
+final class NotificationServiceTests {
+    private lazy var sut = TestableNotificationService { [unowned self] _ in
+        try self.stubbedDecryptionResult.get()
+    }
+
+    private var stubbedDecryptionResult: Result<DecryptedPushNotification, ActionError>!
+
+    @Test
+    func whenReplacingTitleAndBodyWithDecryptedInfo_prefersSenderName() async {
+        let originalContent = prepareContent()
+
+        stubbedDecryptionResult = .success(
+            .email(
+                .init(
+                    subject: "Decrypted subject",
+                    sender: .init(name: "John Doe", address: "john.doe@example.com", group: ""),
+                    messageId: .init(value: "")
+                )
+            )
+        )
+
+        let updatedContent = await sut.transform(originalContent: originalContent)
+
+        #expect(updatedContent.title == "John Doe")
+        #expect(updatedContent.body == "Decrypted subject")
+    }
+
+    @Test
+    func whenReplacingTitleAndBodyWithDecryptedInfo_ifSenderNameIsEmpty_fallsBackToSenderAddress() async {
+        let originalContent = prepareContent()
+
+        stubbedDecryptionResult = .success(
+            .email(
+                .init(
+                    subject: "Decrypted subject",
+                    sender: .init(name: "", address: "john.doe@example.com", group: ""),
+                    messageId: .init(value: "")
+                )
+            )
+        )
+
+        let updatedContent = await sut.transform(originalContent: originalContent)
+
+        #expect(updatedContent.title == "john.doe@example.com")
+        #expect(updatedContent.body == "Decrypted subject")
+    }
+
+    @Test
+    func whenInitialParsingFails_onlyModifiesTheBody() async {
+        let originalContent = prepareContent(userInfo: [:])
+
+        let updatedContent = await sut.transform(originalContent: originalContent)
+
+        #expect(updatedContent.title == "original title")
+        #expect(updatedContent.body == "You received a new message!")
+    }
+
+    @Test
+    func whenDecryptionFails_onlyModifiesTheBody() async {
+        let originalContent = prepareContent()
+        let stubbedError = ActionError.other(.unexpected(.crypto))
+        stubbedDecryptionResult = .failure(stubbedError)
+
+        let updatedContent = await sut.transform(originalContent: originalContent)
+
+        #expect(updatedContent.title == "original title")
+        #expect(updatedContent.body == "You received a new message!")
+    }
+
+    private func prepareContent(
+        userInfo: [AnyHashable: Any] = ["encryptedMessage": "foo", "UID": "123"]
+    ) -> UNNotificationContent {
+        let content = UNMutableNotificationContent()
+        content.title = "original title"
+        content.body = "original body"
+        content.userInfo = userInfo
+        return content
+    }
+}

TestPlans/AllUnitAndSnapshotTests.xctestplan

@@ -16,9 +16,9 @@
   "testTargets" : [
     {
       "target" : {
-        "containerPath" : "container:Modules\/InboxKeychain",
-        "identifier" : "InboxKeychainTests",
-        "name" : "InboxKeychainTests"
+        "containerPath" : "container:Modules\/InboxComposer",
+        "identifier" : "InboxComposerTests",
+        "name" : "InboxComposerTests"
       }
     },
     {
@@ -30,9 +30,9 @@
     },
     {
       "target" : {
-        "containerPath" : "container:Modules\/InboxComposer",
-        "identifier" : "InboxComposerTests",
-        "name" : "InboxComposerTests"
+        "containerPath" : "container:Modules\/InboxKeychain",
+        "identifier" : "InboxKeychainTests",
+        "name" : "InboxKeychainTests"
       }
     },
     {
@@ -48,6 +48,13 @@
         "identifier" : "InboxCoreUITests",
         "name" : "InboxCoreUITests"
       }
+    },
+    {
+      "target" : {
+        "containerPath" : "container:Modules\/TestableNotificationService",
+        "identifier" : "NotificationServiceTests",
+        "name" : "NotificationServiceTests"
+      }
     }
   ],
   "version" : 1

project.yml

@@ -59,6 +59,9 @@ packages:
   SwiftUIIntrospect:
     url: https://github.com/siteline/swiftui-introspect
     from: "1.3.0"
+  TestableNotificationService:
+    path: Modules/TestableNotificationService
+    group: Modules
   ViewInspector:
     url: https://github.com/nalexn/ViewInspector.git
     from: "0.10.1"
@@ -196,7 +199,7 @@ targets:
     type: app-extension
     platform: iOS
     dependencies:
-      - package: InboxCore
+      - package: TestableNotificationService
     sources:
       - path: Modules/NotificationService/Sources
     settings:
@@ -254,7 +257,6 @@ schemes:
       targets:
         NotificationService: all
         ProtonMail: all
-      buildImplicitDependencies: true
     run:
       config: Debug
     test: