Skip to content

Commit 791be46

Browse files
0xFirekeeper0xFirekeeper
authored
Existing Sharded -> Enclave Migration Flow (#74)
1 parent 36b4c2c commit 791be46

File tree

6 files changed

+194
-14
lines changed

6 files changed

+194
-14
lines changed

‎Thirdweb.Console/Program.cs

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@
2929

3030
#region Contract Interaction
3131

32-
var contract = await ThirdwebContract.Create(client: client, address: "0xbc4ca0eda7647a8ab7c2061c2e118a18a936f13d", chain: 1);
33-
var nfts = await contract.ERC721_GetAllNFTs();
34-
Console.WriteLine($"NFTs: {JsonConvert.SerializeObject(nfts, Formatting.Indented)}");
32+
// var contract = await ThirdwebContract.Create(client: client, address: "0xbc4ca0eda7647a8ab7c2061c2e118a18a936f13d", chain: 1);
33+
// var nfts = await contract.ERC721_GetAllNFTs();
34+
// Console.WriteLine($"NFTs: {JsonConvert.SerializeObject(nfts, Formatting.Indented)}");
3535

3636
#endregion
3737

@@ -74,14 +74,34 @@
7474

7575
#region Ecosystem Wallet
7676

77-
// var ecosystemWallet = await EcosystemWallet.Create(client: client, ecosystemId: "ecosystem.the-bonfire", email: "firekeeper+linkeco@thirdweb.com");
77+
// var inAppWallet = await InAppWallet.Create(client: client, authProvider: AuthProvider.Google);
78+
// if (!await inAppWallet.IsConnected())
79+
// {
80+
// _ = await inAppWallet.LoginWithOauth(
81+
// isMobile: false,
82+
// (url) =>
83+
// {
84+
// var psi = new ProcessStartInfo { FileName = url, UseShellExecute = true };
85+
// _ = Process.Start(psi);
86+
// },
87+
// "thirdweb://",
88+
// new InAppWalletBrowser()
89+
// );
90+
// }
91+
92+
// var ecosystemWallet = await EcosystemWallet.Create(
93+
// client: client,
94+
// ecosystemId: "ecosystem.the-bonfire",
95+
// ecosystemPartnerId: "20842d97-be35-4ecc-b51e-9f3ba0843a60",
96+
// email: "firekeeper+shardedsucks@thirdweb.com"
97+
// );
7898

7999
// if (!await ecosystemWallet.IsConnected())
80100
// {
81101
// _ = await ecosystemWallet.SendOTP();
82102
// Console.WriteLine("Enter OTP:");
83103
// var otp = Console.ReadLine();
84-
// _ = await ecosystemWallet.LoginWithOtp(otp: otp);
104+
// _ = await ecosystemWallet.LoginWithOtp(otp);
85105
// }
86106
// var ecosystemWalletAddress = await ecosystemWallet.GetAddress();
87107
// Console.WriteLine($"Ecosystem Wallet address: {ecosystemWalletAddress}");

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

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@ public partial class EcosystemWallet : PrivateKeyWallet
2121

2222
private string _address;
2323

24-
private const string EMBEDDED_WALLET_PATH_2024 = "https://embedded-wallet.thirdweb.com/api/2024-05-05";
25-
private const string EMBEDDED_WALLET_PATH_V1 = "https://embedded-wallet.thirdweb.com/api/v1";
24+
private const string EMBEDDED_WALLET_BASE_PATH = "https://embedded-wallet.thirdweb.com/api";
25+
private const string EMBEDDED_WALLET_PATH_2024 = $"{EMBEDDED_WALLET_BASE_PATH}/2024-05-05";
26+
private const string EMBEDDED_WALLET_PATH_V1 = $"{EMBEDDED_WALLET_BASE_PATH}/v1";
2627
private const string ENCLAVE_PATH = $"{EMBEDDED_WALLET_PATH_V1}/enclave-wallet";
2728

2829
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
146147
}
147148
else
148149
{
149-
// TODO: Implement migration flow from existing sharded InAppWallet to sharded EcosystemWallet to enclave Ecosystem Wallet
150-
throw new InvalidOperationException("Migration flow from existing sharded InAppWallet to enclave Ecosystem Wallet not implemented yet.");
150+
await embeddedWallet.SignOutAsync().ConfigureAwait(false);
151+
throw new InvalidOperationException("Must auth again to perform migration.");
151152
}
152153
}
153154

