Skip to content

Commit 3a782c8

Browse files
0xFirekeeper0xFirekeeper
authored
Unify AWS Date Override Behavior (#121)
1 parent e3a9e72 commit 3a782c8

File tree

2 files changed

+97
-69
lines changed

2 files changed

+97
-69
lines changed

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,6 @@ private async Task<string> PostAuth(Server.VerifyResult result)
264264

265265
private async Task<string> MigrateShardToEnclave(Server.VerifyResult authResult)
266266
{
267-
// TODO: For recovery code, allow old encryption keys as overrides to migrate sharded custom auth?
268267
var (address, encryptedPrivateKeyB64, ivB64, kmsCiphertextB64) = await this.EmbeddedWallet
269268
.GenerateEncryptionDataAsync(authResult.AuthToken, this.LegacyEncryptionKey ?? authResult.RecoveryCode)
270269
.ConfigureAwait(false);

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

Lines changed: 97 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -53,66 +53,18 @@ private static async Task<AwsCredentials> GetTemporaryCredentialsAsync(string id
5353
};
5454
}
5555

56-
private static async Task<JToken> GenerateDataKey(AwsCredentials credentials, IThirdwebHttpClient httpClient, DateTime? dateOverride = null)
56+
private static async Task<JToken> GenerateDataKey(AwsCredentials credentials, IThirdwebHttpClient httpClient)
5757
{
58-
var client = Utils.ReconstructHttpClient(httpClient);
5958
var endpoint = $"https://kms.{AWS_REGION}.amazonaws.com/";
6059

6160
var payloadForGenerateDataKey = new { KeyId = _migrationKeyId, KeySpec = "AES_256" };
61+
var requestBodyString = JsonConvert.SerializeObject(payloadForGenerateDataKey);
6262

63-
var content = new StringContent(JsonConvert.SerializeObject(payloadForGenerateDataKey), Encoding.UTF8, "application/x-amz-json-1.1");
63+
var contentType = "application/x-amz-json-1.1";
6464

65-
client.AddHeader("X-Amz-Target", "TrentService.GenerateDataKey");
65+
var extraHeaders = new Dictionary<string, string> { { "X-Amz-Target", "TrentService.GenerateDataKey" } };
6666

67-
var dateTimeNow = dateOverride ?? DateTime.UtcNow;
68-
var dateStamp = dateTimeNow.ToString("yyyyMMdd");
69-
var amzDateFormat = "yyyyMMddTHHmmssZ";
70-
var amzDate = dateTimeNow.ToString(amzDateFormat);
71-
var canonicalUri = "/";
72-
73-
var canonicalHeaders = $"host:kms.{AWS_REGION}.amazonaws.com\nx-amz-date:{amzDate}\n";
74-
var signedHeaders = "host;x-amz-date";
75-
76-
#if NETSTANDARD
77-
using var sha256 = SHA256.Create();
78-
var payloadHash = ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
79-
#else
80-
var payloadHash = ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(await content.ReadAsStringAsync())));
81-
#endif
82-
83-
var canonicalRequest = $"POST\n{canonicalUri}\n\n{canonicalHeaders}\n{signedHeaders}\n{payloadHash}";
84-
85-
var algorithm = "AWS4-HMAC-SHA256";
86-
var credentialScope = $"{dateStamp}/{AWS_REGION}/kms/aws4_request";
87-
88-
#if NETSTANDARD
89-
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(canonicalRequest)))}";
90-
#else
91-
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(canonicalRequest)))}";
92-
#endif
93-
94-
var signingKey = GetSignatureKey(credentials.SecretAccessKey, dateStamp, AWS_REGION, "kms");
95-
var signature = ToHexString(HMACSHA256(signingKey, stringToSign));
96-
97-
var authorizationHeader = $"{algorithm} Credential={credentials.AccessKeyId}/{credentialScope}, SignedHeaders={signedHeaders}, Signature={signature}";
98-
99-
client.AddHeader("x-amz-date", amzDate);
100-
client.AddHeader("Authorization", authorizationHeader);
101-
client.AddHeader("x-amz-security-token", credentials.SessionToken);
102-
103-
var response = await client.PostAsync(endpoint, content).ConfigureAwait(false);
104-
var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
105-
106-
if (!response.IsSuccessStatusCode)
107-
{
108-
if (dateOverride == null && responseContent.Contains("InvalidSignatureException"))
109-
{
110-
var parsedTime = responseContent.Substring(responseContent.LastIndexOf('(') + 1, amzDate.Length);
111-
return await GenerateDataKey(credentials, httpClient, DateTime.ParseExact(parsedTime, amzDateFormat, System.Globalization.CultureInfo.InvariantCulture).ToUniversalTime())
112-
.ConfigureAwait(false);
113-
}
114-
throw new Exception($"Failed to generate data key: {responseContent}");
115-
}
67+
var responseContent = await PostAwsRequestWithDateOverride(credentials, httpClient, AWS_REGION, "kms", endpoint, "/", "", requestBodyString, contentType, extraHeaders).ConfigureAwait(false);
11668

