1use crate::{DbErr, DbPool, ROOT_DATABASE_NAME, ROOT_DATABASE_ROLE_NAME, models::tenant::Tenant};
29use aws_config::SdkConfig;
30use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError};
31use aws_sigv4::{
32 http_request::{SignableBody, SignableRequest, SigningError, SigningSettings, sign},
33 sign::v4::signing_params,
34};
35use docbox_secrets::{SecretManager, SecretManagerError};
36use moka::{future::Cache, policy::EvictionPolicy};
37use serde::{Deserialize, Serialize};
38use sqlx::{
39 PgPool,
40 postgres::{PgConnectOptions, PgPoolOptions},
41};
42use std::time::Duration;
43use std::{num::ParseIntError, str::ParseBoolError};
44use std::{sync::Arc, time::SystemTime};
45use thiserror::Error;
46use tokio::time::sleep;
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct DatabasePoolCacheConfig {
51 pub host: String,
53 pub port: u16,
55
56 pub root_secret_name: Option<String>,
59
60 #[serde(default)]
63 pub root_iam: bool,
64
65 pub max_connections: Option<u32>,
76
77 pub max_connections_root: Option<u32>,
89
90 pub acquire_timeout: Option<u64>,
95
96 pub idle_timeout: Option<u64>,
102
103 pub pool_timeout: Option<u64>,
108
109 pub cache_duration: Option<u64>,
114
115 pub cache_capacity: Option<u64>,
125
126 pub credentials_cache_duration: Option<u64>,
132
133 pub credentials_cache_capacity: Option<u64>,
138}
139
140impl Default for DatabasePoolCacheConfig {
141 fn default() -> Self {
142 Self {
143 host: Default::default(),
144 port: 5432,
145 root_secret_name: Default::default(),
146 root_iam: false,
147 max_connections: None,
148 max_connections_root: None,
149 acquire_timeout: None,
150 idle_timeout: None,
151 pool_timeout: None,
152 cache_duration: None,
153 cache_capacity: None,
154 credentials_cache_duration: None,
155 credentials_cache_capacity: None,
156 }
157 }
158}
159
160#[derive(Debug, Error)]
161pub enum DatabasePoolCacheConfigError {
162 #[error("missing DOCBOX_DB_HOST environment variable")]
163 MissingDatabaseHost,
164 #[error("missing DOCBOX_DB_PORT environment variable")]
165 MissingDatabasePort,
166 #[error("invalid DOCBOX_DB_PORT environment variable")]
167 InvalidDatabasePort,
168 #[error("missing DOCBOX_DB_CREDENTIAL_NAME environment variable")]
169 MissingDatabaseSecretName,
170 #[error("invalid DOCBOX_DB_POOL_TIMEOUT environment variable")]
171 InvalidPoolTimeout(ParseIntError),
172 #[error("invalid DOCBOX_DB_IDLE_TIMEOUT environment variable")]
173 InvalidIdleTimeout(ParseIntError),
174 #[error("invalid DOCBOX_DB_ACQUIRE_TIMEOUT environment variable")]
175 InvalidAcquireTimeout(ParseIntError),
176 #[error("invalid DOCBOX_DB_CACHE_DURATION environment variable")]
177 InvalidCacheDuration(ParseIntError),
178 #[error("invalid DOCBOX_DB_CACHE_CAPACITY environment variable")]
179 InvalidCacheCapacity(ParseIntError),
180 #[error("invalid DOCBOX_DB_CREDENTIALS_CACHE_DURATION environment variable")]
181 InvalidCredentialsCacheDuration(ParseIntError),
182 #[error("invalid DOCBOX_DB_CREDENTIALS_CACHE_CAPACITY environment variable")]
183 InvalidCredentialsCacheCapacity(ParseIntError),
184 #[error("invalid DOCBOX_DB_ROOT_IAM environment variable")]
185 InvalidRootIam(ParseBoolError),
186}
187
188impl DatabasePoolCacheConfig {
189 pub fn from_env() -> Result<DatabasePoolCacheConfig, DatabasePoolCacheConfigError> {
190 let db_host: String = std::env::var("DOCBOX_DB_HOST")
191 .or(std::env::var("POSTGRES_HOST"))
192 .map_err(|_| DatabasePoolCacheConfigError::MissingDatabaseHost)?;
193 let db_port: u16 = std::env::var("DOCBOX_DB_PORT")
194 .or(std::env::var("POSTGRES_PORT"))
195 .map_err(|_| DatabasePoolCacheConfigError::MissingDatabasePort)?
196 .parse()
197 .map_err(|_| DatabasePoolCacheConfigError::InvalidDatabasePort)?;
198
199 let db_root_secret_name = std::env::var("DOCBOX_DB_CREDENTIAL_NAME").ok();
200 let db_root_iam = std::env::var("DOCBOX_DB_ROOT_IAM")
201 .ok()
202 .map(|value| value.parse::<bool>())
203 .transpose()
204 .map_err(DatabasePoolCacheConfigError::InvalidRootIam)?
205 .unwrap_or_default();
206
207 if !db_root_iam && db_root_secret_name.is_none() {
209 return Err(DatabasePoolCacheConfigError::MissingDatabaseSecretName);
210 }
211
212 let max_connections: Option<u32> = std::env::var("DOCBOX_DB_MAX_CONNECTIONS")
213 .ok()
214 .and_then(|value| value.parse().ok());
215 let max_connections_root: Option<u32> = std::env::var("DOCBOX_DB_MAX_ROOT_CONNECTIONS")
216 .ok()
217 .and_then(|value| value.parse().ok());
218
219 let acquire_timeout: Option<u64> = match std::env::var("DOCBOX_DB_ACQUIRE_TIMEOUT") {
220 Ok(value) => Some(
221 value
222 .parse::<u64>()
223 .map_err(DatabasePoolCacheConfigError::InvalidAcquireTimeout)?,
224 ),
225 Err(_) => None,
226 };
227
228 let pool_timeout: Option<u64> = match std::env::var("DOCBOX_DB_POOL_TIMEOUT") {
229 Ok(value) => Some(
230 value
231 .parse::<u64>()
232 .map_err(DatabasePoolCacheConfigError::InvalidPoolTimeout)?,
233 ),
234 Err(_) => None,
235 };
236
237 let idle_timeout: Option<u64> = match std::env::var("DOCBOX_DB_IDLE_TIMEOUT") {
238 Ok(value) => Some(
239 value
240 .parse::<u64>()
241 .map_err(DatabasePoolCacheConfigError::InvalidIdleTimeout)?,
242 ),
243 Err(_) => None,
244 };
245
246 let cache_duration: Option<u64> = match std::env::var("DOCBOX_DB_CACHE_DURATION") {
247 Ok(value) => Some(
248 value
249 .parse::<u64>()
250 .map_err(DatabasePoolCacheConfigError::InvalidCacheDuration)?,
251 ),
252 Err(_) => None,
253 };
254
255 let cache_capacity: Option<u64> = match std::env::var("DOCBOX_DB_CACHE_CAPACITY") {
256 Ok(value) => Some(
257 value
258 .parse::<u64>()
259 .map_err(DatabasePoolCacheConfigError::InvalidCacheCapacity)?,
260 ),
261 Err(_) => None,
262 };
263
264 let credentials_cache_duration: Option<u64> =
265 match std::env::var("DOCBOX_DB_CREDENTIALS_CACHE_DURATION") {
266 Ok(value) => Some(
267 value
268 .parse::<u64>()
269 .map_err(DatabasePoolCacheConfigError::InvalidCredentialsCacheDuration)?,
270 ),
271 Err(_) => None,
272 };
273
274 let credentials_cache_capacity: Option<u64> =
275 match std::env::var("DOCBOX_DB_CREDENTIALS_CACHE_CAPACITY") {
276 Ok(value) => Some(
277 value
278 .parse::<u64>()
279 .map_err(DatabasePoolCacheConfigError::InvalidCredentialsCacheCapacity)?,
280 ),
281 Err(_) => None,
282 };
283
284 Ok(DatabasePoolCacheConfig {
285 host: db_host,
286 port: db_port,
287 root_iam: db_root_iam,
288 root_secret_name: db_root_secret_name,
289 max_connections,
290 max_connections_root,
291 acquire_timeout,
292 pool_timeout,
293 idle_timeout,
294 cache_duration,
295 cache_capacity,
296 credentials_cache_duration,
297 credentials_cache_capacity,
298 })
299 }
300}
301
302pub struct DatabasePoolCache {
304 aws_config: aws_config::SdkConfig,
306
307 host: String,
309
310 port: u16,
312
313 root_secret_name: Option<String>,
318
319 root_iam: bool,
322
323 cache: Cache<String, DbPool>,
325
326 connect_info_cache: Cache<String, DbSecrets>,
329
330 secrets_manager: SecretManager,
332
333 max_connections: u32,
335 max_connections_root: u32,
337
338 acquire_timeout: Duration,
339 idle_timeout: Duration,
340}
341
342#[derive(Debug, Clone, Serialize, Deserialize)]
344pub struct DbSecrets {
345 pub username: String,
346 pub password: String,
347}
348
349#[derive(Debug, Error)]
350pub enum DbConnectErr {
351 #[error("database credentials not found in secrets manager")]
352 MissingCredentials,
353
354 #[error(transparent)]
355 SecretsManager(Box<SecretManagerError>),
356
357 #[error(transparent)]
358 Db(#[from] DbErr),
359
360 #[error(transparent)]
361 Shared(#[from] Arc<DbConnectErr>),
362
363 #[error("missing aws credentials provider")]
364 MissingCredentialsProvider,
365
366 #[error("failed to provide aws credentials")]
367 AwsCredentials(#[from] CredentialsError),
368
369 #[error("aws configuration missing region")]
370 MissingRegion,
371
372 #[error("failed to build aws signature")]
373 AwsSigner(#[from] signing_params::BuildError),
374
375 #[error("failed to sign aws request")]
376 AwsRequestSign(#[from] SigningError),
377
378 #[error("failed to parse signed aws url")]
379 AwsSignerInvalidUrl(url::ParseError),
380
381 #[error("failed to connect to tenant missing both IAM and secrets fields")]
382 InvalidTenantConfiguration,
383}
384
385impl DatabasePoolCache {
386 pub fn from_config(
387 aws_config: aws_config::SdkConfig,
388 config: DatabasePoolCacheConfig,
389 secrets_manager: SecretManager,
390 ) -> Self {
391 let mut pool_timeout = Duration::from_secs(config.cache_duration.unwrap_or(60 * 60 * 48));
392 let cache_duration = Duration::from_secs(config.cache_duration.unwrap_or(60 * 60 * 48));
393 let credentials_cache_duration =
394 Duration::from_secs(config.credentials_cache_duration.unwrap_or(60 * 60 * 12));
395
396 if config.root_iam && config.pool_timeout.is_none() {
399 tracing::debug!(
400 "IAM database auth is enabled with no pool timeout, setting short pool timeout within token duration"
401 );
402 pool_timeout = Duration::from_secs(60 * 10);
403 }
404
405 let cache_capacity = config.cache_capacity.unwrap_or(50);
406 let credentials_cache_capacity = config.credentials_cache_capacity.unwrap_or(50);
407
408 let cache = Cache::builder()
409 .time_to_live(pool_timeout)
410 .time_to_idle(cache_duration)
411 .max_capacity(cache_capacity)
412 .eviction_policy(EvictionPolicy::tiny_lfu())
413 .async_eviction_listener(|cache_key: Arc<String>, pool: DbPool, _cause| {
414 Box::pin(async move {
415 tracing::debug!(?cache_key, "database pool is no longer in use, closing");
416 pool.close().await
417 })
418 })
419 .build();
420
421 let connect_info_cache = Cache::builder()
422 .time_to_idle(credentials_cache_duration)
423 .max_capacity(credentials_cache_capacity)
424 .eviction_policy(EvictionPolicy::tiny_lfu())
425 .build();
426
427 Self {
428 aws_config,
429 host: config.host,
430 port: config.port,
431 root_secret_name: config.root_secret_name,
432 root_iam: config.root_iam,
433 cache,
434 connect_info_cache,
435 secrets_manager,
436 max_connections: config.max_connections.unwrap_or(10),
437 max_connections_root: config.max_connections_root.unwrap_or(2),
438 idle_timeout: Duration::from_secs(config.idle_timeout.unwrap_or(60 * 10)),
439 acquire_timeout: Duration::from_secs(config.acquire_timeout.unwrap_or(60)),
440 }
441 }
442
443 pub async fn get_root_pool(&self) -> Result<PgPool, DbConnectErr> {
445 match (self.root_secret_name.as_ref(), self.root_iam) {
446 (_, true) => {
447 self.get_pool_iam(ROOT_DATABASE_NAME, ROOT_DATABASE_ROLE_NAME)
448 .await
449 }
450
451 (Some(db_secret_name), _) => self.get_pool(ROOT_DATABASE_NAME, db_secret_name).await,
452
453 _ => Err(DbConnectErr::InvalidTenantConfiguration),
454 }
455 }
456
457 pub async fn get_tenant_pool(&self, tenant: &Tenant) -> Result<DbPool, DbConnectErr> {
459 match (
460 tenant.db_iam_user_name.as_ref(),
461 tenant.db_secret_name.as_ref(),
462 ) {
463 (Some(db_iam_user_name), _) => {
464 self.get_pool_iam(&tenant.db_name, db_iam_user_name).await
465 }
466 (_, Some(db_secret_name)) => self.get_pool(&tenant.db_name, db_secret_name).await,
467
468 _ => Err(DbConnectErr::InvalidTenantConfiguration),
469 }
470 }
471
472 pub async fn close_tenant_pool(&self, tenant: &Tenant) {
475 let cache_key = Self::tenant_cache_key(tenant);
476 if let Some(pool) = self.cache.remove(&cache_key).await {
477 pool.close().await;
478 }
479
480 self.cache.run_pending_tasks().await;
482 }
483
484 fn tenant_cache_key(tenant: &Tenant) -> String {
487 match (
488 tenant.db_secret_name.as_ref(),
489 tenant.db_iam_user_name.as_ref(),
490 ) {
491 (Some(db_secret_name), _) => {
492 format!("secret-{}-{}", &tenant.db_name, db_secret_name)
493 }
494 (_, Some(db_iam_user_name)) => {
495 format!("user-{}-{}", &tenant.db_name, db_iam_user_name)
496 }
497
498 _ => format!("db-{}", &tenant.db_name),
499 }
500 }
501
502 pub async fn flush(&self) {
504 self.cache.invalidate_all();
506 self.connect_info_cache.invalidate_all();
507 self.cache.run_pending_tasks().await;
508 }
509
510 pub async fn close_all(&self) {
512 for (_, value) in self.cache.iter() {
513 value.close().await;
514 }
515
516 self.flush().await;
517 }
518
519 async fn get_pool(&self, db_name: &str, secret_name: &str) -> Result<DbPool, DbConnectErr> {
522 let cache_key = format!("secret-{db_name}-{secret_name}");
523
524 let pool = self
525 .cache
526 .try_get_with(cache_key, async {
527 tracing::debug!(?db_name, "acquiring database pool");
528
529 let pool = self
530 .create_pool(db_name, secret_name)
531 .await
532 .map_err(Arc::new)?;
533
534 Ok(pool)
535 })
536 .await?;
537
538 Ok(pool)
539 }
540
541 async fn get_pool_iam(
544 &self,
545 db_name: &str,
546 db_role_name: &str,
547 ) -> Result<DbPool, DbConnectErr> {
548 let cache_key = format!("user-{db_name}-{db_role_name}");
549
550 let pool = self
551 .cache
552 .try_get_with(cache_key, async {
553 tracing::debug!(?db_name, "acquiring database pool (iam)");
554
555 let pool = self
556 .create_pool_iam(db_name, db_role_name)
557 .await
558 .map_err(Arc::new)?;
559
560 Ok(pool)
561 })
562 .await?;
563
564 Ok(pool)
565 }
566
567 async fn get_credentials(&self, secret_name: &str) -> Result<DbSecrets, DbConnectErr> {
569 if let Some(connect_info) = self.connect_info_cache.get(secret_name).await {
570 return Ok(connect_info);
571 }
572
573 let credentials = self
575 .secrets_manager
576 .parsed_secret::<DbSecrets>(secret_name)
577 .await
578 .map_err(|err| DbConnectErr::SecretsManager(Box::new(err)))?
579 .ok_or(DbConnectErr::MissingCredentials)?;
580
581 self.connect_info_cache
583 .insert(secret_name.to_string(), credentials.clone())
584 .await;
585
586 Ok(credentials)
587 }
588
589 async fn create_pool_iam(
591 &self,
592 db_name: &str,
593 db_role_name: &str,
594 ) -> Result<DbPool, DbConnectErr> {
595 tracing::debug!(?db_name, ?db_role_name, "creating db pool connection");
596
597 let options = iam_pool_connect_options(
598 &self.aws_config,
599 &self.host,
600 self.port,
601 db_name,
602 db_role_name,
603 )
604 .await?;
605
606 let max_connections = match db_name {
607 ROOT_DATABASE_NAME => self.max_connections_root,
608 _ => self.max_connections,
609 };
610
611 let pool = PgPoolOptions::new()
612 .max_connections(max_connections)
613 .acquire_timeout(self.acquire_timeout)
615 .idle_timeout(self.idle_timeout)
617 .connect_with(options)
618 .await
619 .map_err(DbConnectErr::Db)?;
620
621 tokio::spawn(iam_pool_maintenance_task(
622 pool.clone(),
623 self.aws_config.clone(),
624 self.host.clone(),
625 self.port,
626 db_name.to_string(),
627 db_role_name.to_string(),
628 ));
629
630 Ok(pool)
631 }
632
633 async fn create_pool(&self, db_name: &str, secret_name: &str) -> Result<DbPool, DbConnectErr> {
635 tracing::debug!(?db_name, ?secret_name, "creating db pool connection");
636
637 let credentials = self.get_credentials(secret_name).await?;
638 let options = PgConnectOptions::new()
639 .host(&self.host)
640 .port(self.port)
641 .username(&credentials.username)
642 .password(&credentials.password)
643 .database(db_name);
644
645 let max_connections = match db_name {
646 ROOT_DATABASE_NAME => self.max_connections_root,
647 _ => self.max_connections,
648 };
649
650 match PgPoolOptions::new()
651 .max_connections(max_connections)
652 .acquire_timeout(self.acquire_timeout)
654 .idle_timeout(self.idle_timeout)
656 .connect_with(options)
657 .await
658 {
659 Ok(value) => Ok(value),
661 Err(err) => {
662 self.connect_info_cache.remove(secret_name).await;
664 Err(DbConnectErr::Db(err))
665 }
666 }
667 }
668}
669
670async fn iam_pool_connect_options(
671 aws_config: &SdkConfig,
672 host: &str,
673 port: u16,
674 db_name: &str,
675 db_role_name: &str,
676) -> Result<PgConnectOptions, DbConnectErr> {
677 let token = create_rds_signed_token(aws_config, host, port, db_role_name).await?;
678
679 let options = PgConnectOptions::new()
680 .host(host)
681 .port(port)
682 .username(db_role_name)
683 .password(&token)
684 .database(db_name);
685
686 Ok(options)
687}
688
689async fn create_rds_signed_token(
690 aws_config: &SdkConfig,
691 host: &str,
692 port: u16,
693 user: &str,
694) -> Result<String, DbConnectErr> {
695 let credentials_provider = aws_config
696 .credentials_provider()
697 .ok_or(DbConnectErr::MissingCredentialsProvider)?;
698 let credentials = credentials_provider.provide_credentials().await?;
699 let identity = credentials.into();
700 let region = aws_config.region().ok_or(DbConnectErr::MissingRegion)?;
701
702 let mut signing_settings = SigningSettings::default();
703 signing_settings.expires_in = Some(Duration::from_secs(60 * 15));
704 signing_settings.signature_location = aws_sigv4::http_request::SignatureLocation::QueryParams;
705
706 let signing_params = aws_sigv4::sign::v4::SigningParams::builder()
707 .identity(&identity)
708 .region(region.as_ref())
709 .name("rds-db")
710 .time(SystemTime::now())
711 .settings(signing_settings)
712 .build()?;
713
714 let url = format!("https://{host}:{port}/?Action=connect&DBUser={user}");
715
716 let signable_request =
717 SignableRequest::new("GET", &url, std::iter::empty(), SignableBody::Bytes(&[]))?;
718
719 let (signing_instructions, _signature) =
720 sign(signable_request, &signing_params.into())?.into_parts();
721
722 let mut url = url::Url::parse(&url).map_err(DbConnectErr::AwsSignerInvalidUrl)?;
723 for (name, value) in signing_instructions.params() {
724 url.query_pairs_mut().append_pair(name, value);
725 }
726
727 let response = url.to_string().split_off("https://".len());
728 Ok(response)
729}
730
731async fn iam_pool_maintenance_task(
734 db: DbPool,
735 aws_config: SdkConfig,
736 host: String,
737 port: u16,
738 db_name: String,
739 db_role_name: String,
740) {
741 let interval = Duration::from_secs(60 * 10);
742
743 loop {
744 if db.is_closed() {
745 return;
746 }
747
748 match iam_pool_connect_options(&aws_config, &host, port, &db_name, &db_role_name).await {
749 Ok(options) => {
750 db.set_connect_options(options);
751 }
752 Err(error) => {
753 tracing::error!(?error, "failed to refresh IAM pool connect options");
754 }
755 }
756
757 sleep(interval).await;
758 }
759}