@@ -181,6 +182,7 @@ private static async Task<string> GenerateWallet(IThirdwebHttpClient httpClient)
181182
private async Task<string> PostAuth(Server.VerifyResult result)
182183
{
183184
this._httpClient.AddHeader("Authorization", $"Bearer embedded-wallet-token:{result.AuthToken}");
185+
184186
string address;
185187
if (result.IsNewUser)
186188
{
@@ -195,8 +197,7 @@ private async Task<string> PostAuth(Server.VerifyResult result)
195197
}
196198
else
197199
{
198-
// TODO: Implement migration flow from existing sharded InAppWallet to sharded EcosystemWallet to enclave Ecosystem Wallet
199-
throw new InvalidOperationException("Migration flow from existing sharded InAppWallet to enclave Ecosystem Wallet not implemented yet.");
200+
address = await this.MigrateShardToEnclave(result).ConfigureAwait(false);
200201
}
201202
}
202203

@@ -212,6 +213,28 @@ private async Task<string> PostAuth(Server.VerifyResult result)
212213
}
213214
}
214215

216+
private async Task<string> MigrateShardToEnclave(Server.VerifyResult authResult)
217+
{
218+
// TODO: For recovery code, allow old encryption keys as overrides to migrate sharded custom auth?
219+
var (address, encryptedPrivateKeyB64, ivB64, kmsCiphertextB64) = await this._embeddedWallet.GenerateEncryptionDataAsync(authResult.AuthToken, authResult.RecoveryCode).ConfigureAwait(false);
220+
221+
var url = $"{ENCLAVE_PATH}/migrate";
222+
var payload = new
223+
{
224+
address,
225+
encryptedPrivateKeyB64,
226+
ivB64,
227+
kmsCiphertextB64
228+
};
229+
var requestContent = new StringContent(JsonConvert.SerializeObject(payload), Encoding.UTF8, "application/json");
230+
231+
var response = await this._httpClient.PostAsync(url, requestContent).ConfigureAwait(false);
232+
_ = response.EnsureSuccessStatusCode();
233+
234+
var userStatus = await GetUserStatus(this._httpClient).ConfigureAwait(false);
235+
return userStatus.Wallets[0].Address;
236+
}
237+
215238
#endregion
216239

217240
#region Wallet Specific
@@ -526,7 +549,6 @@ public async Task<string> LoginWithSiwe(BigInteger chainId)
526549
{
527550
sessionId = Guid.NewGuid().ToString();
528551
}
529-
Console.WriteLine($"Guest Session ID: {sessionId}");
530552
var serverRes = await this._embeddedWallet.SignInWithGuestAsync(sessionId).ConfigureAwait(false);
531553
return serverRes;
532554
}

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

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using System.Security.Cryptography;
22
using System.Text;
33
using Newtonsoft.Json;
4+
using Newtonsoft.Json.Linq;
45

56
namespace Thirdweb.EWS;
67

@@ -9,13 +10,20 @@ internal class AWS
910
private const string AWS_REGION = "us-west-2";
1011

1112
private static readonly string _recoverySharePasswordLambdaFunctionNameV2 = $"arn:aws:lambda:{AWS_REGION}:324457261097:function:lambda-thirdweb-auth-enc-key-prod-ThirdwebAuthEncKeyFunction";
13+
private static readonly string _migrationKeyId = $"arn:aws:kms:{AWS_REGION}:324457261097:key/ccfb9ecd-f45d-4f37-864a-25fe72dcb49e";
1214

1315
internal static async Task<MemoryStream> InvokeRecoverySharePasswordLambdaAsync(string identityId, string token, string invokePayload, Type thirdwebHttpClientType)
1416
{
1517
var credentials = await GetTemporaryCredentialsAsync(identityId, token, thirdwebHttpClientType).ConfigureAwait(false);
1618
return await InvokeLambdaWithTemporaryCredentialsAsync(credentials, invokePayload, thirdwebHttpClientType, _recoverySharePasswordLambdaFunctionNameV2).ConfigureAwait(false);
1719
}
1820

21+
internal static async Task<JToken> GenerateDataKey(string identityId, string token, Type thirdwebHttpClientType)
22+
{
23+
var credentials = await GetTemporaryCredentialsAsync(identityId, token, thirdwebHttpClientType).ConfigureAwait(false);
24+
return await GenerateDataKey(credentials, thirdwebHttpClientType).ConfigureAwait(false);
25+
}
26+
1927
private static async Task<AwsCredentials> GetTemporaryCredentialsAsync(string identityId, string token, Type thirdwebHttpClientType)
2028
{
2129
var client = thirdwebHttpClientType.GetConstructor(Type.EmptyTypes).Invoke(null) as IThirdwebHttpClient;
@@ -45,6 +53,72 @@ private static async Task<AwsCredentials> GetTemporaryCredentialsAsync(string id
4553
};
4654
}
4755

