Fix Serialisation of BIP32 Keystore

Sources/web3swift/KeystoreManager/BIP32Keystore.swift

@@ -14,21 +14,18 @@ public class BIP32Keystore: AbstractKeystore {
 
     public var addresses: [EthereumAddress]? {
         get {
-            if self.paths.count == 0 {
+            let addresses = self.addressStorage.addresses
+            if addresses.count == 0 {
                 return nil
             }
-            var allAccounts = [EthereumAddress]()
-            for (_, address) in paths {
-                allAccounts.append(address)
-            }
-            return allAccounts
+            return addresses
         }
     }
 
     public var isHDKeystore: Bool = true
 
     public func UNSAFE_getPrivateKeyData(password: String, account: EthereumAddress) throws -> Data {
-        if let key = self.paths.keyForValue(value: account) {
+        if let key = addressStorage.path(by: account) {
             guard let decryptedRootNode = try? self.getPrefixNodeData(password) else {throw AbstractKeystoreError.encryptionError("Failed to decrypt a keystore")}
             guard let rootNode = HDNode(decryptedRootNode) else {throw AbstractKeystoreError.encryptionError("Failed to deserialize a root node")}
             guard rootNode.depth == (self.rootPrefix.components(separatedBy: "/").count - 1) else {throw AbstractKeystoreError.encryptionError("Derivation depth mismatch")}
@@ -44,8 +41,10 @@ public class BIP32Keystore: AbstractKeystore {
     // --------------
 
     public var keystoreParams: KeystoreParamsBIP32?
-    public var paths: [String:EthereumAddress] = [String:EthereumAddress]()
     public var rootPrefix: String
+
+    private (set) var addressStorage: PathAddressStorage
+
     public convenience init?(_ jsonString: String) {
         let lowercaseJSON = jsonString.lowercased()
         guard let jsonData = lowercaseJSON.data(using: .utf8) else {return nil}
@@ -57,9 +56,7 @@ public class BIP32Keystore: AbstractKeystore {
         if (keystorePars.version != 3) {return nil}
         if (keystorePars.crypto.version != nil && keystorePars.crypto.version != "1") {return nil}
         if (!keystorePars.isHDWallet) {return nil}
-        for (p, ad) in keystorePars.pathToAddress {
-            paths[p] = EthereumAddress(ad)
-        }
+        addressStorage = PathAddressStorage(pathAddressPairs: keystorePars.pathAddressPairs)
         if keystorePars.rootPath == nil {
             keystorePars.rootPath = HDNode.defaultPathPrefix
         }
@@ -74,6 +71,7 @@ public class BIP32Keystore: AbstractKeystore {
     }
 
     public init? (seed: Data, password: String = "web3swift", prefixPath: String = HDNode.defaultPathMetamaskPrefix, aesMode: String = "aes-128-cbc") throws {
+        addressStorage = PathAddressStorage()
         guard let rootNode = HDNode(seed: seed)?.derive(path: prefixPath, derivePrivateKey: true) else {return nil}
         self.rootPrefix = prefixPath
         try createNewAccount(parentNode: rootNode, password: password)
@@ -93,7 +91,7 @@ public class BIP32Keystore: AbstractKeystore {
 
     func createNewAccount(parentNode: HDNode, password: String = "web3swift") throws {
         var newIndex = UInt32(0)
-        for (p, _) in paths {
+        for p in addressStorage.paths {
             guard let idx = UInt32(p.components(separatedBy: "/").last!) else {continue}
             if idx >= newIndex {
                 newIndex = idx + 1
@@ -108,7 +106,7 @@ public class BIP32Keystore: AbstractKeystore {
         } else {
             newPath = prefixPath + "/" + String(newNode.index)
         }
-        paths[newPath] = newAddress
+        addressStorage.add(address: newAddress, for: newPath)
     }
 
     public func createNewCustomChildAccount(password: String = "web3swift", path: String) throws {
@@ -143,7 +141,7 @@ public class BIP32Keystore: AbstractKeystore {
         } else {
             newPath = prefixPath + "/" + pathAppendix!
         }
-        paths[newPath] = newAddress
+        addressStorage.add(address: newAddress, for: newPath)
         guard let serializedRootNode = rootNode.serialize(serializePublic: false) else {throw AbstractKeystoreError.keyDerivationError}
         try encryptDataToStorage(password, data: serializedRootNode, aesMode: self.keystoreParams!.crypto.cipher)
     }
@@ -183,12 +181,8 @@ public class BIP32Keystore: AbstractKeystore {
         let kdfparams = KdfParamsV3(salt: saltData.toHexString(), dklen: dkLen, n: N, p: P, r: R, c: nil, prf: nil)
         let cipherparams = CipherParamsV3(iv: IV.toHexString())
         let crypto = CryptoParamsV3(ciphertext: encryptedKeyData.toHexString(), cipher: aesMode, cipherparams: cipherparams, kdf: "scrypt", kdfparams: kdfparams, mac: mac.toHexString(), version: nil)
-        var pathToAddress = [String:String]()
-        for (path, address) in paths {
-            pathToAddress[path] = address.address
-        }
         var keystorePars = KeystoreParamsBIP32(crypto: crypto, id: UUID().uuidString.lowercased(), version: 3)
-        keystorePars.pathToAddress = pathToAddress
+        keystorePars.pathAddressPairs = addressStorage.toPathAddressPairs()
         keystorePars.rootPath = self.rootPrefix
         keystoreParams = keystorePars
     }

Sources/web3swift/KeystoreManager/BIP32KeystoreJSONStructure.swift

@@ -6,20 +6,25 @@
 
 import Foundation
 
+public struct PathAddressPair: Codable {
+    let path: String
+    let address: String
+}
+
 public struct KeystoreParamsBIP32: Decodable, Encodable {
     var crypto: CryptoParamsV3
     var id: String?
     var version: Int = 32
     var isHDWallet: Bool
-    var pathToAddress: [String:String]
+    var pathAddressPairs: [PathAddressPair]
     var rootPath: String?
 
     public init(crypto cr: CryptoParamsV3, id i: String, version ver: Int, rootPath: String? = nil) {
         crypto = cr
         id = i
         version = ver
         isHDWallet = true
-        pathToAddress = [String:String]()
+        pathAddressPairs = [PathAddressPair]()
         self.rootPath = rootPath
     }
 

Sources/web3swift/KeystoreManager/PathAddressStorage.swift

@@ -0,0 +1,56 @@
+//
+//  PathAddressStorage.swift
+//  web3swift
+//
+//  Created by Andrew Podkovyrin on 08.08.2020.
+//  Copyright © 2020 Matter Labs. All rights reserved.
+//
+
+import Foundation
+
+public struct PathAddressStorage {
+    private(set) var addresses: [EthereumAddress]
+    private(set) var paths: [String]
+
+    init() {
+        addresses = []
+        paths = []
+    }
+
+    mutating func add(address: EthereumAddress, for path: String) {
+        addresses.append(address)
+        paths.append(path)
+    }
+
+    func path(by address: EthereumAddress) -> String? {
+        guard let index = addresses.firstIndex(of: address) else { return nil }
+        return paths[index]
+    }
+}
+
+extension PathAddressStorage {
+    init(pathAddressPairs: [PathAddressPair]) {
+        var addresses = [EthereumAddress]()
+        var paths = [String]()
+        for pair in pathAddressPairs {
+            guard let address = EthereumAddress(pair.address) else { continue }
+            addresses.append(address)
+            paths.append(pair.path)
+        }
+
+        assert(addresses.count == paths.count)
+
+        self.addresses = addresses
+        self.paths = paths
+    }
+
+    func toPathAddressPairs() -> [PathAddressPair] {
+        var pathAddressPairs = [PathAddressPair]()
+        for (index, path) in paths.enumerated() {
+            let address = addresses[index]
+            let pair = PathAddressPair(path: path, address: address.address)
+            pathAddressPairs.append(pair)
+        }
+        return pathAddressPairs
+    }
+}

Tests/web3swiftTests/web3swift_keystores_Tests.swift

@@ -137,7 +137,7 @@ class web3swift_Keystores_tests: XCTestCase {
         let account = keystore!.addresses![1]
         let key = try! keystore!.UNSAFE_getPrivateKeyData(password: "", account: account)
         XCTAssertNotNil(key)
-        print(keystore!.paths)
+        print(keystore!.addressStorage.paths)
     }
 
     func testByBIP32keystoreSaveAndDeriva() {
@@ -155,9 +155,8 @@ class web3swift_Keystores_tests: XCTestCase {
         print(keystore!.addresses![1].address)
         print(recreatedStore!.addresses![0].address)
         print(recreatedStore!.addresses![1].address)
-        // This will fail. It wont fail if use scrypt from pod 'scrypt', '2.0', not from CryptoSwift
-        XCTAssert(keystore?.addresses![0] == recreatedStore?.addresses![1])
-        XCTAssert(keystore?.addresses![1] == recreatedStore?.addresses![0])
+        XCTAssert(keystore?.addresses![0] == recreatedStore?.addresses![0])
+        XCTAssert(keystore?.addresses![1] == recreatedStore?.addresses![1])
     }
 
     //    func testPBKDF2() {

web3swift.xcodeproj/project.pbxproj

@@ -53,6 +53,7 @@
 		13C3392521B6C62400F33F5E /* secp256k1_ec_mult_static_context.h in Headers */ = {isa = PBXBuildFile; fileRef = 13C338F621B6C62400F33F5E /* secp256k1_ec_mult_static_context.h */; };
 		13C3392621B6C62400F33F5E /* scratch.h in Headers */ = {isa = PBXBuildFile; fileRef = 13C338F721B6C62400F33F5E /* scratch.h */; };
 		13C3392821B6C68900F33F5E /* secp256k1.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 13C3388E21B6C2DD00F33F5E /* secp256k1.framework */; settings = {ATTRIBUTES = (Weak, ); }; };
+		2AC22E362525C2000072F037 /* PathAddressStorage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2AC22E352525C2000072F037 /* PathAddressStorage.swift */; };
 		3A7EA35E2280EA9A005120C2 /* Encodable+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 3A7EA35D2280EA9A005120C2 /* Encodable+Extensions.swift */; };
 		3A7EA3602280EB27005120C2 /* Decodable+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 3A7EA35F2280EB27005120C2 /* Decodable+Extensions.swift */; };
 		3AA8151C2276E42F00F5DB52 /* EventFiltering.swift in Sources */ = {isa = PBXBuildFile; fileRef = 3AA815172276E42F00F5DB52 /* EventFiltering.swift */; };
@@ -249,6 +250,7 @@
 		13C338F721B6C62400F33F5E /* scratch.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = scratch.h; sourceTree = "<group>"; };
 		13CE02B021FC846800CE7148 /* RELEASE_GUIDE.md */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = net.daringfireball.markdown; path = RELEASE_GUIDE.md; sourceTree = "<group>"; };
 		13CE02B121FC846900CE7148 /* BUILD_GUIDE.md */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = net.daringfireball.markdown; path = BUILD_GUIDE.md; sourceTree = "<group>"; };
+		2AC22E352525C2000072F037 /* PathAddressStorage.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PathAddressStorage.swift; sourceTree = "<group>"; };
 		3A7EA35D2280EA9A005120C2 /* Encodable+Extensions.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Encodable+Extensions.swift"; sourceTree = "<group>"; };
 		3A7EA35F2280EB27005120C2 /* Decodable+Extensions.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Decodable+Extensions.swift"; sourceTree = "<group>"; };
 		3AA815172276E42F00F5DB52 /* EventFiltering.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = EventFiltering.swift; sourceTree = "<group>"; };
@@ -641,6 +643,7 @@
 		3AA815302276E44100F5DB52 /* KeystoreManager */ = {
 			isa = PBXGroup;
 			children = (
+				2AC22E352525C2000072F037 /* PathAddressStorage.swift */,
 				3AA815312276E44100F5DB52 /* KeystoreManager.swift */,
 				3AA815322276E44100F5DB52 /* IBAN.swift */,
 				3AA815332276E44100F5DB52 /* BIP39.swift */,
@@ -1304,6 +1307,7 @@
 				3AA815D62276E44100F5DB52 /* Web3+ReadingTransaction.swift in Sources */,
 				3AA815AC2276E44100F5DB52 /* NameHash.swift in Sources */,
 				3AA8151C2276E42F00F5DB52 /* EventFiltering.swift in Sources */,
+				2AC22E362525C2000072F037 /* PathAddressStorage.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};