Existing Sharded -> Enclave Migration Flow (#74)

Thirdweb.Console/Program.cs

@@ -29,9 +29,9 @@
 
 #region Contract Interaction
 
-var contract = await ThirdwebContract.Create(client: client, address: "0xbc4ca0eda7647a8ab7c2061c2e118a18a936f13d", chain: 1);
-var nfts = await contract.ERC721_GetAllNFTs();
-Console.WriteLine($"NFTs: {JsonConvert.SerializeObject(nfts, Formatting.Indented)}");
+// var contract = await ThirdwebContract.Create(client: client, address: "0xbc4ca0eda7647a8ab7c2061c2e118a18a936f13d", chain: 1);
+// var nfts = await contract.ERC721_GetAllNFTs();
+// Console.WriteLine($"NFTs: {JsonConvert.SerializeObject(nfts, Formatting.Indented)}");
 
 #endregion
 
@@ -74,14 +74,34 @@
 
 #region Ecosystem Wallet
 
-// var ecosystemWallet = await EcosystemWallet.Create(client: client, ecosystemId: "ecosystem.the-bonfire", email: "[email protected]");
+// var inAppWallet = await InAppWallet.Create(client: client, authProvider: AuthProvider.Google);
+// if (!await inAppWallet.IsConnected())
+// {
+//     _ = await inAppWallet.LoginWithOauth(
+//         isMobile: false,
+//         (url) =>
+//         {
+//             var psi = new ProcessStartInfo { FileName = url, UseShellExecute = true };
+//             _ = Process.Start(psi);
+//         },
+//         "thirdweb://",
+//         new InAppWalletBrowser()
+//     );
+// }
+
+// var ecosystemWallet = await EcosystemWallet.Create(
+//     client: client,
+//     ecosystemId: "ecosystem.the-bonfire",
+//     ecosystemPartnerId: "20842d97-be35-4ecc-b51e-9f3ba0843a60",
+//     email: "[email protected]"
+// );
 
 // if (!await ecosystemWallet.IsConnected())
 // {
 //     _ = await ecosystemWallet.SendOTP();
 //     Console.WriteLine("Enter OTP:");
 //     var otp = Console.ReadLine();
-//     _ = await ecosystemWallet.LoginWithOtp(otp: otp);
+//     _ = await ecosystemWallet.LoginWithOtp(otp);
 // }
 // var ecosystemWalletAddress = await ecosystemWallet.GetAddress();
 // Console.WriteLine($"Ecosystem Wallet address: {ecosystemWalletAddress}");

Thirdweb/Thirdweb.Wallets/InAppWallet/EcosystemWallet/EcosystemWallet.cs

@@ -21,8 +21,9 @@ public partial class EcosystemWallet : PrivateKeyWallet
 
     private string _address;
 
-    private const string EMBEDDED_WALLET_PATH_2024 = "https://embedded-wallet.thirdweb.com/api/2024-05-05";
-    private const string EMBEDDED_WALLET_PATH_V1 = "https://embedded-wallet.thirdweb.com/api/v1";
+    private const string EMBEDDED_WALLET_BASE_PATH = "https://embedded-wallet.thirdweb.com/api";
+    private const string EMBEDDED_WALLET_PATH_2024 = $"{EMBEDDED_WALLET_BASE_PATH}/2024-05-05";
+    private const string EMBEDDED_WALLET_PATH_V1 = $"{EMBEDDED_WALLET_BASE_PATH}/v1";
     private const string ENCLAVE_PATH = $"{EMBEDDED_WALLET_PATH_V1}/enclave-wallet";
 
     private EcosystemWallet(ThirdwebClient client, EmbeddedWallet embeddedWallet, IThirdwebHttpClient httpClient, string email, string phoneNumber, string authProvider, IThirdwebWallet siweSigner)
@@ -146,8 +147,8 @@ private static async Task<string> ResumeEnclaveSession(IThirdwebHttpClient httpC
         }
         else
         {
-            // TODO: Implement migration flow from existing sharded InAppWallet to sharded EcosystemWallet to enclave Ecosystem Wallet
-            throw new InvalidOperationException("Migration flow from existing sharded InAppWallet to enclave Ecosystem Wallet not implemented yet.");
+            await embeddedWallet.SignOutAsync().ConfigureAwait(false);
+            throw new InvalidOperationException("Must auth again to perform migration.");
         }
     }
 
