Skip to main content

rusticity_core/
config.rs

1use anyhow::Result;
2
3/// ECR Public is a global service only available in us-east-1.
4pub const ECR_PUBLIC_REGION: &str = "us-east-1";
5
6#[derive(Clone, Debug)]
7pub struct AwsConfig {
8    pub region: String,
9    pub account_id: String,
10    pub role_arn: String,
11    pub region_auto_detected: bool,
12    /// Resolved SDK config — reused by all client factories to avoid re-running credential chain.
13    sdk_config: aws_config::SdkConfig,
14}
15
16impl AwsConfig {
17    pub async fn new(region: Option<String>) -> Result<Self> {
18        Self::new_with_timeout(region, std::time::Duration::from_secs(10)).await
19    }
20
21    pub async fn new_with_timeout(
22        region: Option<String>,
23        timeout: std::time::Duration,
24    ) -> Result<Self> {
25        // Check for region early to avoid IMDS timeout
26        if region.is_none()
27            && std::env::var("AWS_REGION").is_err()
28            && std::env::var("AWS_DEFAULT_REGION").is_err()
29        {
30            return Err(anyhow::anyhow!("Missing Region"));
31        }
32
33        let mut config_loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
34
35        // Use profile from AWS_PROFILE env var if set
36        let profile_to_use = std::env::var("AWS_PROFILE").ok();
37        if let Some(ref profile) = profile_to_use {
38            if !profile.is_empty() {
39                tracing::info!("Using AWS profile: {}", profile);
40                config_loader = config_loader.profile_name(profile);
41            }
42        }
43
44        if let Some(r) = &region {
45            config_loader = config_loader.region(aws_config::Region::new(r.clone()));
46        }
47
48        // Load config once — stored and reused for all client factories.
49        let sdk_config = tokio::time::timeout(timeout, config_loader.load())
50            .await
51            .map_err(|_| anyhow::anyhow!("Timeout loading AWS config"))?;
52
53        // Double-check region is set
54        if sdk_config.region().is_none() {
55            return Err(anyhow::anyhow!("Missing Region"));
56        }
57
58        // Try to get identity with timeout
59        let (account_id, role_arn) =
60            match tokio::time::timeout(timeout, Self::try_get_identity(&sdk_config)).await {
61                Ok(Ok((acc, role))) => {
62                    tracing::info!("Loaded identity: account={}, role={}", acc, role);
63                    (acc, role)
64                }
65                Ok(Err(e)) => {
66                    tracing::error!("Failed to get identity: {}", e);
67                    return Err(e);
68                }
69                Err(_) => return Err(anyhow::anyhow!("Timeout getting AWS identity")),
70            };
71
72        let (region_str, auto_detected) = match sdk_config.region() {
73            Some(r) => (r.as_ref().to_string(), false),
74            None => {
75                let fastest = Self::find_fastest_region().await?;
76                (fastest, true)
77            }
78        };
79
80        Ok(Self {
81            region: region_str,
82            account_id,
83            role_arn,
84            region_auto_detected: auto_detected,
85            sdk_config,
86        })
87    }
88
89    async fn try_get_identity(config: &aws_config::SdkConfig) -> Result<(String, String)> {
90        let sts_client = aws_sdk_sts::Client::new(config);
91        let identity = sts_client.get_caller_identity().send().await?;
92        let account_id = identity.account().unwrap_or("").to_string();
93        let role_arn = identity.arn().unwrap_or("").to_string();
94        Ok((account_id, role_arn))
95    }
96
97    async fn find_fastest_region() -> Result<String> {
98        use std::time::Instant;
99
100        let regions = [
101            "us-east-1",
102            "us-east-2",
103            "us-west-1",
104            "us-west-2",
105            "af-south-1",
106            "ap-east-1",
107            "ap-south-1",
108            "ap-south-2",
109            "ap-northeast-1",
110            "ap-northeast-2",
111            "ap-northeast-3",
112            "ap-southeast-1",
113            "ap-southeast-2",
114            "ap-southeast-3",
115            "ap-southeast-4",
116            "ca-central-1",
117            "ca-west-1",
118            "eu-central-1",
119            "eu-central-2",
120            "eu-west-1",
121            "eu-west-2",
122            "eu-west-3",
123            "eu-north-1",
124            "eu-south-1",
125            "eu-south-2",
126            "il-central-1",
127            "me-central-1",
128            "me-south-1",
129            "sa-east-1",
130        ];
131
132        let mut tasks = Vec::new();
133
134        for &region in &regions {
135            let region_name = region.to_string();
136            tasks.push(tokio::spawn(async move {
137                let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
138                    .region(aws_config::Region::new(region_name.clone()))
139                    .load()
140                    .await;
141                let s3 = aws_sdk_s3::Client::new(&config);
142                let start = Instant::now();
143                match tokio::time::timeout(
144                    std::time::Duration::from_secs(2),
145                    s3.list_buckets().send(),
146                )
147                .await
148                {
149                    Ok(Ok(_)) => Some((region_name, start.elapsed())),
150                    _ => Some((region_name, std::time::Duration::from_secs(9999))),
151                }
152            }));
153        }
154
155        let results = futures::future::join_all(tasks).await;
156        let mut latencies: Vec<(String, std::time::Duration)> = results
157            .into_iter()
158            .filter_map(|r| r.ok().flatten())
159            .collect();
160
161        latencies.sort_by_key(|(_, d)| *d);
162
163        latencies
164            .first()
165            .map(|(r, _)| r.clone())
166            .ok_or_else(|| anyhow::anyhow!("Could not determine fastest region"))
167    }
168
169    pub fn dummy(region: Option<String>) -> Self {
170        // Build a minimal SdkConfig with just the region — no credential resolution.
171        let region_str = region.unwrap_or_default();
172        let sdk_config = aws_config::SdkConfig::builder()
173            .behavior_version(aws_config::BehaviorVersion::latest())
174            .region(aws_config::Region::new(region_str.clone()))
175            .build();
176        Self {
177            region: region_str,
178            account_id: "".to_string(),
179            role_arn: "".to_string(),
180            region_auto_detected: false,
181            sdk_config,
182        }
183    }
184
185    pub async fn get_account_for_profile(profile: &str, region: &str) -> Result<String> {
186        let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
187            .profile_name(profile)
188            .region(aws_config::Region::new(region.to_string()))
189            .load()
190            .await;
191
192        let sts_client = aws_sdk_sts::Client::new(&config);
193        let identity = sts_client.get_caller_identity().send().await?;
194        Ok(identity.account().unwrap_or("").to_string())
195    }
196
197    // ── Client factories ──────────────────────────────────────────────────────
198    // All factories reuse self.sdk_config — credentials are already resolved,
199    // so these are instantaneous (no network calls).
200
201    pub fn s3_client(&self) -> aws_sdk_s3::Client {
202        aws_sdk_s3::Client::new(&self.sdk_config)
203    }
204
205    pub async fn s3_client_with_region(&self, region: &str) -> aws_sdk_s3::Client {
206        // Needs a different region — build a new config with the same credentials.
207        let config = aws_config::SdkConfig::builder()
208            .behavior_version(aws_config::BehaviorVersion::latest())
209            .region(aws_config::Region::new(region.to_string()))
210            .credentials_provider(
211                self.sdk_config
212                    .credentials_provider()
213                    .expect("credentials must be set")
214                    .clone(),
215            )
216            .build();
217        aws_sdk_s3::Client::new(&config)
218    }
219
220    pub fn cloudformation_client(&self) -> aws_sdk_cloudformation::Client {
221        aws_sdk_cloudformation::Client::new(&self.sdk_config)
222    }
223
224    pub fn cloudtrail_client(&self) -> aws_sdk_cloudtrail::Client {
225        aws_sdk_cloudtrail::Client::new(&self.sdk_config)
226    }
227
228    pub fn lambda_client(&self) -> aws_sdk_lambda::Client {
229        aws_sdk_lambda::Client::new(&self.sdk_config)
230    }
231
232    pub fn iam_client(&self) -> aws_sdk_iam::Client {
233        aws_sdk_iam::Client::new(&self.sdk_config)
234    }
235
236    pub fn ecr_client(&self) -> aws_sdk_ecr::Client {
237        aws_sdk_ecr::Client::new(&self.sdk_config)
238    }
239
240    pub fn kms_client(&self) -> aws_sdk_kms::Client {
241        aws_sdk_kms::Client::new(&self.sdk_config)
242    }
243
244    pub fn efs_client(&self) -> aws_sdk_efs::Client {
245        aws_sdk_efs::Client::new(&self.sdk_config)
246    }
247
248    pub fn fsx_client(&self) -> aws_sdk_fsx::Client {
249        aws_sdk_fsx::Client::new(&self.sdk_config)
250    }
251
252    pub fn ecr_public_client(&self) -> aws_sdk_ecrpublic::Client {
253        // ECR Public is a global service only available in us-east-1.
254        // Reuse stored credentials but override region.
255        let config = aws_config::SdkConfig::builder()
256            .behavior_version(aws_config::BehaviorVersion::latest())
257            .region(aws_config::Region::new(ECR_PUBLIC_REGION))
258            .credentials_provider(
259                self.sdk_config
260                    .credentials_provider()
261                    .expect("credentials must be set")
262                    .clone(),
263            )
264            .build();
265        aws_sdk_ecrpublic::Client::new(&config)
266    }
267
268    pub fn cloudwatch_client(&self) -> aws_sdk_cloudwatch::Client {
269        aws_sdk_cloudwatch::Client::new(&self.sdk_config)
270    }
271
272    pub fn sqs_client(&self) -> aws_sdk_sqs::Client {
273        aws_sdk_sqs::Client::new(&self.sdk_config)
274    }
275
276    pub fn cloudwatch_logs_client(&self) -> aws_sdk_cloudwatchlogs::Client {
277        aws_sdk_cloudwatchlogs::Client::new(&self.sdk_config)
278    }
279
280    pub fn pipes_client(&self) -> aws_sdk_pipes::Client {
281        aws_sdk_pipes::Client::new(&self.sdk_config)
282    }
283
284    pub fn ec2_client(&self) -> aws_sdk_ec2::Client {
285        aws_sdk_ec2::Client::new(&self.sdk_config)
286    }
287
288    pub fn apigateway_client(&self) -> aws_sdk_apigateway::Client {
289        aws_sdk_apigateway::Client::new(&self.sdk_config)
290    }
291
292    pub fn apigatewayv2_client(&self) -> aws_sdk_apigatewayv2::Client {
293        aws_sdk_apigatewayv2::Client::new(&self.sdk_config)
294    }
295
296    pub fn alarms_client(&self) -> aws_sdk_cloudwatch::Client {
297        aws_sdk_cloudwatch::Client::new(&self.sdk_config)
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304
305    #[test]
306    fn test_dummy_config_with_region() {
307        let region = "us-west-2";
308        let config = AwsConfig::dummy(Some(region.to_string()));
309        assert_eq!(config.region, region);
310        assert_eq!(config.account_id, "");
311        assert!(!config.region_auto_detected);
312    }
313
314    #[test]
315    fn test_dummy_config_without_region() {
316        let config = AwsConfig::dummy(None);
317        assert_eq!(config.region, "");
318        assert_eq!(config.account_id, "");
319        assert!(!config.region_auto_detected);
320    }
321
322    #[tokio::test]
323    async fn test_new_fails_without_credentials() {
324        // Clear AWS env vars to simulate no credentials
325        std::env::remove_var("AWS_ACCESS_KEY_ID");
326        std::env::remove_var("AWS_SECRET_ACCESS_KEY");
327        std::env::remove_var("AWS_SESSION_TOKEN");
328        std::env::set_var("AWS_PROFILE", "nonexistent-profile-test");
329
330        let result = AwsConfig::new(Some("us-east-1".to_string())).await;
331
332        // Should fail with credentials error before attempting region detection
333        assert!(result.is_err());
334        let err_str = format!("{}", result.unwrap_err());
335        // Can be credentials error or dispatch failure (both indicate auth issues)
336        assert!(
337            err_str.contains("credentials")
338                || err_str.contains("profile")
339                || err_str.contains("dispatch"),
340            "Expected auth error, got: {}",
341            err_str
342        );
343    }
344
345    #[tokio::test]
346    async fn test_timeout_is_configurable() {
347        std::env::remove_var("AWS_ACCESS_KEY_ID");
348        std::env::remove_var("AWS_SECRET_ACCESS_KEY");
349        std::env::set_var("AWS_PROFILE", "nonexistent-profile-test");
350
351        // Test with very short timeout
352        let result = AwsConfig::new_with_timeout(
353            Some("us-east-1".to_string()),
354            std::time::Duration::from_millis(100),
355        )
356        .await;
357
358        assert!(result.is_err());
359    }
360
361    #[test]
362    fn test_dummy_config_preserves_values() {
363        let region = "eu-west-1";
364        let config = AwsConfig::dummy(Some(region.to_string()));
365        assert_eq!(config.region, region);
366        assert_eq!(config.account_id, "");
367        assert_eq!(config.role_arn, "");
368        assert!(!config.region_auto_detected);
369    }
370
371    #[test]
372    fn test_ecr_public_region_is_us_east_1() {
373        // ECR Public is a global service — must always use us-east-1 regardless of user region.
374        assert_eq!(ECR_PUBLIC_REGION, "us-east-1");
375    }
376
377    #[test]
378    fn test_client_factories_are_sync() {
379        // All client factories (except s3_client_with_region) must be sync —
380        // they should not do any async work after credentials are resolved.
381        let config = AwsConfig::dummy(Some("us-east-1".to_string()));
382        let _ = config.s3_client();
383        let _ = config.ecr_client();
384        let _ = config.cloudwatch_client();
385        let _ = config.lambda_client();
386        let _ = config.iam_client();
387        let _ = config.sqs_client();
388        let _ = config.ec2_client();
389    }
390}