56+
private static async Task<JToken> GenerateDataKey(AwsCredentials credentials, Type thirdwebHttpClientType)
57+
{
58+
var client = thirdwebHttpClientType.GetConstructor(Type.EmptyTypes).Invoke(null) as IThirdwebHttpClient;
59+
var endpoint = $"https://kms.{AWS_REGION}.amazonaws.com/";
60+
61+
var payloadForGenerateDataKey = new { KeyId = _migrationKeyId, KeySpec = "AES_256" };
62+
63+
var content = new StringContent(JsonConvert.SerializeObject(payloadForGenerateDataKey), Encoding.UTF8, "application/x-amz-json-1.1");
64+
65+
client.AddHeader("X-Amz-Target", "TrentService.GenerateDataKey");
66+
67+
var dateTimeNow = DateTime.UtcNow;
68+
var dateStamp = dateTimeNow.ToString("yyyyMMdd");
69+
var amzDate = dateTimeNow.ToString("yyyyMMddTHHmmssZ");
70+
var canonicalUri = "/";
71+
72+
var canonicalHeaders = $"host:kms.{AWS_REGION}.amazonaws.com\nx-amz-date:{amzDate}\n";
73+
var signedHeaders = "host;x-amz-date";
74+
75+
#if NETSTANDARD
76+
using var sha256 = SHA256.Create();
77+
var payloadHash = ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
78+
#else
79+
var payloadHash = ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
80+
#endif
81+
82+
var canonicalRequest = $"POST\n{canonicalUri}\n\n{canonicalHeaders}\n{signedHeaders}\n{payloadHash}";
83+
84+
var algorithm = "AWS4-HMAC-SHA256";
85+
var credentialScope = $"{dateStamp}/{AWS_REGION}/kms/aws4_request";
86+
87+
#if NETSTANDARD
88+
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(canonicalRequest)))}";
89+
#else
90+
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(canonicalRequest)))}";
91+
#endif
92+
93+
var signingKey = GetSignatureKey(credentials.SecretAccessKey, dateStamp, AWS_REGION, "kms");
94+
var signature = ToHexString(HMACSHA256(signingKey, stringToSign));
95+
96+
var authorizationHeader = $"{algorithm} Credential={credentials.AccessKeyId}/{credentialScope}, SignedHeaders={signedHeaders}, Signature={signature}";
97+
98+
client.AddHeader("x-amz-date", amzDate);
99+
client.AddHeader("Authorization", authorizationHeader);
100+
client.AddHeader("x-amz-security-token", credentials.SessionToken);
101+
102+
var response = await client.PostAsync(endpoint, content).ConfigureAwait(false);
103+
var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
104+
105+
if (!response.IsSuccessStatusCode)
106+
{
107+
throw new Exception($"Failed to generate data key: {responseContent}");
108+
}
109+
110+
var responseObject = JToken.Parse(responseContent);
111+
var plaintextKeyBlob = responseObject["Plaintext"];
112+
var cipherTextBlob = responseObject["CiphertextBlob"];
113+
114+
if (plaintextKeyBlob == null || cipherTextBlob == null)
115+
{
116+
throw new Exception("No migration key found. Please try again.");
117+
}
118+
119+
return responseObject;
120+
}
121+
48122
private static async Task<MemoryStream> InvokeLambdaWithTemporaryCredentialsAsync(AwsCredentials credentials, string invokePayload, Type thirdwebHttpClientType, string lambdaFunction)
49123
{
50124
var endpoint = $"https://lambda.{AWS_REGION}.amazonaws.com/2015-03-31/functions/{lambdaFunction}/invocations";

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ internal class UserWallet
148148
}
149149

150150
[DataContract]
151-
private class IdTokenResponse
151+
internal class IdTokenResponse
152152
{
153153
[DataMember(Name = "token")]
154154
internal string Token { get; set; }

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System.Net.Http.Headers;
22
using Newtonsoft.Json;
3+
using Newtonsoft.Json.Linq;
34

45
namespace Thirdweb.EWS;
56

@@ -31,6 +32,8 @@ internal abstract class ServerBase
3132
internal abstract Task<Server.VerifyResult> VerifyOAuthAsync(string authResultStr);
3233

3334
internal abstract Task<Server.VerifyResult> VerifyAuthEndpointAsync(string payload);
35+
36+
internal abstract Task<JToken> GenerateEncryptedKeyResultAsync(string authToken);
3437
}
3538

