From 633b449d0a89c29121cec1f00d01cff077475f06 Mon Sep 17 00:00:00 2001
From: charlie632 <charlie.caoa@gmail.com>
Date: Wed, 12 Jul 2023 22:30:36 -0600
Subject: [PATCH] feat: ensure valid randomness of PrivateKey

---
 Sources/zkp/secp256k1.swift | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/Sources/zkp/secp256k1.swift b/Sources/zkp/secp256k1.swift
index 5a8bb78..82d399d 100644
--- 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
         }