@@ -181,6 +182,7 @@ private static async Task<string> GenerateWallet(IThirdwebHttpClient httpClient)
     private async Task<string> PostAuth(Server.VerifyResult result)
     {
         this._httpClient.AddHeader("Authorization", $"Bearer embedded-wallet-token:{result.AuthToken}");
+
         string address;
         if (result.IsNewUser)
         {
@@ -195,8 +197,7 @@ private async Task<string> PostAuth(Server.VerifyResult result)
             }
             else
             {
-                // TODO: Implement migration flow from existing sharded InAppWallet to sharded EcosystemWallet to enclave Ecosystem Wallet
-                throw new InvalidOperationException("Migration flow from existing sharded InAppWallet to enclave Ecosystem Wallet not implemented yet.");
+                address = await this.MigrateShardToEnclave(result).ConfigureAwait(false);
             }
         }
 
@@ -212,6 +213,28 @@ private async Task<string> PostAuth(Server.VerifyResult result)
         }
     }
 
+    private async Task<string> MigrateShardToEnclave(Server.VerifyResult authResult)
+    {
+        // TODO: For recovery code, allow old encryption keys as overrides to migrate sharded custom auth?
+        var (address, encryptedPrivateKeyB64, ivB64, kmsCiphertextB64) = await this._embeddedWallet.GenerateEncryptionDataAsync(authResult.AuthToken, authResult.RecoveryCode).ConfigureAwait(false);
+
+        var url = $"{ENCLAVE_PATH}/migrate";
+        var payload = new
+        {
+            address,
+            encryptedPrivateKeyB64,
+            ivB64,
+            kmsCiphertextB64
+        };
+        var requestContent = new StringContent(JsonConvert.SerializeObject(payload), Encoding.UTF8, "application/json");
+
+        var response = await this._httpClient.PostAsync(url, requestContent).ConfigureAwait(false);
+        _ = response.EnsureSuccessStatusCode();
+
+        var userStatus = await GetUserStatus(this._httpClient).ConfigureAwait(false);
+        return userStatus.Wallets[0].Address;
+    }
+
     #endregion
 
     #region Wallet Specific
@@ -526,7 +549,6 @@ public async Task<string> LoginWithSiwe(BigInteger chainId)
         {
             sessionId = Guid.NewGuid().ToString();
         }
-        Console.WriteLine($"Guest Session ID: {sessionId}");
         var serverRes = await this._embeddedWallet.SignInWithGuestAsync(sessionId).ConfigureAwait(false);
         return serverRes;
     }

Thirdweb/Thirdweb.Wallets/InAppWallet/EmbeddedWallet.Authentication/AWS.cs

@@ -1,6 +1,7 @@
 using System.Security.Cryptography;
 using System.Text;
 using Newtonsoft.Json;
+using Newtonsoft.Json.Linq;
 
 namespace Thirdweb.EWS;
 
@@ -9,13 +10,20 @@ internal class AWS
     private const string AWS_REGION = "us-west-2";
 
     private static readonly string _recoverySharePasswordLambdaFunctionNameV2 = $"arn:aws:lambda:{AWS_REGION}:324457261097:function:lambda-thirdweb-auth-enc-key-prod-ThirdwebAuthEncKeyFunction";
+    private static readonly string _migrationKeyId = $"arn:aws:kms:{AWS_REGION}:324457261097:key/ccfb9ecd-f45d-4f37-864a-25fe72dcb49e";
 
     internal static async Task<MemoryStream> InvokeRecoverySharePasswordLambdaAsync(string identityId, string token, string invokePayload, Type thirdwebHttpClientType)
     {
         var credentials = await GetTemporaryCredentialsAsync(identityId, token, thirdwebHttpClientType).ConfigureAwait(false);
         return await InvokeLambdaWithTemporaryCredentialsAsync(credentials, invokePayload, thirdwebHttpClientType, _recoverySharePasswordLambdaFunctionNameV2).ConfigureAwait(false);
     }
 
+    internal static async Task<JToken> GenerateDataKey(string identityId, string token, Type thirdwebHttpClientType)
+    {
+        var credentials = await GetTemporaryCredentialsAsync(identityId, token, thirdwebHttpClientType).ConfigureAwait(false);
+        return await GenerateDataKey(credentials, thirdwebHttpClientType).ConfigureAwait(false);
+    }
+
     private static async Task<AwsCredentials> GetTemporaryCredentialsAsync(string identityId, string token, Type thirdwebHttpClientType)
     {
         var client = thirdwebHttpClientType.GetConstructor(Type.EmptyTypes).Invoke(null) as IThirdwebHttpClient;
@@ -45,6 +53,72 @@ private static async Task<AwsCredentials> GetTemporaryCredentialsAsync(string id
         };
     }
 
