From cbaae3fb081260a2fa10cc5331d524ef7ff6abe5 Mon Sep 17 00:00:00 2001
From: Jordan Rose <jrose@signal.org>
Date: Mon, 20 May 2024 10:02:47 -0700
Subject: [PATCH] Tighten up checks for PublicKey.init(keyBytes:)

For historical reasons libsignal's PublicKey.init(_:) allows trailing
data, because the type byte specifies how much of the data is relevant
anyway. init(keyBytes:), however, should not be so lax, since it's
assuming the key data is a Curve25519 public point.
---
 SignalServiceKit/Curve25519/PublicKey.swift           |  4 ++++
 SignalServiceKit/tests/Curve25519/ECKeyPairTest.swift | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/SignalServiceKit/Curve25519/PublicKey.swift b/SignalServiceKit/Curve25519/PublicKey.swift
index 671a0525e6..fb24fa23e3 100644
--- a/SignalServiceKit/Curve25519/PublicKey.swift
+++ b/SignalServiceKit/Curve25519/PublicKey.swift
@@ -8,10 +8,14 @@ import LibSignalClient
 
 extension PublicKey {
     public convenience init(keyData: Data) throws {
+        if keyData.count != Constants.keyLengthDJB {
+            throw SignalError.invalidKey("invalid number of public key bytes (expected \(Constants.keyLengthDJB), was \(keyData.count))")
+        }
         try self.init([Constants.keyTypeDJB] + keyData)
     }
 
     public enum Constants {
         public static let keyTypeDJB: UInt8 = 0x05
+        public static let keyLengthDJB: Int = 32
     }
 }
diff --git a/SignalServiceKit/tests/Curve25519/ECKeyPairTest.swift b/SignalServiceKit/tests/Curve25519/ECKeyPairTest.swift
index 38b6a8fc0e..73a6346508 100644
--- a/SignalServiceKit/tests/Curve25519/ECKeyPairTest.swift
+++ b/SignalServiceKit/tests/Curve25519/ECKeyPairTest.swift
@@ -9,6 +9,16 @@ import XCTest
 @testable import SignalServiceKit
 
 final class ECKeyPairTest: XCTestCase {
+    func testInvalidPublicSerializationLibsignalExtensions() throws {
+        let keyPair = ECKeyPair.generateKeyPair()
+
+        let serializedPublic = keyPair.publicKey
+        let deserializedPublic = try PublicKey(keyData: serializedPublic)
+        XCTAssertEqual(deserializedPublic, keyPair.identityKeyPair.publicKey)
+        XCTAssertThrowsError(try PublicKey(keyData: serializedPublic.dropFirst()))
+        XCTAssertThrowsError(try PublicKey(keyData: serializedPublic + [0x41]))
+    }
+
     func testEncodeDecode() throws {
         let privateKey = try PrivateKey(Array(repeating: 0, count: 31) + [0x41])
         let keyPair = ECKeyPair(IdentityKeyPair(publicKey: privateKey.publicKey, privateKey: privateKey))