LCOV - code coverage report
Current view: top level - proxy/src/redis - connection_with_credentials_provider.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 0.0 % 166 0
Test Date: 2025-03-12 00:01:28 Functions: 0.0 % 30 0

            Line data    Source code
       1              : use std::sync::Arc;
       2              : use std::time::Duration;
       3              : 
       4              : use futures::FutureExt;
       5              : use redis::aio::{ConnectionLike, MultiplexedConnection};
       6              : use redis::{ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult};
       7              : use tokio::task::JoinHandle;
       8              : use tracing::{debug, error, info, warn};
       9              : 
      10              : use super::elasticache::CredentialsProvider;
      11              : 
      12              : enum Credentials {
      13              :     Static(ConnectionInfo),
      14              :     Dynamic(Arc<CredentialsProvider>, redis::ConnectionAddr),
      15              : }
      16              : 
      17              : impl Clone for Credentials {
      18            0 :     fn clone(&self) -> Self {
      19            0 :         match self {
      20            0 :             Credentials::Static(info) => Credentials::Static(info.clone()),
      21            0 :             Credentials::Dynamic(provider, addr) => {
      22            0 :                 Credentials::Dynamic(Arc::clone(provider), addr.clone())
      23              :             }
      24              :         }
      25            0 :     }
      26              : }
      27              : 
      28              : /// A wrapper around `redis::MultiplexedConnection` that automatically refreshes the token.
      29              : /// Provides PubSub connection without credentials refresh.
      30              : pub struct ConnectionWithCredentialsProvider {
      31              :     credentials: Credentials,
      32              :     // TODO: with more load on the connection, we should consider using a connection pool
      33              :     con: Option<MultiplexedConnection>,
      34              :     refresh_token_task: Option<JoinHandle<()>>,
      35              :     mutex: tokio::sync::Mutex<()>,
      36              : }
      37              : 
      38              : impl Clone for ConnectionWithCredentialsProvider {
      39            0 :     fn clone(&self) -> Self {
      40            0 :         Self {
      41            0 :             credentials: self.credentials.clone(),
      42            0 :             con: None,
      43            0 :             refresh_token_task: None,
      44            0 :             mutex: tokio::sync::Mutex::new(()),
      45            0 :         }
      46            0 :     }
      47              : }
      48              : 
      49              : impl ConnectionWithCredentialsProvider {
      50            0 :     pub fn new_with_credentials_provider(
      51            0 :         host: String,
      52            0 :         port: u16,
      53            0 :         credentials_provider: Arc<CredentialsProvider>,
      54            0 :     ) -> Self {
      55            0 :         Self {
      56            0 :             credentials: Credentials::Dynamic(
      57            0 :                 credentials_provider,
      58            0 :                 redis::ConnectionAddr::TcpTls {
      59            0 :                     host,
      60            0 :                     port,
      61            0 :                     insecure: false,
      62            0 :                     tls_params: None,
      63            0 :                 },
      64            0 :             ),
      65            0 :             con: None,
      66            0 :             refresh_token_task: None,
      67            0 :             mutex: tokio::sync::Mutex::new(()),
      68            0 :         }
      69            0 :     }
      70              : 
      71            0 :     pub fn new_with_static_credentials<T: IntoConnectionInfo>(params: T) -> Self {
      72            0 :         Self {
      73            0 :             credentials: Credentials::Static(
      74            0 :                 params
      75            0 :                     .into_connection_info()
      76            0 :                     .expect("static configured redis credentials should be a valid format"),
      77            0 :             ),
      78            0 :             con: None,
      79            0 :             refresh_token_task: None,
      80            0 :             mutex: tokio::sync::Mutex::new(()),
      81            0 :         }
      82            0 :     }
      83              : 
      84            0 :     async fn ping(con: &mut MultiplexedConnection) -> RedisResult<()> {
      85            0 :         redis::cmd("PING").query_async(con).await
      86            0 :     }
      87              : 
      88            0 :     pub(crate) async fn connect(&mut self) -> anyhow::Result<()> {
      89            0 :         let _guard = self.mutex.lock().await;
      90            0 :         if let Some(con) = self.con.as_mut() {
      91            0 :             match Self::ping(con).await {
      92              :                 Ok(()) => {
      93            0 :                     return Ok(());
      94              :                 }
      95            0 :                 Err(e) => {
      96            0 :                     warn!("Error during PING: {e:?}");
      97              :                 }
      98              :             }
      99              :         } else {
     100            0 :             info!("Connection is not established");
     101              :         }
     102            0 :         info!("Establishing a new connection...");
     103            0 :         self.con = None;
     104            0 :         if let Some(f) = self.refresh_token_task.take() {
     105            0 :             f.abort();
     106            0 :         }
     107            0 :         let mut con = self
     108            0 :             .get_client()
     109            0 :             .await?
     110            0 :             .get_multiplexed_tokio_connection()
     111            0 :             .await?;
     112            0 :         if let Credentials::Dynamic(credentials_provider, _) = &self.credentials {
     113            0 :             let credentials_provider = credentials_provider.clone();
     114            0 :             let con2 = con.clone();
     115            0 :             let f = tokio::spawn(async move {
     116            0 :                 Self::keep_connection(con2, credentials_provider)
     117            0 :                     .await
     118            0 :                     .inspect_err(|e| debug!("keep_connection failed: {e}"))
     119            0 :                     .ok();
     120            0 :             });
     121            0 :             self.refresh_token_task = Some(f);
     122            0 :         }
     123            0 :         match Self::ping(&mut con).await {
     124              :             Ok(()) => {
     125            0 :                 info!("Connection succesfully established");
     126              :             }
     127            0 :             Err(e) => {
     128            0 :                 warn!("Connection is broken. Error during PING: {e:?}");
     129              :             }
     130              :         }
     131            0 :         self.con = Some(con);
     132            0 :         Ok(())
     133            0 :     }
     134              : 
     135            0 :     async fn get_connection_info(&self) -> anyhow::Result<ConnectionInfo> {
     136            0 :         match &self.credentials {
     137            0 :             Credentials::Static(info) => Ok(info.clone()),
     138            0 :             Credentials::Dynamic(provider, addr) => {
     139            0 :                 let (username, password) = provider.provide_credentials().await?;
     140            0 :                 Ok(ConnectionInfo {
     141            0 :                     addr: addr.clone(),
     142            0 :                     redis: RedisConnectionInfo {
     143            0 :                         db: 0,
     144            0 :                         username: Some(username),
     145            0 :                         password: Some(password.clone()),
     146            0 :                     },
     147            0 :                 })
     148              :             }
     149              :         }
     150            0 :     }
     151              : 
     152            0 :     async fn get_client(&self) -> anyhow::Result<redis::Client> {
     153            0 :         let client = redis::Client::open(self.get_connection_info().await?)?;
     154            0 :         Ok(client)
     155            0 :     }
     156              : 
     157              :     // PubSub does not support credentials refresh.
     158              :     // Requires manual reconnection every 12h.
     159            0 :     pub(crate) async fn get_async_pubsub(&self) -> anyhow::Result<redis::aio::PubSub> {
     160            0 :         Ok(self.get_client().await?.get_async_pubsub().await?)
     161            0 :     }
     162              : 
     163              :     // The connection lives for 12h.
     164              :     // It can be prolonged with sending `AUTH` commands with the refreshed token.
     165              :     // https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/auth-iam.html#auth-iam-limits
     166            0 :     async fn keep_connection(
     167            0 :         mut con: MultiplexedConnection,
     168            0 :         credentials_provider: Arc<CredentialsProvider>,
     169            0 :     ) -> anyhow::Result<()> {
     170              :         loop {
     171              :             // The connection lives for 12h, for the sanity check we refresh it every hour.
     172            0 :             tokio::time::sleep(Duration::from_secs(60 * 60)).await;
     173            0 :             match Self::refresh_token(&mut con, credentials_provider.clone()).await {
     174              :                 Ok(()) => {
     175            0 :                     info!("Token refreshed");
     176              :                 }
     177            0 :                 Err(e) => {
     178            0 :                     error!("Error during token refresh: {e:?}");
     179              :                 }
     180              :             }
     181              :         }
     182              :     }
     183            0 :     async fn refresh_token(
     184            0 :         con: &mut MultiplexedConnection,
     185            0 :         credentials_provider: Arc<CredentialsProvider>,
     186            0 :     ) -> anyhow::Result<()> {
     187            0 :         let (user, password) = credentials_provider.provide_credentials().await?;
     188            0 :         let _: () = redis::cmd("AUTH")
     189            0 :             .arg(user)
     190            0 :             .arg(password)
     191            0 :             .query_async(con)
     192            0 :             .await?;
     193            0 :         Ok(())
     194            0 :     }
     195              :     /// Sends an already encoded (packed) command into the TCP socket and
     196              :     /// reads the single response from it.
     197            0 :     pub(crate) async fn send_packed_command(
     198            0 :         &mut self,
     199            0 :         cmd: &redis::Cmd,
     200            0 :     ) -> RedisResult<redis::Value> {
     201              :         // Clone connection to avoid having to lock the ArcSwap in write mode
     202            0 :         let con = self.con.as_mut().ok_or(redis::RedisError::from((
     203            0 :             redis::ErrorKind::IoError,
     204            0 :             "Connection not established",
     205            0 :         )))?;
     206            0 :         con.send_packed_command(cmd).await
     207            0 :     }
     208              : 
     209              :     /// Sends multiple already encoded (packed) command into the TCP socket
     210              :     /// and reads `count` responses from it.  This is used to implement
     211              :     /// pipelining.
     212            0 :     pub(crate) async fn send_packed_commands(
     213            0 :         &mut self,
     214            0 :         cmd: &redis::Pipeline,
     215            0 :         offset: usize,
     216            0 :         count: usize,
     217            0 :     ) -> RedisResult<Vec<redis::Value>> {
     218              :         // Clone shared connection future to avoid having to lock the ArcSwap in write mode
     219            0 :         let con = self.con.as_mut().ok_or(redis::RedisError::from((
     220            0 :             redis::ErrorKind::IoError,
     221            0 :             "Connection not established",
     222            0 :         )))?;
     223            0 :         con.send_packed_commands(cmd, offset, count).await
     224            0 :     }
     225              : }
     226              : 
     227              : impl ConnectionLike for ConnectionWithCredentialsProvider {
     228            0 :     fn req_packed_command<'a>(
     229            0 :         &'a mut self,
     230            0 :         cmd: &'a redis::Cmd,
     231            0 :     ) -> redis::RedisFuture<'a, redis::Value> {
     232            0 :         (async move { self.send_packed_command(cmd).await }).boxed()
     233            0 :     }
     234              : 
     235            0 :     fn req_packed_commands<'a>(
     236            0 :         &'a mut self,
     237            0 :         cmd: &'a redis::Pipeline,
     238            0 :         offset: usize,
     239            0 :         count: usize,
     240            0 :     ) -> redis::RedisFuture<'a, Vec<redis::Value>> {
     241            0 :         (async move { self.send_packed_commands(cmd, offset, count).await }).boxed()
     242            0 :     }
     243              : 
     244            0 :     fn get_db(&self) -> i64 {
     245            0 :         0
     246            0 :     }
     247              : }
        

Generated by: LCOV version 2.1-beta