fix: issue planetarium#1696

Modify GenerateKeyParam to return a new value when .D.ToByteArrayUnsigned().Length != 32 Repeat until == 32
zeta6 · Jun 18, 2022 · b2507e9 · b2507e9
1 parent 6b2ccce
commit b2507e9
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 1 deletion.
diff --git a/Libplanet.Tests/Crypto/PrivateKeyTest.cs b/Libplanet.Tests/Crypto/PrivateKeyTest.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
 using System.Text;
@@ -325,5 +326,36 @@ public void HexStringConstructor()
                 actual.ToByteArray()
             );
         }
+
+        [Fact]
+        public void PrivateKeyShorterThan32Byte()
+        {
+            var bs = ByteUtil.ParseHex(
+                "5c352665057f1dfb90845d08c5712b8dd5a64dc4e2e4d089a89f30cbd469b5"
+            );
+            var privateKeyWith31Bytes = new PrivateKey(bs, true);
+            var newPrivateKey = new PrivateKey(
+                privateKeyWith31Bytes.ToByteArray()
+            );
+
+            Assert.Equal(privateKeyWith31Bytes, newPrivateKey);
+        }
+
+        [Fact]
+        public void PrivateKeyGenerateLongerThan31Bytes()
+        {
+            var faults = new List<int>();
+            for (int i = 0; i < 3000; i++)
+            {
+                var pk = new PrivateKey();
+
+                if (pk.ByteArray.Length < 32)
+                {
+                    faults.Add(i);
+                }
+            }
+
+            Assert.Empty(faults);
+        }
     }
 }
diff --git a/Libplanet/Crypto/GenerateKeyParamTriesExceedException.cs b/Libplanet/Crypto/GenerateKeyParamTriesExceedException.cs
@@ -0,0 +1,30 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace Libplanet.Crypto
+{
+    [Serializable]
+    public class GenerateKeyParamTriesExceedException : Exception
+    {
+        public GenerateKeyParamTriesExceedException()
+        {
+        }
+
+        public GenerateKeyParamTriesExceedException(string message)
+          : base(message)
+        {
+        }
+
+        public GenerateKeyParamTriesExceedException(string message, Exception innerException)
+          : base(message, innerException)
+        {
+        }
+
+        protected GenerateKeyParamTriesExceedException(
+          SerializationInfo info, StreamingContext context
+        )
+          : base(info, context)
+        {
+        }
+    }
+}
diff --git a/Libplanet/Crypto/PrivateKey.cs b/Libplanet/Crypto/PrivateKey.cs
@@ -375,7 +375,25 @@ private static ECPrivateKeyParameters GenerateKeyParam()
                 new ECKeyGenerationParameters(ecParams, secureRandom);
             gen.Init(keyGenParam);
 
-            return (ECPrivateKeyParameters)gen.GenerateKeyPair().Private;
+            const int maxTries = 3000;
+            int tries = 0;
+            ECPrivateKeyParameters result;
+
+            while (tries < maxTries)
+            {
+                result = (ECPrivateKeyParameters)gen.GenerateKeyPair().Private;
+                if (result.D.ToByteArrayUnsigned().Length == KeyByteSize)
+                {
+                    return result;
+                }
+
+                tries++;
+            }
+
+            throw new GenerateKeyParamTriesExceedException(
+                "Can't find appropriate parameter for private key" +
+                $"(maxTries: {maxTries})"
+            );
         }
 
         private static ECPrivateKeyParameters GenerateKeyFromBytes(byte[] privateKey)