From 0636d9c1e584d52afae96a03db466dcfccfe3c40 Mon Sep 17 00:00:00 2001 From: erikzhang Date: Sun, 22 Mar 2020 21:44:59 +0800 Subject: [PATCH 1/2] Optimize the wallet UT --- src/neo/Wallets/SQLite/UserWallet.cs | 98 ++++++++++--------- .../Wallets/SQLite/UT_UserWallet.cs | 3 +- 2 files changed, 52 insertions(+), 49 deletions(-) diff --git a/src/neo/Wallets/SQLite/UserWallet.cs b/src/neo/Wallets/SQLite/UserWallet.cs index 36b865a2e3..92866959af 100644 --- a/src/neo/Wallets/SQLite/UserWallet.cs +++ b/src/neo/Wallets/SQLite/UserWallet.cs @@ -42,56 +42,58 @@ public override Version Version } /// - /// Constructor + /// Open an existing wallet /// /// Path /// Password Key - /// True for create the wallet - /// Scrypt initialization value (only if create=True) - 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(); - 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 + /// + /// Create a new wallet + /// + /// Path + /// Password Key + /// Scrypt initialization value + 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(); + 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 index 87993f12cd..0737554184 100644 --- 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] From 809bdcf936442e7f16ab7b18a298d9115c12cd93 Mon Sep 17 00:00:00 2001 From: Jin Qiao Date: Mon, 23 Mar 2020 13:58:36 +0800 Subject: [PATCH 2/2] Optimize wallet UT test --- .../Wallets/SQLite/UT_UserWallet.cs | 154 ++++++++---------- 1 file changed, 70 insertions(+), 84 deletions(-) diff --git a/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs b/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs index 0737554184..2a1cc699b0 100644 --- a/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs +++ b/tests/neo.UnitTests/Wallets/SQLite/UT_UserWallet.cs @@ -16,76 +16,46 @@ namespace Neo.UnitTests.Wallets.SQLite [TestClass] public class UT_UserWallet { - private string path; - private UserWallet wallet; + private static string path; + private static UserWallet wallet; + private static WalletAccount account; public static string GetRandomPath() { string threadName = Thread.CurrentThread.ManagedThreadId.ToString(); return Path.GetFullPath(string.Format("Wallet_{0}", new Random().Next(1, 1000000).ToString("X8")) + threadName); } - [TestInitialize] - public void Setup() + [ClassInitialize] + public static void Setup(TestContext ctx) { path = GetRandomPath(); wallet = UserWallet.Create(path, "123456", new ScryptParameters(0, 0, 0)); + byte[] privateKey = new byte[32]; + using (RandomNumberGenerator rng = RandomNumberGenerator.Create()) + { + rng.GetBytes(privateKey); + } + account = wallet.CreateAccount(privateKey); } - [TestCleanup] - public void Cleanup() + [ClassCleanup] + public static void Cleanup() { TestUtils.DeleteFile(path); } [TestMethod] - public void TestGetName() - { - wallet.Name.Should().Be(Path.GetFileNameWithoutExtension(path)); - } - - [TestMethod] - public void TestGetVersion() - { - Action action = () => wallet.Version.ToString(); - action.Should().NotThrow(); - } - - [TestMethod] - public void TestCreateAndOpenSecureString() + public void TestChangePassword() { - string myPath = GetRandomPath(); - var ss = new SecureString(); - ss.AppendChar('a'); - ss.AppendChar('b'); - ss.AppendChar('c'); - - var w1 = UserWallet.Create(myPath, ss); - w1.Should().NotBeNull(); - - var w2 = UserWallet.Open(myPath, ss); - w2.Should().NotBeNull(); - - ss.AppendChar('d'); - Action action = () => UserWallet.Open(myPath, ss); - action.Should().Throw(); - - TestUtils.DeleteFile(myPath); + wallet.ChangePassword("123455", "654321").Should().BeFalse(); + wallet.ChangePassword("123456", "654321").Should().BeTrue(); + wallet.ChangePassword("654321", "123456").Should().BeTrue(); } [TestMethod] - public void TestOpen() + public void TestContains() { - byte[] privateKey = new byte[32]; - using (RandomNumberGenerator rng = RandomNumberGenerator.Create()) - { - rng.GetBytes(privateKey); - } - var account = wallet.CreateAccount(privateKey); - var w1 = UserWallet.Open(path, "123456"); - w1.Should().NotBeNull(); - - Action action = () => UserWallet.Open(path, "123"); - action.Should().Throw(); + wallet.Contains(account.ScriptHash).Should().BeTrue(); } [TestMethod] @@ -103,6 +73,8 @@ public void TestCreateAccountAndGetByPrivateKey() var account1 = wallet.CreateAccount(privateKey); var dbAccount1 = wallet.GetAccount(account1.ScriptHash); account1.Should().Be(dbAccount1); + wallet.DeleteAccount(account.ScriptHash); + wallet.DeleteAccount(account1.ScriptHash); } [TestMethod] @@ -111,6 +83,7 @@ public void TestCreateAccountByScriptHash() var account = wallet.CreateAccount(UInt160.Parse("0xa6ee944042f3c7ea900481a95d65e4a887320cf0")); var dbAccount = wallet.GetAccount(account.ScriptHash); account.Should().Be(dbAccount); + wallet.DeleteAccount(account.ScriptHash); } [TestMethod] @@ -145,64 +118,77 @@ public void TestCreateAccountBySmartContract() var account2 = wallet.CreateAccount(contract2, key2); var dbAccount2 = wallet.GetAccount(account2.ScriptHash); account2.Should().Be(dbAccount2); + wallet.DeleteAccount(account.ScriptHash); + wallet.DeleteAccount(account2.ScriptHash); } [TestMethod] - public void TestDeleteAccount() + public void TestCreateAndOpenSecureString() { - bool ret = wallet.DeleteAccount(UInt160.Parse("0xa6ee944042f3c7ea900481a95d65e4a887320cf0")); - ret.Should().BeFalse(); + string myPath = GetRandomPath(); + var ss = new SecureString(); + ss.AppendChar('a'); + ss.AppendChar('b'); + ss.AppendChar('c'); - byte[] privateKey = new byte[32]; - using (RandomNumberGenerator rng = RandomNumberGenerator.Create()) + var w1 = UserWallet.Create(myPath, ss, new ScryptParameters(0, 0, 0)); + w1.Should().NotBeNull(); + + var w2 = UserWallet.Open(myPath, ss); + w2.Should().NotBeNull(); + + ss.AppendChar('d'); + Action action = () => UserWallet.Open(myPath, ss); + action.Should().Throw(); + + TestUtils.DeleteFile(myPath); + } + + [TestMethod] + public void TestGetAccounts() + { + var ret = wallet.GetAccounts(); + ret.Should().NotBeEmpty(); + foreach (var dbAccount in ret) { - rng.GetBytes(privateKey); + dbAccount.Should().Be(account); } - var account = wallet.CreateAccount(privateKey); - bool ret2 = wallet.DeleteAccount(account.ScriptHash); - ret2.Should().BeTrue(); } [TestMethod] - public void TestChangePassword() + public void TestGetName() { - wallet.ChangePassword("123455", "654321").Should().BeFalse(); - wallet.ChangePassword("123456", "654321").Should().BeTrue(); + wallet.Name.Should().Be(Path.GetFileNameWithoutExtension(path)); } [TestMethod] - public void TestContains() + public void TestGetVersion() { - byte[] privateKey = new byte[32]; - using (RandomNumberGenerator rng = RandomNumberGenerator.Create()) - { - rng.GetBytes(privateKey); - } - var account = wallet.CreateAccount(privateKey); - wallet.Contains(account.ScriptHash).Should().BeTrue(); + Action action = () => wallet.Version.ToString(); + action.Should().NotThrow(); } [TestMethod] - public void TestGetAccounts() + public void TestOpen() { - var ret = wallet.GetAccounts(); - ret.Should().BeNullOrEmpty(); + var w1 = UserWallet.Open(path, "123456"); + w1.Should().NotBeNull(); - byte[] privateKey = new byte[32]; - using (RandomNumberGenerator rng = RandomNumberGenerator.Create()) - { - rng.GetBytes(privateKey); - } - var account = wallet.CreateAccount(privateKey); - ret = wallet.GetAccounts(); - foreach (var dbAccount in ret) - { - dbAccount.Should().Be(account); - } + Action action = () => UserWallet.Open(path, "123"); + action.Should().Throw(); + } + + [TestMethod] + public void TestToDeleteAccount() + { + bool ret = wallet.DeleteAccount(UInt160.Parse("0xa6ee944042f3c7ea900481a95d65e4a887320cf0")); + ret.Should().BeFalse(); + bool ret2 = wallet.DeleteAccount(account.ScriptHash); + ret2.Should().BeTrue(); } [TestMethod] - public void TestVerifyPassword() + public void TestToVerifyPassword() { wallet.VerifyPassword("123456").Should().BeTrue(); wallet.VerifyPassword("123").Should().BeFalse();