Merge pull request #1711 from p2p-org/feature/pwn-1000

[ETH-1000] Handle register request specifically
p2p-org · Feb 20, 2024 · 4b22e71 · 4b22e71
2 parents d857381 + a5efe54
commit 4b22e71
Showing 1 changed file with 74 additions and 24 deletions.
diff --git a/p2p_wallet/Common/Services/ReferralProgram/ReferralProgramService.swift b/p2p_wallet/Common/Services/ReferralProgram/ReferralProgramService.swift
@@ -3,6 +3,7 @@ import KeyAppBusiness
 import KeyAppNetworking
 import Resolver
 import SolanaSwift
+import Task_retrying
 import TweetNacl
 
 protocol ReferralProgramService {
@@ -14,7 +15,10 @@ protocol ReferralProgramService {
 }
 
 enum ReferralProgramServiceError: Error {
-    case failedSet
+    case unauthorized
+    case timeOut
+    case alreadyRegistered
+    case other
 }
 
 final class ReferralProgramServiceImpl {
@@ -50,69 +54,115 @@ extension ReferralProgramServiceImpl: ReferralProgramService {
     func register() async {
         guard !Defaults.referrerRegistered else { return }
         do {
-            guard let secret = userWallet.wallet?.account.secretKey else { throw ReferralProgramServiceError.failedSet }
+            try await Task.retrying(
+                where: { $0.isRetryable },
+                maxRetryCount: 10,
+                retryDelay: 5, // 5 secs
+                timeoutInSeconds: 60, // wait for 60s if no success then throw .timedOut error
+                operation: { [weak self] _ in
+                    // if there is transaction, send it
+                    try await self?.sendRegisterRequest()
+                }
+            ).value
+        } catch {
+            DefaultLogManager.shared.log(
+                event: "\(ReferralProgramService.self)_register",
+                data: error.localizedDescription,
+                logLevel: LogLevel.error
+            )
+        }
+    }
+
+    func setReferent(from: String) async {
+        guard from != currentUserAddress else { return }
+        do {
+            guard let secret = userWallet.wallet?.account.secretKey
+            else { throw ReferralProgramServiceError.unauthorized }
             let timestamp = Int64(Date().timeIntervalSince1970)
-            let signed = try RegisterUserSignature(
-                user: currentUserAddress, referrent: nil, timestamp: timestamp
+            let signed = try SetReferentSignature(
+                user: currentUserAddress, referent: from, timestamp: timestamp
             )
             .sign(secretKey: secret)
             let _: String? = try await jsonrpcClient.request(
                 baseURL: baseURL,
                 body: .init(
-                    method: "register",
-                    params: RegisterUserRequest(
+                    method: "set_referent",
+                    params: SetReferentRequest(
                         user: currentUserAddress,
+                        referent: from,
                         timedSignature: ReferralTimedSignature(
-                            timestamp: timestamp, signature: signed.toHexString()
+                            timestamp: timestamp,
+                            signature: signed.toHexString()
                         )
                     )
                 )
             )
-            Defaults.referrerRegistered = true
         } catch {
-            if error.localizedDescription.contains("duplicate") {
-                // The code is not unique so we look at the description
-                Defaults.referrerRegistered = true
-            }
             debugPrint(error)
             DefaultLogManager.shared.log(
-                event: "\(ReferralProgramService.self)_register",
+                event: "\(ReferralProgramService.self)_setReferent",
                 data: error.localizedDescription,
                 logLevel: LogLevel.error
             )
         }
     }
 
-    func setReferent(from: String) async {
-        guard from != currentUserAddress else { return }
+    private func sendRegisterRequest() async throws {
         do {
-            guard let secret = userWallet.wallet?.account.secretKey else { throw ReferralProgramServiceError.failedSet }
+            guard let secret = userWallet.wallet?.account.secretKey
+            else { throw ReferralProgramServiceError.unauthorized }
             let timestamp = Int64(Date().timeIntervalSince1970)
-            let signed = try SetReferentSignature(
-                user: currentUserAddress, referent: from, timestamp: timestamp
+            let signed = try RegisterUserSignature(
+                user: currentUserAddress, referrent: nil, timestamp: timestamp
             )
             .sign(secretKey: secret)
             let _: String? = try await jsonrpcClient.request(
                 baseURL: baseURL,
                 body: .init(
-                    method: "set_referent",
-                    params: SetReferentRequest(
+                    method: "register",
+                    params: RegisterUserRequest(
                         user: currentUserAddress,
-                        referent: from,
                         timedSignature: ReferralTimedSignature(
-                            timestamp: timestamp,
-                            signature: signed.toHexString()
+                            timestamp: timestamp, signature: signed.toHexString()
                         )
                     )
                 )
             )
+            Defaults.referrerRegistered = true
         } catch {
             debugPrint(error)
             DefaultLogManager.shared.log(
-                event: "\(ReferralProgramService.self)_setReferent",
+                event: "\(ReferralProgramService.self)_register",
                 data: error.localizedDescription,
                 logLevel: LogLevel.error
             )
+            if error.localizedDescription.contains("duplicate key value violates unique constraint") {
+                // The code is not unique so we look at the description
+                Defaults.referrerRegistered = true
+                throw ReferralProgramServiceError.alreadyRegistered
+            }
+            if error.localizedDescription.contains("timed out") || (error as NSError).code == NSURLErrorTimedOut {
+                throw ReferralProgramServiceError.timeOut
+            }
+            throw ReferralProgramServiceError.other
+        }
+    }
+}
+
+// MARK: - Helpers
+
+private extension Swift.Error {
+    var isRetryable: Bool {
+        switch self {
+        case let error as ReferralProgramServiceError:
+            switch error {
+            case .timeOut:
+                return true
+            default:
+                return false
+            }
+        default:
+            return false
         }
     }
 }