Merge pull request #103 from torusresearch/import_check

add check for import share
torusresearch · Sep 13, 2024 · 23083bf · 23083bf
2 parents 303dc2c + 4bc19c7
commit 23083bf
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 8 deletions.
diff --git a/Sources/TorusUtils/Extensions/String+Extension.swift b/Sources/TorusUtils/Extensions/String+Extension.swift
@@ -36,6 +36,15 @@ extension String {
             return self
         }
     }
+
+    public func addLeading0sForLength128() -> String {
+        if count < 128 {
+            let toAdd = String(repeating: "0", count: 128 - count)
+            return toAdd + self
+        } else {
+            return self
+        }
+    }
 
     public func hexEncodedToString() -> String {
         var finalString = ""

diff --git a/Sources/TorusUtils/Helpers/KeyUtils.swift b/Sources/TorusUtils/Helpers/KeyUtils.swift
@@ -64,7 +64,10 @@ public class KeyUtils {
             publicKeyUnprefixed = publicKeyUnprefixed.strip04Prefix()
         }
 
-        if !(publicKeyUnprefixed.count == 128) {
+
+        if (publicKeyUnprefixed.count <= 128) {
+            publicKeyUnprefixed = publicKeyUnprefixed.addLeading0sForLength128()
+        } else {
             throw TorusUtilError.invalidPubKeySize
         }
 

diff --git a/Sources/TorusUtils/Helpers/NodeUtils.swift b/Sources/TorusUtils/Helpers/NodeUtils.swift
@@ -130,6 +130,7 @@ internal class NodeUtils {
         idToken: String,
         importedShares: [ImportedShare]?,
         apiKey: String = "torus-default",
+        newPrivateKey: String?,
         extraParams: TorusUtilsExtraParams
     ) async throws -> TorusKey {
         let threshold = Int(trunc(Double((endpoints.count / 2) + 1)))
@@ -402,10 +403,10 @@ internal class NodeUtils {
         var sessionTokens: [String?] = []
         var nodeIndexes: [Int?] = []
         var sessionTokenDatas: [SessionToken?] = []
-        var isNewKeys: [String] = []
+        var isNewKeys: [IsNewKeyResponse] = []
 
         for item in shareResponses {
-            isNewKeys.append(item.isNewKey)
+            isNewKeys.append(IsNewKeyResponse(isNewKey: item.isNewKey == "true", publicKeyX: item.keys.first?.publicKey.X ?? ""))
 
             if !item.sessionTokenSigs.isEmpty {
                 if !item.sessionTokenSigMetadata.isEmpty {
@@ -505,7 +506,12 @@ internal class NodeUtils {
             throw TorusUtilError.privateKeyDeriveFailed
         }
 
-        let thresholdIsNewKey: String? = try thresholdSame(arr: isNewKeys, threshold: threshold)
+        var isNewKey = false;
+        for item in isNewKeys {
+            if (item.isNewKey && item.publicKeyX.lowercased() == thresholdPublicKey!.X.lowercased()) {
+                isNewKey = true
+            }
+        }
 
         let oAuthKey = privateKey!
         let oAuthPublicKey = try SecretKey(hex: oAuthKey).toPublic().serialize(compressed: false)
@@ -519,8 +525,7 @@ internal class NodeUtils {
             finalPubKey = oAuthPublicKey
         } else if TorusUtils.isLegacyNetworkRouteMap(network: network) {
             if enableOneKey {
-                let isNewKey = !(thresholdIsNewKey == "true")
-                let nonce = try await MetadataUtils.getOrSetNonce(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffsetResponse, X: thresholdPublicKey!.X, Y: thresholdPublicKey!.Y, privateKey: oAuthKey, getOnly: isNewKey)
+                let nonce = try await MetadataUtils.getOrSetNonce(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffsetResponse, X: thresholdPublicKey!.X, Y: thresholdPublicKey!.Y, privateKey: oAuthKey, getOnly: !isNewKey)
                 metadataNonce = BigInt(nonce.nonce?.addLeading0sForLength64() ?? "0", radix: 16) ?? BigInt(0)
                 typeOfUser = UserType(rawValue: nonce.typeOfUser?.lowercased() ?? "v1")!
                 if typeOfUser == .v2 {
@@ -569,6 +574,15 @@ internal class NodeUtils {
             finalPrivKey = privateKeyWithNonce.magnitude.serialize().hexString.addLeading0sForLength64()
         }
 
+        // This is a sanity check to make doubly sure we are returning the correct private key after importing a share
+        if isImportShareReq {
+            if newPrivateKey == nil {
+                throw TorusUtilError.importShareFailed
+            } else if (!(finalPrivKey == newPrivateKey!.addLeading0sForLength64())) {
+                throw TorusUtilError.importShareFailed
+            }
+        }
+
         var isUpgraded: Bool?
         if typeOfUser == .v2 {
             isUpgraded = metadataNonce == BigInt(0)

diff --git a/Sources/TorusUtils/Helpers/jsonRPC/Responses/IsNewKeyResponse.swift b/Sources/TorusUtils/Helpers/jsonRPC/Responses/IsNewKeyResponse.swift
@@ -0,0 +1,11 @@
+import Foundation
+
+internal struct IsNewKeyResponse: Codable {
+    public var isNewKey: Bool;
+    public var publicKeyX: String;
+
+    public init(isNewKey: Bool, publicKeyX: String) {
+        self.isNewKey = isNewKey
+        self.publicKeyX = publicKeyX
+    }
+}
diff --git a/Sources/TorusUtils/TorusUtils.swift b/Sources/TorusUtils/TorusUtils.swift
@@ -131,7 +131,7 @@ public class TorusUtils {
             extraParams.session_token_exp_second = sessionTime
         }
 
-        return try await NodeUtils.retrieveOrImportShare(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffset, enableOneKey: enableOneKey, allowHost: allowHost, network: network, clientId: clientId, endpoints: endpoints, verifier: verifier, verifierParams: verifierParams, idToken: idToken, importedShares: [], apiKey: apiKey, extraParams: extraParams)
+        return try await NodeUtils.retrieveOrImportShare(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffset, enableOneKey: enableOneKey, allowHost: allowHost, network: network, clientId: clientId, endpoints: endpoints, verifier: verifier, verifierParams: verifierParams, idToken: idToken, importedShares: [], apiKey: apiKey, newPrivateKey: nil, extraParams: extraParams)
     }
 
     /// Retrieves user information, defaulting the user type to .v2
@@ -185,7 +185,7 @@ public class TorusUtils {
             extraParams.session_token_exp_second = sessionTime
         }
 
-        return try await NodeUtils.retrieveOrImportShare(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffset ?? 0, enableOneKey: enableOneKey, allowHost: allowHost, network: network, clientId: clientId, endpoints: endpoints, verifier: verifier, verifierParams: verifierParams, idToken: idToken, importedShares: sharesData, extraParams: extraParams)
+        return try await NodeUtils.retrieveOrImportShare(legacyMetadataHost: legacyMetadataHost, serverTimeOffset: serverTimeOffset ?? 0, enableOneKey: enableOneKey, allowHost: allowHost, network: network, clientId: clientId, endpoints: endpoints, verifier: verifier, verifierParams: verifierParams, idToken: idToken, importedShares: sharesData, newPrivateKey: newPrivateKey, extraParams: extraParams)
     }
 
     /// Retrieves user information