+    private static async Task<JToken> GenerateDataKey(AwsCredentials credentials, Type thirdwebHttpClientType)
+    {
+        var client = thirdwebHttpClientType.GetConstructor(Type.EmptyTypes).Invoke(null) as IThirdwebHttpClient;
+        var endpoint = $"https://kms.{AWS_REGION}.amazonaws.com/";
+
+        var payloadForGenerateDataKey = new { KeyId = _migrationKeyId, KeySpec = "AES_256" };
+
+        var content = new StringContent(JsonConvert.SerializeObject(payloadForGenerateDataKey), Encoding.UTF8, "application/x-amz-json-1.1");
+
+        client.AddHeader("X-Amz-Target", "TrentService.GenerateDataKey");
+
+        var dateTimeNow = DateTime.UtcNow;
+        var dateStamp = dateTimeNow.ToString("yyyyMMdd");
+        var amzDate = dateTimeNow.ToString("yyyyMMddTHHmmssZ");
+        var canonicalUri = "/";
+
+        var canonicalHeaders = $"host:kms.{AWS_REGION}.amazonaws.com\nx-amz-date:{amzDate}\n";
+        var signedHeaders = "host;x-amz-date";
+
+#if NETSTANDARD
+        using var sha256 = SHA256.Create();
+        var payloadHash = ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
+#else
+        var payloadHash = ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
+#endif
+
+        var canonicalRequest = $"POST\n{canonicalUri}\n\n{canonicalHeaders}\n{signedHeaders}\n{payloadHash}";
+
+        var algorithm = "AWS4-HMAC-SHA256";
+        var credentialScope = $"{dateStamp}/{AWS_REGION}/kms/aws4_request";
+
+#if NETSTANDARD
+        var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(canonicalRequest)))}";
+#else
+        var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(canonicalRequest)))}";
+#endif
+
+        var signingKey = GetSignatureKey(credentials.SecretAccessKey, dateStamp, AWS_REGION, "kms");
+        var signature = ToHexString(HMACSHA256(signingKey, stringToSign));
+
+        var authorizationHeader = $"{algorithm} Credential={credentials.AccessKeyId}/{credentialScope}, SignedHeaders={signedHeaders}, Signature={signature}";
+
+        client.AddHeader("x-amz-date", amzDate);
+        client.AddHeader("Authorization", authorizationHeader);
+        client.AddHeader("x-amz-security-token", credentials.SessionToken);
+
+        var response = await client.PostAsync(endpoint, content).ConfigureAwait(false);
+        var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
+
+        if (!response.IsSuccessStatusCode)
+        {
+            throw new Exception($"Failed to generate data key: {responseContent}");
+        }
+
+        var responseObject = JToken.Parse(responseContent);
+        var plaintextKeyBlob = responseObject["Plaintext"];
+        var cipherTextBlob = responseObject["CiphertextBlob"];
+
+        if (plaintextKeyBlob == null || cipherTextBlob == null)
+        {
+            throw new Exception("No migration key found. Please try again.");
+        }
+
+        return responseObject;
+    }
+
     private static async Task<MemoryStream> InvokeLambdaWithTemporaryCredentialsAsync(AwsCredentials credentials, string invokePayload, Type thirdwebHttpClientType, string lambdaFunction)
     {
         var endpoint = $"https://lambda.{AWS_REGION}.amazonaws.com/2015-03-31/functions/{lambdaFunction}/invocations";

Thirdweb/Thirdweb.Wallets/InAppWallet/EmbeddedWallet.Authentication/Server.Types.cs

@@ -148,7 +148,7 @@ internal class UserWallet
     }
 
     [DataContract]
-    private class IdTokenResponse
+    internal class IdTokenResponse
     {
         [DataMember(Name = "token")]
         internal string Token { get; set; }

Thirdweb/Thirdweb.Wallets/InAppWallet/EmbeddedWallet.Authentication/Server.cs

@@ -1,5 +1,6 @@
 using System.Net.Http.Headers;
 using Newtonsoft.Json;
+using Newtonsoft.Json.Linq;
 
 namespace Thirdweb.EWS;
 
@@ -31,6 +32,8 @@ internal abstract class ServerBase
     internal abstract Task<Server.VerifyResult> VerifyOAuthAsync(string authResultStr);
 
     internal abstract Task<Server.VerifyResult> VerifyAuthEndpointAsync(string payload);
+
+    internal abstract Task<JToken> GenerateEncryptedKeyResultAsync(string authToken);
 }
 
 internal partial class Server : ServerBase