11769
var responseObject = JToken.Parse(responseContent);
11870
var plaintextKeyBlob = responseObject["Plaintext"];
@@ -129,54 +81,131 @@ private static async Task<JToken> GenerateDataKey(AwsCredentials credentials, IT
12981
private static async Task<MemoryStream> InvokeLambdaWithTemporaryCredentialsAsync(AwsCredentials credentials, string invokePayload, IThirdwebHttpClient httpClient, string lambdaFunction)
13082
{
13183
var endpoint = $"https://lambda.{AWS_REGION}.amazonaws.com/2015-03-31/functions/{lambdaFunction}/invocations";
132-
var requestBody = new StringContent(invokePayload, Encoding.UTF8, "application/json");
84+
var contentType = "application/json";
85+
86+
var canonicalUri = $"/2015-03-31/functions/{Uri.EscapeDataString(lambdaFunction)}/invocations";
87+
var canonicalQueryString = "";
88+
89+
var extraHeaders = new Dictionary<string, string>();
90+
91+
var responseContent = await PostAwsRequestWithDateOverride(
92+
credentials,
93+
httpClient,
94+
AWS_REGION,
95+
"lambda",
96+
endpoint,
97+
canonicalUri,
98+
canonicalQueryString,
99+
invokePayload,
100+
contentType,
101+
extraHeaders
102+
)
103+
.ConfigureAwait(false);
104+
105+
var memoryStream = new MemoryStream(Encoding.UTF8.GetBytes(responseContent));
106+
return memoryStream;
107+
}
133108

109+
private static async Task<string> PostAwsRequestWithDateOverride(
110+
AwsCredentials credentials,
111+
IThirdwebHttpClient httpClient,
112+
string region,
113+
string service,
114+
string endpoint,
115+
string canonicalUri,
116+
string canonicalQueryString,
117+
string requestBodyString,
118+
string contentType,
119+
Dictionary<string, string> extraHeaders,
120+
DateTime? dateOverride = null
121+
)
122+
{
134123
var client = Utils.ReconstructHttpClient(httpClient);
135124

136-
var dateTimeNow = DateTime.UtcNow;
125+
if (extraHeaders != null)
126+
{
127+
foreach (var kvp in extraHeaders)
128+
{
129+
client.AddHeader(kvp.Key, kvp.Value);
130+
}
131+
}
132+
133+
var dateTimeNow = dateOverride ?? DateTime.UtcNow;
134+
var amzDateFormat = "yyyyMMddTHHmmssZ";
135+
var amzDate = dateTimeNow.ToString(amzDateFormat);
137136
var dateStamp = dateTimeNow.ToString("yyyyMMdd");
138-
var amzDate = dateTimeNow.ToString("yyyyMMddTHHmmssZ");
139137

140-
var canonicalUri = "/2015-03-31/functions/" + Uri.EscapeDataString(lambdaFunction) + "/invocations";
141-
var canonicalQueryString = "";
142-
var canonicalHeaders = $"host:lambda.{AWS_REGION}.amazonaws.com\nx-amz-date:{amzDate}\n";
138+
var canonicalHeaders = $"host:{new Uri(endpoint).Host}\n" + $"x-amz-date:{amzDate}\n";
143139
var signedHeaders = "host;x-amz-date";
140+
144141
#if NETSTANDARD
145142
using var sha256 = SHA256.Create();
146-
var payloadHash = ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(invokePayload)));
143+
var payloadHash = ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(requestBodyString)));
147144
#else
148-
var payloadHash = ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(invokePayload)));
145+
var payloadHash = ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(requestBodyString)));
149146
#endif
147+
150148
var canonicalRequest = $"POST\n{canonicalUri}\n{canonicalQueryString}\n{canonicalHeaders}\n{signedHeaders}\n{payloadHash}";
151149

