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 ecr_public_client(&self) -> aws_sdk_ecrpublic::Client {
241        // ECR Public is a global service only available in us-east-1.
242        // Reuse stored credentials but override region.
243        let config = aws_config::SdkConfig::builder()
244            .behavior_version(aws_config::BehaviorVersion::latest())
245            .region(aws_config::Region::new(ECR_PUBLIC_REGION))
246            .credentials_provider(
247                self.sdk_config
248                    .credentials_provider()
249                    .expect("credentials must be set")
250                    .clone(),
251            )
252            .build();
253        aws_sdk_ecrpublic::Client::new(&config)
254    }
255
256    pub fn cloudwatch_client(&self) -> aws_sdk_cloudwatch::Client {
257        aws_sdk_cloudwatch::Client::new(&self.sdk_config)
258    }
259
260    pub fn sqs_client(&self) -> aws_sdk_sqs::Client {
261        aws_sdk_sqs::Client::new(&self.sdk_config)
262    }
263
264    pub fn cloudwatch_logs_client(&self) -> aws_sdk_cloudwatchlogs::Client {
265        aws_sdk_cloudwatchlogs::Client::new(&self.sdk_config)
266    }
267
268    pub fn pipes_client(&self) -> aws_sdk_pipes::Client {
269        aws_sdk_pipes::Client::new(&self.sdk_config)
270    }
271
272    pub fn ec2_client(&self) -> aws_sdk_ec2::Client {
273        aws_sdk_ec2::Client::new(&self.sdk_config)
274    }
275
276    pub fn apigateway_client(&self) -> aws_sdk_apigateway::Client {
277        aws_sdk_apigateway::Client::new(&self.sdk_config)
278    }
279
280    pub fn apigatewayv2_client(&self) -> aws_sdk_apigatewayv2::Client {
281        aws_sdk_apigatewayv2::Client::new(&self.sdk_config)
282    }
283
284    pub fn alarms_client(&self) -> aws_sdk_cloudwatch::Client {
285        aws_sdk_cloudwatch::Client::new(&self.sdk_config)
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn test_dummy_config_with_region() {
295        let region = "us-west-2";
296        let config = AwsConfig::dummy(Some(region.to_string()));
297        assert_eq!(config.region, region);
298        assert_eq!(config.account_id, "");
299        assert!(!config.region_auto_detected);
300    }
301
302    #[test]
303    fn test_dummy_config_without_region() {
304        let config = AwsConfig::dummy(None);
305        assert_eq!(config.region, "");
306        assert_eq!(config.account_id, "");
307        assert!(!config.region_auto_detected);
308    }
309
310    #[tokio::test]
311    async fn test_new_fails_without_credentials() {
312        // Clear AWS env vars to simulate no credentials
313        std::env::remove_var("AWS_ACCESS_KEY_ID");
314        std::env::remove_var("AWS_SECRET_ACCESS_KEY");
315        std::env::remove_var("AWS_SESSION_TOKEN");
316        std::env::set_var("AWS_PROFILE", "nonexistent-profile-test");
317
318        let result = AwsConfig::new(Some("us-east-1".to_string())).await;
319
320        // Should fail with credentials error before attempting region detection
321        assert!(result.is_err());
322        let err_str = format!("{}", result.unwrap_err());
323        // Can be credentials error or dispatch failure (both indicate auth issues)
324        assert!(
325            err_str.contains("credentials")
326                || err_str.contains("profile")
327                || err_str.contains("dispatch"),
328            "Expected auth error, got: {}",
329            err_str
330        );
331    }
332
333    #[tokio::test]
334    async fn test_timeout_is_configurable() {
335        std::env::remove_var("AWS_ACCESS_KEY_ID");
336        std::env::remove_var("AWS_SECRET_ACCESS_KEY");
337        std::env::set_var("AWS_PROFILE", "nonexistent-profile-test");
338
339        // Test with very short timeout
340        let result = AwsConfig::new_with_timeout(
341            Some("us-east-1".to_string()),
342            std::time::Duration::from_millis(100),
343        )
344        .await;
345
346        assert!(result.is_err());
347    }
348
349    #[test]
350    fn test_dummy_config_preserves_values() {
351        let region = "eu-west-1";
352        let config = AwsConfig::dummy(Some(region.to_string()));
353        assert_eq!(config.region, region);
354        assert_eq!(config.account_id, "");
355        assert_eq!(config.role_arn, "");
356        assert!(!config.region_auto_detected);
357    }
358
359    #[test]
360    fn test_ecr_public_region_is_us_east_1() {
361        // ECR Public is a global service — must always use us-east-1 regardless of user region.
362        assert_eq!(ECR_PUBLIC_REGION, "us-east-1");
363    }
364
365    #[test]
366    fn test_client_factories_are_sync() {
367        // All client factories (except s3_client_with_region) must be sync —
368        // they should not do any async work after credentials are resolved.
369        let config = AwsConfig::dummy(Some("us-east-1".to_string()));
370        let _ = config.s3_client();
371        let _ = config.ecr_client();
372        let _ = config.cloudwatch_client();
373        let _ = config.lambda_client();
374        let _ = config.iam_client();
375        let _ = config.sqs_client();
376        let _ = config.ec2_client();
377    }
378}