Merge pull request #49 from lazaroclapp/lazaro_add_param_support

Add support for SSM Parameter Store parameters

src/lib.rs

@@ -163,6 +163,22 @@ impl Manager {
             connection: self.connection.clone(),
         }
     }
+    /// Get a representation of a parameter that matches a given parameter name.
+    ///
+    /// For parameters of type `SecureString`, `with_decryption` must be set to `true.
+    /// Additionally, the lambda role must have the `kms:Decrypt` permission.
+    ///
+    /// Note that this does not return the value of the parameter; see [`Parameter`] for how to get it.
+    pub fn get_parameter(&self, param_name: &str, with_decryption: bool) -> Parameter {
+        Parameter {
+            query: format!(
+                "name={}&withDecryption={}",
+                param_name,
+                with_decryption
+            ),
+            connection: self.connection.clone(),
+        }
+    }
 }
 
 impl Default for Manager {
@@ -184,22 +200,35 @@ struct Connection {
 }
 
 impl Connection {
-    async fn get_secret(&self, query: &str) -> Result<String> {
-        Ok(self.client
-            .get(format!("http://localhost:{port}/secretsmanager/get?{query}", port = self.port))
+    async fn get_from_request(&self, url: &str) -> Result<reqwest::Response> {
+        self.client
+            .get(url)
             .header(TOKEN_HEADER_NAME, &self.token)
             .send()
             .await
             .context(
                 "could not communicate with the Secrets Manager extension (are you not running in AWS Lambda with the 'AWS-Parameters-and-Secrets-Lambda-Extension' version 2 layer?)"
             )?
             .error_for_status()
-            .context("received an error response from the Secrets Manager extension")?
-            .json::<ExtensionResponse>()
+            .context("received an error response from the Secrets Manager extension")
+    }
+
+    async fn get_secret(&self, query: &str) -> Result<String> {
+        let url = format!("http://localhost:{port}/secretsmanager/get?{query}", port = self.port);
+        Ok(self.get_from_request(&url).await?
+            .json::<ExtensionResponseSecret>()
             .await
             .context("invalid JSON received from Secrets Manager extension")?
             .secret_string)
     }
+
+    async fn get_parameter(&self, query: &str) -> Result<ExtensionResponseParam> {
+        let url = format!("http://localhost:{port}/systemsmanager/parameters/get?{query}", port = self.port);
+        self.get_from_request(&url).await?
+            .json::<ExtensionResponseParam>()
+            .await
+            .context("invalid JSON received from Secrets Manager extension")
+    }
 }
 
 /// A representation of a secret in Secrets Manager.
@@ -248,11 +277,83 @@ impl PartialEq for Secret {
 impl Eq for Secret {}
 
 #[derive(Deserialize)]
-struct ExtensionResponse {
+struct ExtensionResponseSecret {
     #[serde(rename = "SecretString")]
     secret_string: String,
 }
 
+/// A representation of a parameter in Parameter Store in SSM.
+#[derive(Debug, Clone)]
+pub struct Parameter {
+    query: String,
+    connection: Arc<Connection>,
+}
+
+impl Parameter {
+    /// Get the plaintext value of this parameter.
+    pub async fn get_raw(&self) -> Result<String> {
+        Ok(self.get_as_full_extension_response().await?.parameter.value)
+    }
+
+    /// Get the value of this parameter, represented as a strongly-typed T.
+    pub async fn get_typed<T: DeserializeOwned>(&self) -> Result<T> {
+        let raw = self.get_raw().await?;
+        Ok(serde_json::from_str(&raw)?)
+    }
+
+    /// Get the full response from the AWS lambda extension, including parameter type / version / ARN
+    /// info.
+    ///
+    /// Rarely used, see [`Self::get_raw()`] and [`Self::get_typed()`] for ways to retrieve string
+    /// and JSON parameters, respectively.
+    pub async fn get_as_full_extension_response(&self) -> Result<ExtensionResponseParam> {
+        self.connection.get_parameter(&self.query).await
+    }
+}
+
+impl PartialEq for Parameter {
+    fn eq(&self, other: &Self) -> bool {
+        self.query == other.query
+    }
+}
+
+impl Eq for Parameter {}
+
+/// The response from the AWS Lambda extension when successfully queried for a paramater at
+/// endpoint `/systemsmanager/parameters/get/?name=...`.
+#[derive(Deserialize)]
+pub struct ExtensionResponseParam {
+    /// The parameter returned.
+    #[serde(rename = "Parameter")]
+    pub parameter: ExtensionResponseParameterField
+}
+
+/// A parameter returned by the AWS Lambda extension, as structured JSON
+#[derive(Deserialize)]
+pub struct ExtensionResponseParameterField {
+    /// The parameter's ARN (Amazon Resource Name) full path.
+    #[serde(rename = "ARN")]
+    pub arn: String,
+    /// The data type of the parameter (e.g. text)
+    #[serde(rename = "DataType")]
+    pub data_type: String,
+    /// The date the parameter was last modified
+    #[serde(rename = "LastModifiedDate")]
+    pub last_modified_date: String,
+    /// The parameter's name.
+    #[serde(rename = "Name")]
+    pub name: String,
+    /// The date the parameter's type (e.g. `String`, `StringList`, or `SecureString`).
+    #[serde(rename = "Type")]
+    pub r#type: String,
+    /// The value of the parameter (this is the field that gets returned by [`Parameter::get_raw()`]).
+    #[serde(rename = "Value")]
+    pub value: String,
+    /// The date the parameter's version.
+    #[serde(rename = "Version")]
+    pub version: u64
+}
+
 /// A query for a specific [`Secret`] in AWS Secrets Manager. See [`Manager::get_secret`] for usage.
 /// 
 /// # Sealed
@@ -385,7 +486,11 @@ mod tests {
 
     use super::*;
 
+    const SECRETS_ENDPOINT: &'static str = "/secretsmanager/get";
+    const PARAMETERS_ENDPOINT: &'static str  = "/systemsmanager/parameters/get";
+
     struct MockServerConfig<'a> {
+        endpoint: &'a str,
         query: HashMap<&'a str, &'a str>,
         status: u16,
         response: &'a str,
@@ -395,7 +500,7 @@ mod tests {
         let server = MockServer::start();
 
         let mock = server.mock(|when, then| {
-            let mut when = when.method("GET").path("/secretsmanager/get");
+            let mut when = when.method("GET").path(config.endpoint);
 
             for (name, value) in config.query {
                 when = when.query_param(name, value);
@@ -411,6 +516,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_raw_secret() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{\"SecretString\": \"xyz\"}",
@@ -433,6 +539,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_raw_secret_from_version_id() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret", "versionId" => "some-version"},
             status: 200,
             response: "{\"SecretString\": \"xyz\"}",
@@ -459,6 +566,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_raw_secret_from_version_stage() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret", "versionStage" => "some-stage"},
             status: 200,
             response: "{\"SecretString\": \"xyz\"}",
@@ -485,6 +593,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_single_secret() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{\"SecretString\": \"{\\\"name\\\": \\\"value\\\"}\"}",
@@ -516,6 +625,7 @@ mod tests {
         }
 
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{\"SecretString\": \"{\\\"name\\\": \\\"value\\\"}\"}",
@@ -552,6 +662,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_invalid_json() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{",
@@ -601,6 +712,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_server_returns_non_200_status_code() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 500,
             response: "",
@@ -663,6 +775,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_single_secret_not_found() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{\"SecretString\": \"{}\"}",
@@ -692,6 +805,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager_get_single_secret_incorrect_type() {
         let config = MockServerConfig {
+            endpoint: SECRETS_ENDPOINT,
             query: hashmap! {"secretId" => "some-secret"},
             status: 200,
             response: "{\"SecretString\": \"{\\\"name\\\": 1}\"}",
@@ -717,4 +831,76 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn test_manager_get_ssm_raw_parameter() {
+        let config = MockServerConfig {
+            endpoint: PARAMETERS_ENDPOINT,
+            query: hashmap! {"name" => "/some/path/to/a/param", "withDecryption" => "false"},
+            status: 200,
+            response: "{
+                \"Parameter\": {
+                    \"ARN\": \"arn:aws:ssm:us-east-1:000000000000:parameter/some/path/to/a/param\",
+                    \"DataType\": \"text\",
+                    \"LastModifiedDate\": \"2024-03-01T17:53:36.314Z\",
+                    \"Name\": \"/some/path/to/a/param\",
+                    \"Selector\": null,
+                    \"SourceResult\": null,
+                    \"Type\": \"String\",
+                    \"Value\": \"Some param\",
+                    \"Version\": 1
+                },
+                \"ResultMetadata\": {}
+            }",
+        };
+
+        with_mock_server(config, |port| async move {
+            let manager = ManagerBuilder::new()
+                .with_port(port)
+                .with_token(String::from("TOKEN"))
+                .build()
+                .unwrap();
+
+            let param_value = manager.get_parameter("/some/path/to/a/param", false).get_raw().await.unwrap();
+
+            assert_eq!(String::from("Some param"), param_value);
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn test_manager_get_ssm_raw_parameter_secure_string() {
+        let config = MockServerConfig {
+            endpoint: PARAMETERS_ENDPOINT,
+            query: hashmap! {"name" => "/some/path/to/a/param", "withDecryption" => "true"},
+            status: 200,
+            response: "{
+                \"Parameter\": {
+                    \"ARN\": \"arn:aws:ssm:us-east-1:000000000000:parameter/some/path/to/a/param\",
+                    \"DataType\": \"text\",
+                    \"LastModifiedDate\": \"2024-03-01T17:53:36.314Z\",
+                    \"Name\": \"/some/path/to/a/param\",
+                    \"Selector\": null,
+                    \"SourceResult\": null,
+                    \"Type\": \"SecureString\",
+                    \"Value\": \"Some encrypted string (now decrypted)\",
+                    \"Version\": 1
+                },
+                \"ResultMetadata\": {}
+            }",
+        };
+
+        with_mock_server(config, |port| async move {
+            let manager = ManagerBuilder::new()
+                .with_port(port)
+                .with_token(String::from("TOKEN"))
+                .build()
+                .unwrap();
+
+            let param_value = manager.get_parameter("/some/path/to/a/param", true).get_raw().await.unwrap();
+
+            assert_eq!(String::from("Some encrypted string (now decrypted)"), param_value);
+        })
+        .await;
+    }
 }