@@ -302,6 +305,12 @@ internal override async Task<VerifyResult> VerifyOAuthAsync(string authResultStr
 
     #region Misc
 
+    internal override async Task<JToken> GenerateEncryptedKeyResultAsync(string authToken)
+    {
+        var webExchangeResult = await this.FetchCognitoIdTokenAsync(authToken).ConfigureAwait(false);
+        return await AWS.GenerateDataKey(webExchangeResult.IdentityId, webExchangeResult.Token, _thirdwebHttpClientType).ConfigureAwait(false);
+    }
+
     private async Task<VerifyResult> InvokeAuthResultLambdaAsync(AuthResultType authResult)
     {
         var authToken = authResult.StoredToken.CookieString;

Thirdweb/Thirdweb.Wallets/InAppWallet/EmbeddedWallet/EmbeddedWallet.Misc.cs

@@ -1,3 +1,4 @@
+using System.Security.Cryptography;
 using Nethereum.Web3.Accounts;
 
 namespace Thirdweb.EWS;
@@ -101,7 +102,7 @@ private User MakeUserAsync(string emailAddress, string phoneNumber, Account acco
         return (account, deviceShare);
     }
 
-    private async Task<(Account account, string deviceShare)> RecoverAccountAsync(string authToken, string recoveryCode)
+    internal async Task<(Account account, string deviceShare)> RecoverAccountAsync(string authToken, string recoveryCode)
     {
         (var authShare, var encryptedRecoveryShare) = await this._server.FetchAuthAndRecoverySharesAsync(authToken).ConfigureAwait(false);
 
@@ -113,6 +114,60 @@ private User MakeUserAsync(string emailAddress, string phoneNumber, Account acco
         return (account, deviceShare);
     }
 
+    internal async Task<(string address, string encryptedPrivateKeyB64, string ivB64, string kmsCiphertextB64)> GenerateEncryptionDataAsync(string authToken, string recoveryCode)
+    {
+        var (account, _) = await this.RecoverAccountAsync(authToken, recoveryCode).ConfigureAwait(false);
+        var address = account.Address;
+
+        var encryptedKeyResult = await this._server.GenerateEncryptedKeyResultAsync(authToken).ConfigureAwait(false);
+
+        var plainTextBase64 = encryptedKeyResult["Plaintext"]?.ToString();
+        var cipherTextBlobBase64 = encryptedKeyResult["CiphertextBlob"]?.ToString();
+
+        if (string.IsNullOrEmpty(plainTextBase64) || string.IsNullOrEmpty(cipherTextBlobBase64))
+        {
+            throw new InvalidOperationException("No migration key found. Please try again.");
+        }
+
+        var iv = new byte[16];
+        using (var rng = RandomNumberGenerator.Create())
+        {
+            rng.GetBytes(iv);
+        }
+
+        var privateKey = account.PrivateKey;
+        var utf8WithoutBom = new System.Text.UTF8Encoding(encoderShouldEmitUTF8Identifier: true);
+        var privateKeyBytes = utf8WithoutBom.GetBytes(privateKey);
+
+        byte[] encryptedPrivateKeyBytes;
+        try
+        {
+            using var aes = Aes.Create();
+            aes.KeySize = 256;
+            aes.BlockSize = 128;
+            aes.Key = Convert.FromBase64String(plainTextBase64);
+            aes.IV = iv;
+            aes.Mode = CipherMode.CBC;
+            aes.Padding = PaddingMode.PKCS7;
+
+            using var encryptor = aes.CreateEncryptor();
+            encryptedPrivateKeyBytes = encryptor.TransformFinalBlock(privateKeyBytes, 0, privateKeyBytes.Length);
+        }
+        catch (Exception ex)
+        {
+            throw new InvalidOperationException("Encryption failed.", ex);
+        }
+
+        var encryptedData = new byte[iv.Length + encryptedPrivateKeyBytes.Length];
+        iv.CopyTo(encryptedData, 0);
+        encryptedPrivateKeyBytes.CopyTo(encryptedData, iv.Length);
+
+        var encryptedDataB64 = Convert.ToBase64String(encryptedData);
+        var ivB64 = Convert.ToBase64String(iv);
+
+        return (address, encryptedDataB64, ivB64, cipherTextBlobBase64);
+    }
+
     public class VerifyResult
     {
         public User User { get; }