feat: ensure valid randomness of PrivateKey

21-DOT-DEV · Jul 13, 2023 · 633b449 · 633b449
1 parent 40f1f66
commit 633b449
Showing 1 changed file with 16 additions and 14 deletions.
diff --git a/Sources/zkp/secp256k1.swift b/Sources/zkp/secp256k1.swift
@@ -40,7 +40,8 @@ public extension secp256k1 {
         public static func create(_ context: Context = .none) throws -> OpaquePointer {
             var randomBytes = SecureBytes(count: secp256k1.ByteDetails.count).bytes
             guard let context = secp256k1_context_create(context.rawValue),
-                  secp256k1_context_randomize(context, &randomBytes).boolValue else {
+                  secp256k1_context_randomize(context, &randomBytes).boolValue
+            else {
                 throw secp256k1Error.underlyingCryptoError
             }
 
@@ -151,16 +152,14 @@ extension secp256k1 {
 
     /// Backing initialization that creates a random secp256k1 private key for signing
     @usableFromInline init(format: secp256k1.Format = .compressed) throws {
-        let privateKey = SecureBytes(count: secp256k1.ByteDetails.count)
-        self.keyParity = 0
-        self.format = format
-        self.privateBytes = privateKey
-        self.publicBytes = try PublicKeyImplementation.generate(bytes: &privateBytes, format: format)
-        self.xonlyBytes = try XonlyKeyImplementation.generate(
-            bytes: publicBytes,
-            keyParity: &keyParity,
-            format: format
-        )
+        for _ in 0 ..< 10 {
+            let randomBytes = SecureBytes(count: secp256k1.ByteDetails.count)
+            if let privateKey = try? PrivateKeyImplementation(dataRepresentation: Data(randomBytes), format: format) {
+                self = privateKey
+                return
+            }
+        }
+        fatalError("Looped more than 10 times trying to generate a key")
     }
 
     /// Backing initialization that creates a secp256k1 private key for signing from a data representation.
@@ -224,7 +223,8 @@ extension secp256k1 {
             var bytes = [UInt8](repeating: 0, count: keyLength)
 
             guard secp256k1_ec_pubkey_negate(context, &key).boolValue,
-                  secp256k1_ec_pubkey_serialize(context, &bytes, &keyLength, &key, format.rawValue).boolValue else {
+                  secp256k1_ec_pubkey_serialize(context, &bytes, &keyLength, &key, format.rawValue).boolValue
+            else {
                 throw secp256k1Error.underlyingCryptoError
             }
 
@@ -299,7 +299,8 @@ extension secp256k1 {
         signature.dataRepresentation.copyToUnsafeMutableBytes(of: &recoverySignature.data)
 
         guard secp256k1_ecdsa_recover(context, &pubKey, &recoverySignature, Array(digest)).boolValue,
-              secp256k1_ec_pubkey_serialize(context, &pubBytes, &pubKeyLen, &pubKey, format.rawValue).boolValue else {
+              secp256k1_ec_pubkey_serialize(context, &pubBytes, &pubKeyLen, &pubKey, format.rawValue).boolValue
+        else {
             throw secp256k1Error.underlyingCryptoError
         }
 
@@ -402,7 +403,8 @@ extension secp256k1 {
 
         guard secp256k1_ec_pubkey_parse(context, &pubKey, publicBytes, format.length).boolValue,
               secp256k1_xonly_pubkey_from_pubkey(context, &xonlyPubKey, &keyParity, &pubKey).boolValue,
-              secp256k1_xonly_pubkey_serialize(context, &xonlyBytes, &xonlyPubKey).boolValue else {
+              secp256k1_xonly_pubkey_serialize(context, &xonlyBytes, &xonlyPubKey).boolValue
+        else {
             throw secp256k1Error.underlyingCryptoError
         }