3639
internal partial class Server : ServerBase
@@ -302,6 +305,12 @@ internal override async Task<VerifyResult> VerifyOAuthAsync(string authResultStr
302305

303306
#region Misc
304307

308+
internal override async Task<JToken> GenerateEncryptedKeyResultAsync(string authToken)
309+
{
310+
var webExchangeResult = await this.FetchCognitoIdTokenAsync(authToken).ConfigureAwait(false);
311+
return await AWS.GenerateDataKey(webExchangeResult.IdentityId, webExchangeResult.Token, _thirdwebHttpClientType).ConfigureAwait(false);
312+
}
313+
305314
private async Task<VerifyResult> InvokeAuthResultLambdaAsync(AuthResultType authResult)
306315
{
307316
var authToken = authResult.StoredToken.CookieString;

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

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Security.Cryptography;
12
using Nethereum.Web3.Accounts;
23

34
namespace Thirdweb.EWS;
@@ -101,7 +102,7 @@ private User MakeUserAsync(string emailAddress, string phoneNumber, Account acco
101102
return (account, deviceShare);
102103
}
103104

104-
private async Task<(Account account, string deviceShare)> RecoverAccountAsync(string authToken, string recoveryCode)
105+
internal async Task<(Account account, string deviceShare)> RecoverAccountAsync(string authToken, string recoveryCode)
105106
{
106107
(var authShare, var encryptedRecoveryShare) = await this._server.FetchAuthAndRecoverySharesAsync(authToken).ConfigureAwait(false);
107108

@@ -113,6 +114,60 @@ private User MakeUserAsync(string emailAddress, string phoneNumber, Account acco
113114
return (account, deviceShare);
114115
}
115116

117+
internal async Task<(string address, string encryptedPrivateKeyB64, string ivB64, string kmsCiphertextB64)> GenerateEncryptionDataAsync(string authToken, string recoveryCode)
118+
{
119+
var (account, _) = await this.RecoverAccountAsync(authToken, recoveryCode).ConfigureAwait(false);
120+
var address = account.Address;
121+
122+
var encryptedKeyResult = await this._server.GenerateEncryptedKeyResultAsync(authToken).ConfigureAwait(false);
123+
124+
var plainTextBase64 = encryptedKeyResult["Plaintext"]?.ToString();
125+
var cipherTextBlobBase64 = encryptedKeyResult["CiphertextBlob"]?.ToString();
126+
127+
if (string.IsNullOrEmpty(plainTextBase64) || string.IsNullOrEmpty(cipherTextBlobBase64))
128+
{
129+
throw new InvalidOperationException("No migration key found. Please try again.");
130+
}
131+
132+
var iv = new byte[16];
133+
using (var rng = RandomNumberGenerator.Create())
134+
{
135+
rng.GetBytes(iv);
136+
}
137+
138+
var privateKey = account.PrivateKey;
139+
var utf8WithoutBom = new System.Text.UTF8Encoding(encoderShouldEmitUTF8Identifier: true);
140+
var privateKeyBytes = utf8WithoutBom.GetBytes(privateKey);
141+
142+
byte[] encryptedPrivateKeyBytes;
143+
try
144+
{
145+
using var aes = Aes.Create();
146+
aes.KeySize = 256;
147+
aes.BlockSize = 128;
148+
aes.Key = Convert.FromBase64String(plainTextBase64);
149+
aes.IV = iv;
150+
aes.Mode = CipherMode.CBC;
151+
aes.Padding = PaddingMode.PKCS7;
152+
153+
using var encryptor = aes.CreateEncryptor();
154+
encryptedPrivateKeyBytes = encryptor.TransformFinalBlock(privateKeyBytes, 0, privateKeyBytes.Length);
155+
}
156+
catch (Exception ex)
157+
{
158+
throw new InvalidOperationException("Encryption failed.", ex);
159+
}
160+
161+
var encryptedData = new byte[iv.Length + encryptedPrivateKeyBytes.Length];
162+
iv.CopyTo(encryptedData, 0);
163+
encryptedPrivateKeyBytes.CopyTo(encryptedData, iv.Length);
164+
165+
var encryptedDataB64 = Convert.ToBase64String(encryptedData);
166+
var ivB64 = Convert.ToBase64String(iv);
167+
168+
return (address, encryptedDataB64, ivB64, cipherTextBlobBase64);
169+
}
170+
116171
public class VerifyResult
117172
{
118173
public User User { get; }

0 commit comments

Comments
 (0)