1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use super::cache_item::CacheItem;
use super::config::CacheConfig;
use aws_sdk_secretsmanager::error::GetSecretValueError;
use aws_sdk_secretsmanager::{Client as SecretsManagerClient, SdkError};
use lru::LruCache;

/// Client for in-process caching of secret values from AWS Secrets Manager.
///
/// An LRU (least-recently used) caching scheme is used that provides
/// O(1) insertions and O(1) lookups for cached values.
pub struct SecretCache {
    client: SecretsManagerClient,
    config: CacheConfig,
    cache: LruCache<String, CacheItem<String>>,
}

impl SecretCache {
    /// Returns a new SecretsCache using the default Cache Configuration options.
    pub fn new(client: SecretsManagerClient) -> Self {
        SecretCache::new_cache(client, CacheConfig::new())
    }

    /// Returns a new SecretsCache using a provided custom Cache Configuration.
    pub fn new_with_config(client: SecretsManagerClient, config: CacheConfig) -> Self {
        SecretCache::new_cache(client, config)
    }

    fn new_cache(client: SecretsManagerClient, config: CacheConfig) -> Self {
        let cache = LruCache::new(config.max_cache_size);
        Self {
            client,
            config,
            cache,
        }
    }

    /// Returns a builder for getting secret strings.
    ///
    /// Retrieve the secret value with send()
    pub fn get_secret_string(&mut self, secret_id: String) -> GetSecretStringBuilder {
        GetSecretStringBuilder::new(self, secret_id)
    }
}

/// A builder for the get_secret_string method.
pub struct GetSecretStringBuilder<'a> {
    secret_cache: &'a mut SecretCache,
    secret_id: String,
    force_refresh: bool,
}

impl<'a> GetSecretStringBuilder<'a> {
    pub fn new(secret_cache: &'a mut SecretCache, secret_id: String) -> Self {
        GetSecretStringBuilder {
            secret_cache,
            secret_id,
            force_refresh: false,
        }
    }

    /// Forces a refresh of the secret.
    ///
    /// Forces the secret to be fetched from AWS and updates the cache with the fresh value.
    /// This is required when the cached secret is out of date but not expired, for example due to rotation.
    pub fn force_refresh(mut self) -> Self {
        self.force_refresh = true;
        self
    }

    /// Fetches the secret value from the cache.
    ///
    /// If the secret value exists in the cache and hasn't expired it will be immediately returned.
    /// The secret will be fetched by calling AWS Secrets Manager and updated in the cache if:
    /// - the secret value hasn't been stored in the cache
    /// - the secret stored in the cache but has expired
    /// - the force_refresh option was provided
    ///
    /// Values are stored in the cache with the cache_item_ttl from the CacheConfig.
    pub async fn send(&mut self) -> Result<String, SdkError<GetSecretValueError>> {
        if !self.force_refresh {
            if let Some(cache_item) = self.secret_cache.cache.get(&self.secret_id) {
                if !cache_item.is_expired() {
                    return Ok(cache_item.value.clone());
                }
            }
        }

        match self.fetch_secret().await {
            Ok(secret_value) => {
                let cache_item = CacheItem::new(
                    secret_value.clone(),
                    self.secret_cache.config.cache_item_ttl,
                );
                self.secret_cache
                    .cache
                    .put(self.secret_id.clone(), cache_item);
                Ok(secret_value)
            }
            Err(e) => Err(e),
        }
    }

    async fn fetch_secret(&mut self) -> Result<String, SdkError<GetSecretValueError>> {
        match self
            .secret_cache
            .client
            .get_secret_value()
            .secret_id(self.secret_id.clone())
            .version_stage(self.secret_cache.config.version_stage.clone())
            .send()
            .await
        {
            Ok(resp) => return Ok(resp.secret_string.as_deref().unwrap().to_string()),
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aws_sdk_secretsmanager::{Client as SecretsManagerClient, Config, Credentials, Region};

    #[test]
    fn get_secret_string_builder_defaults() {
        let mock_secrets_manager_client = get_mock_secretsmanager_client();
        let mut secret_cache = SecretCache::new(mock_secrets_manager_client);

        let builder = GetSecretStringBuilder::new(&mut secret_cache, "service/secret".to_string());

        assert_eq!(builder.secret_id, "service/secret");
        assert!(!builder.force_refresh);
    }

    #[test]
    fn get_secret_string_builder_force_refresh() {
        let mock_secrets_manager_client = get_mock_secretsmanager_client();
        let mut secret_cache = SecretCache::new(mock_secrets_manager_client);

        let builder = GetSecretStringBuilder::new(&mut secret_cache, "service/secret".to_string())
            .force_refresh();

        assert_eq!(builder.secret_id, "service/secret");
        assert!(builder.force_refresh);
    }

    // provides a mocked AWS SecretsManager client for testing
    fn get_mock_secretsmanager_client() -> SecretsManagerClient {
        let conf = Config::builder()
            .region(Region::new("ap-southeast-2"))
            .credentials_provider(Credentials::new("asdf", "asdf", None, None, "test"))
            .build();

        SecretsManagerClient::from_conf(conf)
    }
}