swift: PreKeyRecords lazily validate their keys

Convert the *Key and keyPair properties on *PreKeyRecord to throwing methods. Since the usual constructor is strongly typed, these should only throw when the serialized data is corrupted. But that is a possibility.
howjmay · Aug 20, 2024 · 2f7f83d · 2f7f83d
1 parent 5e81e01
commit 2f7f83d
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 42 deletions.
diff --git a/swift/Sources/LibSignalClient/state/KyberPreKeyRecord.swift b/swift/Sources/LibSignalClient/state/KyberPreKeyRecord.swift
@@ -69,32 +69,26 @@ public class KyberPreKeyRecord: ClonableHandleOwner {
         }
     }
 
-    public var keyPair: KEMKeyPair {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_kyber_pre_key_record_get_key_pair($0, nativeHandle)
-                }
+    public func keyPair() throws -> KEMKeyPair {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_kyber_pre_key_record_get_key_pair($0, nativeHandle)
             }
         }
     }
 
-    public var publicKey: KEMPublicKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_kyber_pre_key_record_get_public_key($0, nativeHandle)
-                }
+    public func publicKey() throws -> KEMPublicKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_kyber_pre_key_record_get_public_key($0, nativeHandle)
             }
         }
     }
 
-    public var secretKey: KEMSecretKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_kyber_pre_key_record_get_secret_key($0, nativeHandle)
-                }
+    public func secretKey() throws -> KEMSecretKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_kyber_pre_key_record_get_secret_key($0, nativeHandle)
             }
         }
     }

diff --git a/swift/Sources/LibSignalClient/state/PreKeyRecord.swift b/swift/Sources/LibSignalClient/state/PreKeyRecord.swift
@@ -60,22 +60,18 @@ public class PreKeyRecord: ClonableHandleOwner {
         }
     }
 
-    public var publicKey: PublicKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_pre_key_record_get_public_key($0, nativeHandle)
-                }
+    public func publicKey() throws -> PublicKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_pre_key_record_get_public_key($0, nativeHandle)
             }
         }
     }
 
-    public var privateKey: PrivateKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_pre_key_record_get_private_key($0, nativeHandle)
-                }
+    public func privateKey() throws -> PrivateKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_pre_key_record_get_private_key($0, nativeHandle)
             }
         }
     }

diff --git a/swift/Sources/LibSignalClient/state/SignedPreKeyRecord.swift b/swift/Sources/LibSignalClient/state/SignedPreKeyRecord.swift
@@ -77,22 +77,18 @@ public class SignedPreKeyRecord: ClonableHandleOwner {
         }
     }
 
-    public var publicKey: PublicKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_signed_pre_key_record_get_public_key($0, nativeHandle)
-                }
+    public func publicKey() throws -> PublicKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_signed_pre_key_record_get_public_key($0, nativeHandle)
             }
         }
     }
 
-    public var privateKey: PrivateKey {
-        return withNativeHandle { nativeHandle in
-            failOnError {
-                try invokeFnReturningNativeHandle {
-                    signal_signed_pre_key_record_get_private_key($0, nativeHandle)
-                }
+    public func privateKey() throws -> PrivateKey {
+        return try withNativeHandle { nativeHandle in
+            try invokeFnReturningNativeHandle {
+                signal_signed_pre_key_record_get_private_key($0, nativeHandle)
             }
         }
     }

diff --git a/swift/Tests/LibSignalClientTests/SessionRecordTests.swift b/swift/Tests/LibSignalClientTests/SessionRecordTests.swift
@@ -0,0 +1,51 @@
+//
+// Copyright 2024 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+import LibSignalClient
+import XCTest
+
+class SessionRecordTests: TestCaseBase {
+    func testBadPreKeyRecords() {
+        XCTAssertThrowsError(try PreKeyRecord(bytes: [0]))
+        XCTAssertThrowsError(try SignedPreKeyRecord(bytes: [0]))
+        XCTAssertThrowsError(try KyberPreKeyRecord(bytes: [0]))
+
+        // The keys in records are lazily parsed, which means malformed keys aren't caught right away.
+        // The following payloads were generated via protoscope:
+        // % protoscope -s | base64
+        // The fields are described in storage.proto in the libsignal-protocol crate.
+        do {
+            // 1: 42
+            // 2: {}
+            // 3: {}
+            let record = try! PreKeyRecord(bytes: Data(base64Encoded: "CCoSABoA")!)
+            XCTAssertThrowsError(try record.publicKey())
+            XCTAssertThrowsError(try record.privateKey())
+        }
+
+        do {
+            // 1: 42
+            // 2: {}
+            // 3: {}
+            // 4: {}
+            // 5: 0i64
+            let record = try! SignedPreKeyRecord(bytes: Data(base64Encoded: "CCoSABoAIgApAAAAAAAAAAA=")!)
+            XCTAssertThrowsError(try record.publicKey())
+            XCTAssertThrowsError(try record.privateKey())
+        }
+
+        do {
+            // 1: 42
+            // 2: {}
+            // 3: {}
+            // 4: {}
+            // 5: 0i64
+            let record = try! KyberPreKeyRecord(bytes: Data(base64Encoded: "CCoSABoAIgApAAAAAAAAAAA=")!)
+            XCTAssertThrowsError(try record.publicKey())
+            XCTAssertThrowsError(try record.secretKey())
+            XCTAssertThrowsError(try record.keyPair())
+        }
+    }
+}