Optimize the wallet UT (neo-project#1496)

Tommo-L · Jun 22, 2020 · 8fdf778 · 8fdf778
1 parent 90b2564
commit 8fdf778
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 49 deletions.
diff --git a/src/neo/Wallets/SQLite/UserWallet.cs b/src/neo/Wallets/SQLite/UserWallet.cs
@@ -42,56 +42,58 @@ public override Version Version
         }
 
         /// <summary>
-        /// Constructor
+        /// Open an existing wallet
         /// </summary>
         /// <param name="path">Path</param>
         /// <param name="passwordKey">Password Key</param>
-        /// <param name="create">True for create the wallet</param>
-        /// <param name="scrypt">Scrypt initialization value (only if create=True)</param>
-        private UserWallet(string path, byte[] passwordKey, bool create, ScryptParameters scrypt = null)
+        private UserWallet(string path, byte[] passwordKey)
         {
             this.path = path;
+            this.salt = LoadStoredData("Salt");
+            byte[] passwordHash = LoadStoredData("PasswordHash");
+            if (passwordHash != null && !passwordHash.SequenceEqual(passwordKey.Concat(salt).ToArray().Sha256()))
+                throw new CryptographicException();
+            this.iv = LoadStoredData("IV");
+            this.masterKey = LoadStoredData("MasterKey").AesDecrypt(passwordKey, iv);
+            this.scrypt = new ScryptParameters
+                (
+                BitConverter.ToInt32(LoadStoredData("ScryptN")),
+                BitConverter.ToInt32(LoadStoredData("ScryptR")),
+                BitConverter.ToInt32(LoadStoredData("ScryptP"))
+                );
+            this.accounts = LoadAccounts();
+        }
 
-            if (create)
-            {
-                this.iv = new byte[16];
-                this.salt = new byte[20];
-                this.masterKey = new byte[32];
-                this.scrypt = scrypt ?? ScryptParameters.Default;
-                this.accounts = new Dictionary<UInt160, UserWalletAccount>();
-                using (RandomNumberGenerator rng = RandomNumberGenerator.Create())
-                {
-                    rng.GetBytes(iv);
-                    rng.GetBytes(salt);
-                    rng.GetBytes(masterKey);
-                }
-                Version version = Assembly.GetExecutingAssembly().GetName().Version;
-                BuildDatabase();
-                SaveStoredData("IV", iv);
-                SaveStoredData("Salt", salt);
-                SaveStoredData("PasswordHash", passwordKey.Concat(salt).ToArray().Sha256());
-                SaveStoredData("MasterKey", masterKey.AesEncrypt(passwordKey, iv));
-                SaveStoredData("Version", new[] { version.Major, version.Minor, version.Build, version.Revision }.Select(p => BitConverter.GetBytes(p)).SelectMany(p => p).ToArray());
-                SaveStoredData("ScryptN", BitConverter.GetBytes(this.scrypt.N));
-                SaveStoredData("ScryptR", BitConverter.GetBytes(this.scrypt.R));
-                SaveStoredData("ScryptP", BitConverter.GetBytes(this.scrypt.P));
-            }
-            else
+        /// <summary>
+        /// Create a new wallet
+        /// </summary>
+        /// <param name="path">Path</param>
+        /// <param name="passwordKey">Password Key</param>
+        /// <param name="scrypt">Scrypt initialization value</param>
+        private UserWallet(string path, byte[] passwordKey, ScryptParameters scrypt)
+        {
+            this.path = path;
+            this.iv = new byte[16];
+            this.salt = new byte[20];
+            this.masterKey = new byte[32];
+            this.scrypt = scrypt;
+            this.accounts = new Dictionary<UInt160, UserWalletAccount>();
+            using (RandomNumberGenerator rng = RandomNumberGenerator.Create())
             {
-                this.salt = LoadStoredData("Salt");
-                byte[] passwordHash = LoadStoredData("PasswordHash");
-                if (passwordHash != null && !passwordHash.SequenceEqual(passwordKey.Concat(salt).ToArray().Sha256()))
-                    throw new CryptographicException();
-                this.iv = LoadStoredData("IV");
-                this.masterKey = LoadStoredData("MasterKey").AesDecrypt(passwordKey, iv);
-                this.scrypt = new ScryptParameters
-                    (
-                    BitConverter.ToInt32(LoadStoredData("ScryptN")),
-                    BitConverter.ToInt32(LoadStoredData("ScryptR")),
-                    BitConverter.ToInt32(LoadStoredData("ScryptP"))
-                    );
-                this.accounts = LoadAccounts();
+                rng.GetBytes(iv);
+                rng.GetBytes(salt);
+                rng.GetBytes(masterKey);
             }
+            Version version = Assembly.GetExecutingAssembly().GetName().Version;
+            BuildDatabase();
+            SaveStoredData("IV", iv);
+            SaveStoredData("Salt", salt);
+            SaveStoredData("PasswordHash", passwordKey.Concat(salt).ToArray().Sha256());
+            SaveStoredData("MasterKey", masterKey.AesEncrypt(passwordKey, iv));
+            SaveStoredData("Version", new[] { version.Major, version.Minor, version.Build, version.Revision }.Select(p => BitConverter.GetBytes(p)).SelectMany(p => p).ToArray());
+            SaveStoredData("ScryptN", BitConverter.GetBytes(this.scrypt.N));
+            SaveStoredData("ScryptR", BitConverter.GetBytes(this.scrypt.R));
+            SaveStoredData("ScryptP", BitConverter.GetBytes(this.scrypt.P));
         }
 
         private void AddAccount(UserWalletAccount account)
@@ -192,14 +194,14 @@ public override bool Contains(UInt160 scriptHash)
             }
         }
 
-        public static UserWallet Create(string path, string password)
+        public static UserWallet Create(string path, string password, ScryptParameters scrypt = null)
         {
-            return new UserWallet(path, password.ToAesKey(), true);
+            return new UserWallet(path, password.ToAesKey(), scrypt ?? ScryptParameters.Default);
         }
 
-        public static UserWallet Create(string path, SecureString password)
+        public static UserWallet Create(string path, SecureString password, ScryptParameters scrypt = null)
         {
-            return new UserWallet(path, password.ToAesKey(), true);
+            return new UserWallet(path, password.ToAesKey(), scrypt ?? ScryptParameters.Default);
         }
 
         public override WalletAccount CreateAccount(byte[] privateKey)
@@ -326,12 +328,12 @@ private byte[] LoadStoredData(string name)
 
         public static UserWallet Open(string path, string password)
         {
-            return new UserWallet(path, password.ToAesKey(), false);
+            return new UserWallet(path, password.ToAesKey());
         }
 
         public static UserWallet Open(string path, SecureString password)
         {
-            return new UserWallet(path, password.ToAesKey(), false);
+            return new UserWallet(path, password.ToAesKey());
         }
 
         private void SaveStoredData(string name, byte[] value)

diff --git a/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs b/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs
@@ -2,6 +2,7 @@
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using Neo.SmartContract;
 using Neo.Wallets;
+using Neo.Wallets.NEP6;
 using Neo.Wallets.SQLite;
 using System;
 using System.IO;
@@ -27,7 +28,7 @@ public static string GetRandomPath()
         public void Setup()
         {
             path = GetRandomPath();
-            wallet = UserWallet.Create(path, "123456");
+            wallet = UserWallet.Create(path, "123456", new ScryptParameters(0, 0, 0));
         }
 
         [TestCleanup]