Line data    Source code

       1              : use std::time::{Duration, SystemTime};
       2              : 
       3              : use aws_config::meta::credentials::CredentialsProviderChain;
       4              : use aws_sdk_iam::config::ProvideCredentials;
       5              : use aws_sigv4::http_request::{
       6              :     self, SignableBody, SignableRequest, SignatureLocation, SigningSettings,
       7              : };
       8              : use tracing::info;
       9              : 
      10              : #[derive(Debug)]
      11              : pub struct AWSIRSAConfig {
      12              :     region: String,
      13              :     service_name: String,
      14              :     cluster_name: String,
      15              :     user_id: String,
      16              :     token_ttl: Duration,
      17              :     action: String,
      18              : }
      19              : 
      20              : impl AWSIRSAConfig {
      21            0 :     pub fn new(region: String, cluster_name: Option<String>, user_id: Option<String>) -> Self {
      22            0 :         AWSIRSAConfig {
      23            0 :             region,
      24            0 :             service_name: "elasticache".to_string(),
      25            0 :             cluster_name: cluster_name.unwrap_or_default(),
      26            0 :             user_id: user_id.unwrap_or_default(),
      27            0 :             // "The IAM authentication token is valid for 15 minutes"
      28            0 :             // https://docs.aws.amazon.com/memorydb/latest/devguide/auth-iam.html#auth-iam-limits
      29            0 :             token_ttl: Duration::from_secs(15 * 60),
      30            0 :             action: "connect".to_string(),
      31            0 :         }
      32            0 :     }
      33              : }
      34              : 
      35              : /// Credentials provider for AWS elasticache authentication.
      36              : ///
      37              : /// Official documentation:
      38              : /// <https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/auth-iam.html>
      39              : ///
      40              : /// Useful resources:
      41              : /// <https://aws.amazon.com/blogs/database/simplify-managing-access-to-amazon-elasticache-for-redis-clusters-with-iam/>
      42              : pub struct CredentialsProvider {
      43              :     config: AWSIRSAConfig,
      44              :     credentials_provider: CredentialsProviderChain,
      45              : }
      46              : 
      47              : impl CredentialsProvider {
      48            0 :     pub fn new(config: AWSIRSAConfig, credentials_provider: CredentialsProviderChain) -> Self {
      49            0 :         CredentialsProvider {
      50            0 :             config,
      51            0 :             credentials_provider,
      52            0 :         }
      53            0 :     }
      54            0 :     pub(crate) async fn provide_credentials(&self) -> anyhow::Result<(String, String)> {
      55            0 :         let aws_credentials = self
      56            0 :             .credentials_provider
      57            0 :             .provide_credentials()
      58            0 :             .await?
      59            0 :             .into();
      60            0 :         info!("AWS credentials successfully obtained");
      61            0 :         info!("Connecting to Redis with configuration: {:?}", self.config);
      62            0 :         let mut settings = SigningSettings::default();
      63            0 :         settings.signature_location = SignatureLocation::QueryParams;
      64            0 :         settings.expires_in = Some(self.config.token_ttl);
      65            0 :         let signing_params = aws_sigv4::sign::v4::SigningParams::builder()
      66            0 :             .identity(&aws_credentials)
      67            0 :             .region(&self.config.region)
      68            0 :             .name(&self.config.service_name)
      69            0 :             .time(SystemTime::now())
      70            0 :             .settings(settings)
      71            0 :             .build()?
      72            0 :             .into();
      73            0 :         let auth_params = [
      74            0 :             ("Action", &self.config.action),
      75            0 :             ("User", &self.config.user_id),
      76            0 :         ];
      77            0 :         let auth_params = url::form_urlencoded::Serializer::new(String::new())
      78            0 :             .extend_pairs(auth_params)
      79            0 :             .finish();
      80            0 :         let auth_uri = http::Uri::builder()
      81            0 :             .scheme("http")
      82            0 :             .authority(self.config.cluster_name.as_bytes())
      83            0 :             .path_and_query(format!("/?{auth_params}"))
      84            0 :             .build()?;
      85            0 :         info!("{}", auth_uri);
      86              : 
      87              :         // Convert the HTTP request into a signable request
      88            0 :         let signable_request = SignableRequest::new(
      89            0 :             "GET",
      90            0 :             auth_uri.to_string(),
      91            0 :             std::iter::empty(),
      92            0 :             SignableBody::Bytes(&[]),
      93            0 :         )?;
      94              : 
      95              :         // Sign and then apply the signature to the request
      96            0 :         let (si, _) = http_request::sign(signable_request, &signing_params)?.into_parts();
      97            0 :         let mut signable_request = http::Request::builder()
      98            0 :             .method("GET")
      99            0 :             .uri(auth_uri)
     100            0 :             .body(())?;
     101            0 :         si.apply_to_request_http1x(&mut signable_request);
     102            0 :         Ok((
     103            0 :             self.config.user_id.clone(),
     104            0 :             signable_request
     105            0 :                 .uri()
     106            0 :                 .to_string()
     107            0 :                 .replacen("http://", "", 1),
     108            0 :         ))
     109            0 :     }
     110              : }