Mutli Wallet Support (#434)

* update package * add ability to add and remove accounts * add tests for it * bump th epod
xmtp · Nov 21, 2024 · 9830fe8 · 9830fe8
1 parent 5ce4f85
commit 9830fe8
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 43 deletions.
diff --git a/Sources/XMTPiOS/Client.swift b/Sources/XMTPiOS/Client.swift
@@ -164,7 +164,7 @@ public final class Client {
 		)
 	}
 
-	static func initFFiClient(
+	private static func initFFiClient(
 		accountAddress: String,
 		options: ClientOptions,
 		signingKey: SigningKey?,
@@ -213,28 +213,8 @@ public final class Client {
 		if let signatureRequest = ffiClient.signatureRequest() {
 			if let signingKey = signingKey {
 				do {
-					if signingKey.type == WalletType.SCW {
-						guard let chainId = signingKey.chainId else {
-							throw ClientError.creationError(
-								"Chain id must be present to sign Smart Contract Wallet"
-							)
-						}
-						let signedData = try await signingKey.signSCW(
-							message: signatureRequest.signatureText())
-						try await signatureRequest.addScwSignature(
-							signatureBytes: signedData,
-							address: signingKey.address.lowercased(),
-							chainId: UInt64(chainId),
-							blockNumber: signingKey.blockNumber.flatMap {
-								$0 >= 0 ? UInt64($0) : nil
-							})
-
-					} else {
-						let signedData = try await signingKey.sign(
-							message: signatureRequest.signatureText())
-						try await signatureRequest.addEcdsaSignature(
-							signatureBytes: signedData.rawData)
-					}
+					try await handleSignature(
+						for: signatureRequest, signingKey: signingKey)
 					try await ffiClient.registerIdentity(
 						signatureRequest: signatureRequest)
 				} catch {
@@ -249,11 +229,36 @@ public final class Client {
 			}
 		}
 
-		print("LibXMTP \(getVersionInfo())")
-
 		return (ffiClient, dbURL)
 	}
 
+	private static func handleSignature(
+		for signatureRequest: FfiSignatureRequest,
+		signingKey: SigningKey
+	) async throws {
+		if signingKey.type == .SCW {
+			guard let chainId = signingKey.chainId else {
+				throw ClientError.creationError(
+					"Chain id must be present to sign Smart Contract Wallet")
+			}
+			let signedData = try await signingKey.signSCW(
+				message: signatureRequest.signatureText())
+			try await signatureRequest.addScwSignature(
+				signatureBytes: signedData,
+				address: signingKey.address.lowercased(),
+				chainId: UInt64(chainId),
+				blockNumber: signingKey.blockNumber.flatMap {
+					$0 >= 0 ? UInt64($0) : nil
+				}
+			)
+		} else {
+			let signedData = try await signingKey.sign(
+				message: signatureRequest.signatureText())
+			try await signatureRequest.addEcdsaSignature(
+				signatureBytes: signedData.rawData)
+		}
+	}
+
 	public static func getOrCreateInboxId(
 		api: ClientOptions.Api, address: String
 	) async throws -> String {
@@ -287,6 +292,55 @@ public final class Client {
 		self.environment = environment
 	}
 
+	public func addAccount(recoveryAccount: SigningKey, newAccount: SigningKey)
+		async throws
+	{
+		let signatureRequest = try await ffiClient.addWallet(
+			existingWalletAddress: recoveryAccount.address.lowercased(),
+			newWalletAddress: newAccount.address.lowercased())
+		do {
+			try await Client.handleSignature(
+				for: signatureRequest, signingKey: recoveryAccount)
+			try await Client.handleSignature(
+				for: signatureRequest, signingKey: newAccount)
+			try await ffiClient.applySignatureRequest(
+				signatureRequest: signatureRequest)
+		} catch {
+			throw ClientError.creationError(
+				"Failed to sign the message: \(error.localizedDescription)")
+		}
+	}
+
+	public func removeAccount(
+		recoveryAccount: SigningKey, addressToRemove: String
+	) async throws {
+		let signatureRequest = try await ffiClient.revokeWallet(
+			walletAddress: addressToRemove.lowercased())
+		do {
+			try await Client.handleSignature(
+				for: signatureRequest, signingKey: recoveryAccount)
+			try await ffiClient.applySignatureRequest(
+				signatureRequest: signatureRequest)
+		} catch {
+			throw ClientError.creationError(
+				"Failed to sign the message: \(error.localizedDescription)")
+		}
+	}
+
+	public func revokeAllOtherInstallations(signingKey: SigningKey) async throws
+	{
+		let signatureRequest = try await ffiClient.revokeAllOtherInstallations()
+		do {
+			try await Client.handleSignature(
+				for: signatureRequest, signingKey: signingKey)
+			try await ffiClient.applySignatureRequest(
+				signatureRequest: signatureRequest)
+		} catch {
+			throw ClientError.creationError(
+				"Failed to sign the message: \(error.localizedDescription)")
+		}
+	}
+
 	public func canMessage(address: String) async throws -> Bool {
 		let canMessage = try await ffiClient.canMessage(accountAddresses: [
 			address
@@ -320,7 +374,7 @@ public final class Client {
 	public func inboxIdFromAddress(address: String) async throws -> String? {
 		return try await ffiClient.findInboxId(address: address.lowercased())
 	}
-	
+
 	public func signWithInstallationKey(message: String) throws -> Data {
 		return try ffiClient.signWithInstallationKey(text: message)
 	}
@@ -404,22 +458,6 @@ public final class Client {
 		try await ffiClient.sendSyncRequest(kind: .consent)
 	}
 
-	public func revokeAllOtherInstallations(signingKey: SigningKey) async throws
-	{
-		let signatureRequest = try await ffiClient.revokeAllOtherInstallations()
-		do {
-			let signedData = try await signingKey.sign(
-				message: signatureRequest.signatureText())
-			try await signatureRequest.addEcdsaSignature(
-				signatureBytes: signedData.rawData)
-			try await ffiClient.applySignatureRequest(
-				signatureRequest: signatureRequest)
-		} catch {
-			throw ClientError.creationError(
-				"Failed to sign the message: \(error.localizedDescription)")
-		}
-	}
-
 	public func inboxState(refreshFromNetwork: Bool) async throws -> InboxState
 	{
 		return InboxState(

diff --git a/Tests/XMTPTests/ClientTests.swift b/Tests/XMTPTests/ClientTests.swift
@@ -326,4 +326,76 @@ class ClientTests: XCTestCase {
 			states.last!.recoveryAddress.lowercased(),
 			fixtures.caro.walletAddress.lowercased())
 	}
+
+	func testAddAccounts() async throws {
+		let fixtures = try await fixtures()
+		let alix2Wallet = try PrivateKey.generate()
+		let alix3Wallet = try PrivateKey.generate()
+
+		try await fixtures.alixClient.addAccount(
+			recoveryAccount: fixtures.alix, newAccount: alix2Wallet)
+		try await fixtures.alixClient.addAccount(
+			recoveryAccount: fixtures.alix, newAccount: alix3Wallet)
+
+		let state = try await fixtures.alixClient.inboxState(
+			refreshFromNetwork: true)
+		XCTAssertEqual(state.installations.count, 1)
+		XCTAssertEqual(state.addresses.count, 3)
+		XCTAssertEqual(
+			state.recoveryAddress.lowercased(),
+			fixtures.alixClient.address.lowercased())
+		XCTAssertEqual(
+			state.addresses.sorted(),
+			[
+				alix2Wallet.address.lowercased(),
+				alix3Wallet.address.lowercased(),
+				fixtures.alixClient.address.lowercased(),
+			].sorted()
+		)
+	}
+
+	func testRemovingAccounts() async throws {
+		let fixtures = try await fixtures()
+		let alix2Wallet = try PrivateKey.generate()
+		let alix3Wallet = try PrivateKey.generate()
+
+		try await fixtures.alixClient.addAccount(
+			recoveryAccount: fixtures.alix, newAccount: alix2Wallet)
+		try await fixtures.alixClient.addAccount(
+			recoveryAccount: fixtures.alix, newAccount: alix3Wallet)
+
+		var state = try await fixtures.alixClient.inboxState(
+			refreshFromNetwork: true)
+		XCTAssertEqual(state.addresses.count, 3)
+		XCTAssertEqual(
+			state.recoveryAddress.lowercased(),
+			fixtures.alixClient.address.lowercased())
+
+		try await fixtures.alixClient.removeAccount(
+			recoveryAccount: fixtures.alix, addressToRemove: alix2Wallet.address
+		)
+
+		state = try await fixtures.alixClient.inboxState(
+			refreshFromNetwork: true)
+		XCTAssertEqual(state.addresses.count, 2)
+		XCTAssertEqual(
+			state.recoveryAddress.lowercased(),
+			fixtures.alixClient.address.lowercased())
+		XCTAssertEqual(
+			state.addresses.sorted(),
+			[
+				alix3Wallet.address.lowercased(),
+				fixtures.alixClient.address.lowercased(),
+			].sorted()
+		)
+		XCTAssertEqual(state.installations.count, 1)
+
+		// Cannot remove the recovery address
+		await assertThrowsAsyncError(
+			try await fixtures.alixClient.removeAccount(
+				recoveryAccount: alix3Wallet,
+				addressToRemove: fixtures.alixClient.address
+			))
+	}
+
 }
diff --git a/XMTP.podspec b/XMTP.podspec
@@ -1,6 +1,6 @@
 Pod::Spec.new do |spec|
   spec.name         = "XMTP"
-  spec.version      = "3.0.6"
+  spec.version      = "3.0.7"
   spec.summary      = "XMTP SDK Cocoapod"
 
   spec.description  = <<-DESC