From ef9b0e4bef9afb433bde2a82b30ab0aae1963393 Mon Sep 17 00:00:00 2001
From: Yang Chun Ung <qooraven@gmail.com>
Date: Thu, 15 Aug 2019 11:22:39 +0900
Subject: [PATCH] Sanity check in PrivateKey constructor using bytes

---
 Libplanet.Tests/Crypto/PrivateKeyTest.cs | 12 ++++++++++
 Libplanet/Crypto/PrivateKey.cs           | 29 ++++++++++++++++++++----
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/Libplanet.Tests/Crypto/PrivateKeyTest.cs b/Libplanet.Tests/Crypto/PrivateKeyTest.cs
index b4c113ef753..f9ac4516606 100644
--- a/Libplanet.Tests/Crypto/PrivateKeyTest.cs
+++ b/Libplanet.Tests/Crypto/PrivateKeyTest.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Text;
 using Libplanet.Crypto;
 using Xunit;
@@ -20,6 +21,17 @@ public void BytesTest()
             Assert.Equal(bs, key.ByteArray);
         }
 
+        [Fact]
+        public void BytesSanityCheckTest()
+        {
+            var bs = new byte[]
+            {
+                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+            };
+            Assert.Throws<InvalidOperationException>(() => new PrivateKey(bs));
+        }
+
         [Fact]
         public void PublicKeyTest()
         {
diff --git a/Libplanet/Crypto/PrivateKey.cs b/Libplanet/Crypto/PrivateKey.cs
index 6ff9663b1d0..2216130740c 100644
--- a/Libplanet/Crypto/PrivateKey.cs
+++ b/Libplanet/Crypto/PrivateKey.cs
@@ -69,11 +69,7 @@ public PrivateKey()
         /// <seealso cref="ByteArray"/>
         public PrivateKey(byte[] privateKey)
             : this(
-                new ECPrivateKeyParameters(
-                    "ECDSA",
-                    new BigInteger(1, privateKey),
-                    GetECParameters()
-                )
+                GenerateKeyFromBytes(privateKey)
             )
         {
         }
@@ -269,6 +265,29 @@ private static ECPrivateKeyParameters GenerateKeyParam()
             return gen.GenerateKeyPair().Private as ECPrivateKeyParameters;
         }
 
+        private static ECPrivateKeyParameters GenerateKeyFromBytes(byte[] privateKey)
+        {
+            var param = new ECPrivateKeyParameters(
+                "ECDSA",
+                new BigInteger(1, privateKey),
+                GetECParameters()
+            );
+
+            var key = new PrivateKey(param);
+            try
+            {
+                var publicKey = key.PublicKey;
+            }
+            catch (ArgumentException)
+            {
+                throw new InvalidOperationException(
+                    "Infinity is not a valid public key for ECDH"
+                );
+            }
+
+            return param;
+        }
+
         private ECPoint CalculatePoint(ECPublicKeyParameters pubKeyParams)
         {
             ECDomainParameters dp = keyParam.Parameters;