add length checks to secp256k1 functions

kukai-wallet · Dec 19, 2023 · b0747e6 · b0747e6
1 parent 21bb16c
commit b0747e6
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 3 deletions.
diff --git a/Sources/KukaiCryptoSwift/KeyPair.swift b/Sources/KukaiCryptoSwift/KeyPair.swift
@@ -125,6 +125,10 @@ public struct KeyPair {
 
 	/// Helper method to take a secp256k1 private key (for a regualr keypair) and use it to create a public key for the same curve
 	public static func secp256k1PublicKey(fromPrivateKeyBytes pkBytes: [UInt8]) -> PublicKey? {
+		if pkBytes.count != 64 {
+			return nil
+		}
+
 		var publicKey = secp256k1_pubkey()
 		var outputLength = 33
 		var publicKeyBytes = [UInt8](repeating: 0, count: outputLength)
@@ -148,6 +152,10 @@ public struct KeyPair {
 
 	/// Helper method to uncompress a secp256k1 public key
 	public static func secp256k1PublicKey_uncompressed(fromBytes: [UInt8]) -> [UInt8] {
+		if fromBytes.count != 32 {
+			return []
+		}
+
 		var publicKey = secp256k1_pubkey()
 		var outputLength = 65
 		var outputBytes = [UInt8](repeating: 0, count: outputLength)

diff --git a/Sources/KukaiCryptoSwift/PrivateKey.swift b/Sources/KukaiCryptoSwift/PrivateKey.swift
@@ -99,7 +99,7 @@ public struct PrivateKey: Codable {
 				let signatureLength = 64
 				var output = [UInt8](repeating: 0, count: signatureLength)
 
-				guard let context = secp256k1_context_create(UInt32(SECP256K1_CONTEXT_SIGN)) else {
+				guard let context = secp256k1_context_create(UInt32(SECP256K1_CONTEXT_SIGN)), self.bytes.count == 64 else {
 					return nil
 				}
 

diff --git a/Sources/KukaiCryptoSwift/PublicKey.swift b/Sources/KukaiCryptoSwift/PublicKey.swift
@@ -87,7 +87,7 @@ public struct PublicKey: Codable {
 				return Sodium.shared.sign.verify(message: message, publicKey: self.bytes, signature: signature)
 
 			case .secp256k1:
-				guard let context = secp256k1_context_create(UInt32(SECP256K1_CONTEXT_VERIFY)) else {
+				guard let context = secp256k1_context_create(UInt32(SECP256K1_CONTEXT_VERIFY)), signature.count == 64 else {
 					return false
 				}
 

diff --git a/Tests/KukaiCryptoSwiftTests/KeyPairTests.swift b/Tests/KukaiCryptoSwiftTests/KeyPairTests.swift
@@ -154,6 +154,6 @@ final class KeyPairTests: XCTestCase {
 
 
 		let pubKeySafety = KeyPair.secp256k1PublicKey(fromPrivateKeyBytes: signatureBytes)
-		XCTAssert(pubKeySafety?.bytes.count == 33, (pubKeySafety?.bytes.count ?? 0).description)
+		XCTAssert(pubKeySafety == nil, (pubKeySafety?.bytes.count ?? 0).description)
 	}
 }