152150
var algorithm = "AWS4-HMAC-SHA256";
153-
var credentialScope = $"{dateStamp}/{AWS_REGION}/lambda/aws4_request";
151+
var credentialScope = $"{dateStamp}/{region}/{service}/aws4_request";
154152
#if NETSTANDARD
155153
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(sha256.ComputeHash(Encoding.UTF8.GetBytes(canonicalRequest)))}";
156154
#else
157155
var stringToSign = $"{algorithm}\n{amzDate}\n{credentialScope}\n{ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(canonicalRequest)))}";
158156
#endif
159157

160-
var signingKey = GetSignatureKey(credentials.SecretAccessKey, dateStamp, AWS_REGION, "lambda");
158+
var signingKey = GetSignatureKey(credentials.SecretAccessKey, dateStamp, region, service);
161159
var signature = ToHexString(HMACSHA256(signingKey, stringToSign));
162160

163161
var authorizationHeader = $"{algorithm} Credential={credentials.AccessKeyId}/{credentialScope}, SignedHeaders={signedHeaders}, Signature={signature}";
164162

165163
client.AddHeader("x-amz-date", amzDate);
166164
client.AddHeader("Authorization", authorizationHeader);
167-
client.AddHeader("x-amz-security-token", credentials.SessionToken);
168165

169-
var response = await client.PostAsync(endpoint, requestBody).ConfigureAwait(false);
166+
if (!string.IsNullOrEmpty(credentials.SessionToken))
167+
{
168+
client.AddHeader("x-amz-security-token", credentials.SessionToken);
169+
}
170+
171+
var content = new StringContent(requestBodyString, Encoding.UTF8, contentType);
170172

173+
var response = await client.PostAsync(endpoint, content).ConfigureAwait(false);
171174
var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
172175

173176
if (!response.IsSuccessStatusCode)
174177
{
175-
throw new Exception($"Lambda invocation failed: {responseContent}");
178+
if (dateOverride == null && responseContent.Contains("Signature expired"))
179+
{
180+
var idx = responseContent.LastIndexOf('(');
181+
if (idx > -1)
182+
{
183+
var parsedTimeString = responseContent.Substring(idx + 1, amzDate.Length);
184+
var serverTime = DateTime.ParseExact(parsedTimeString, amzDateFormat, System.Globalization.CultureInfo.InvariantCulture).ToUniversalTime();
185+
186+
Console.WriteLine($"Server time: {serverTime}");
187+
188+
return await PostAwsRequestWithDateOverride(
189+
credentials,
190+
httpClient,
191+
region,
192+
service,
193+
endpoint,
194+
canonicalUri,
195+
canonicalQueryString,
196+
requestBodyString,
197+
contentType,
198+
extraHeaders,
199+
serverTime
200+
)
201+
.ConfigureAwait(false);
202+
}
203+
}
204+
205+
throw new Exception($"AWS request failed: {responseContent}");
176206
}
177207

178-
var memoryStream = new MemoryStream(Encoding.UTF8.GetBytes(responseContent));
179-
return memoryStream;
208+
return responseContent;
180209
}
181210

182211
private static byte[] HMACSHA256(byte[] key, string data)

0 commit comments

Comments
 (0)