Skip to content

Commit 799a1a8

Browse files
committed
Enhance AWS KMS integration and error handling
- Updated AWS KMS credential extraction to use a more consistent header naming convention. - Improved error handling for retryable RPC errors in the EoaExecutorWorkerError, allowing for better job management. - Refactored AWS KMS credential validation to ensure proper extraction of key IDs and regions from ARNs. These changes aim to streamline AWS KMS operations and enhance the robustness of error handling in the executor worker.
1 parent 6bcbafc commit 799a1a8

File tree

3 files changed

+35
-9
lines changed

3 files changed

+35
-9
lines changed

core/src/credentials.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use alloy::primitives::ChainId;
22
use alloy_signer_aws::AwsSigner;
3-
use aws_config::BehaviorVersion;
3+
use aws_config::{BehaviorVersion, Region};
44
use aws_credential_types::provider::future::ProvideCredentials as ProvideCredentialsFuture;
55
use aws_sdk_kms::config::{Credentials, ProvideCredentials};
66
use serde::{Deserialize, Serialize};
@@ -48,6 +48,7 @@ impl AwsKmsCredential {
4848
pub async fn get_signer(&self, chain_id: Option<ChainId>) -> Result<AwsSigner, EngineError> {
4949
let config = aws_config::defaults(BehaviorVersion::latest())
5050
.credentials_provider(self.clone())
51+
.region(Region::new(self.region.clone()))
5152
.load()
5253
.await;
5354
let client = aws_sdk_kms::Client::new(&config);

executors/src/eoa/worker/error.rs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,17 @@ impl EoaExecutorWorkerError {
9494
inner_error: TransactionStoreError::LockLost { .. },
9595
..
9696
} => JobError::Fail(self),
97+
EoaExecutorWorkerError::RpcError { .. } => {
98+
if is_retryable_preparation_error(&self) {
99+
JobError::Nack {
100+
error: self,
101+
delay: Some(Duration::from_secs(10)),
102+
position: RequeuePosition::Last,
103+
}
104+
} else {
105+
JobError::Fail(self)
106+
}
107+
}
97108
_ => JobError::Nack {
98109
error: self,
99110
delay: Some(Duration::from_secs(10)),
@@ -219,6 +230,11 @@ pub fn is_retryable_rpc_error(kind: &RpcErrorKind) -> bool {
219230
match kind {
220231
RpcErrorKind::TransportHttpError { status, .. } if *status >= 400 && *status < 500 => false,
221232
RpcErrorKind::UnsupportedFeature { .. } => false,
233+
RpcErrorKind::ErrorResp(resp) => {
234+
let message = resp.message.to_lowercase();
235+
// if the error message contains "invalid chain", it's not retryable
236+
!message.contains("invalid chain")
237+
}
222238
_ => true,
223239
}
224240
}

server/src/http/extractors.rs

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ const HEADER_THIRDWEB_SERVICE_KEY: &str = "x-thirdweb-service-key";
2222
const HEADER_WALLET_ACCESS_TOKEN: &str = "x-wallet-access-token";
2323
const HEADER_VAULT_ACCESS_TOKEN: &str = "x-vault-access-token";
2424
const HEADER_AWS_KMS_ARN: &str = "x-aws-kms-arn";
25-
const HEADER_AWS_KMS_ACCESS_KEY_ID: &str = "x-aws-kms-access-key-id";
25+
const HEADER_AWS_ACCESS_KEY_ID: &str = "x-aws-access-key-id";
2626
const HEADER_AWS_SECRET_ACCESS_KEY: &str = "x-aws-secret-access-key";
2727

2828
/// Extractor for RPC credentials from headers
@@ -147,7 +147,7 @@ impl SigningCredentialsExtractor {
147147
/// Try to extract AWS KMS credentials from headers
148148
fn try_extract_aws_kms(parts: &Parts) -> Result<Option<AwsKmsCredential>, ApiEngineError> {
149149
let arn = Self::get_header_value(parts, HEADER_AWS_KMS_ARN);
150-
let access_key_id = Self::get_header_value(parts, HEADER_AWS_KMS_ACCESS_KEY_ID);
150+
let access_key_id = Self::get_header_value(parts, HEADER_AWS_ACCESS_KEY_ID);
151151
let secret_access_key = Self::get_header_value(parts, HEADER_AWS_SECRET_ACCESS_KEY);
152152

153153
match (arn, access_key_id, secret_access_key) {
@@ -180,12 +180,19 @@ impl SigningCredentialsExtractor {
180180
}
181181

182182
// Extract and validate key ID
183-
let key_id = parsed_arn.resource.to_string();
184-
if key_id.is_empty() {
185-
return Err(ApiEngineError(EngineError::ValidationError {
186-
message: "KMS ARN must contain a valid key ID in the resource part".to_string(),
187-
}));
188-
}
183+
let key_id = parsed_arn
184+
.resource
185+
.path_split()
186+
.last()
187+
.map(|id| {
188+
dbg!(&id);
189+
id.to_string()
190+
})
191+
.ok_or_else(|| {
192+
ApiEngineError(EngineError::ValidationError {
193+
message: "KMS ARN must contain a valid key ID in the resource part".to_string(),
194+
})
195+
})?;
189196

190197
// Extract and validate region
191198
let region = parsed_arn.region.ok_or_else(|| {
@@ -194,6 +201,8 @@ impl SigningCredentialsExtractor {
194201
})
195202
})?;
196203

204+
dbg!(&region);
205+
197206
Ok((key_id, region.to_string()))
198207
}
199208

0 commit comments

Comments
 (0)