Add Content-Type header to Agent responses

aws_secretsmanager_agent/src/config.rs

@@ -180,8 +180,7 @@
     ///
     /// # Returns
     ///
-    /// * `log_to_file` - `true` if writing logs to a file (default), `false` if writing logs to
-    /// stdout/stderr
+    /// * `log_to_file` - `true` if writing logs to a file (default), `false` if writing logs to stdout/stderr
     pub fn log_to_file(&self) -> bool {
         self.log_to_file
     }

aws_secretsmanager_agent/src/constants.rs

@@ -13,7 +13,7 @@ pub const EMPTY_ENV_LIST_MSG: &str =
 pub const BAD_PREFIX_MSG: &str =
     "The path prefix specified in the configuration file must begin with /.";
 
-/// Other constants that are used across the code base.
+// Other constants that are used across the code base.
 
 // The application name.
 pub const APPNAME: &str = "aws-secrets-manager-agent";

aws_secretsmanager_agent/src/main.rs

@@ -744,7 +744,7 @@ mod tests {
     #[tokio::test]
     async fn path_refresh_success() {
         let req = "/v1/My/Test?versionStage=AWSPENDING&refreshNow=0";
-        let (status, body) = run_request(&req).await;
+        let (status, body) = run_request(req).await;
         assert_eq!(status, StatusCode::OK);
         validate_response_extra("My/Test", DEFAULT_VERSION, vec!["AWSPENDING"], body);
     }

aws_secretsmanager_agent/src/parse.rs

@@ -98,7 +98,7 @@ mod tests {
         assert_eq!(query.secret_id, secret_id);
         assert_eq!(query.version_id, None);
         assert_eq!(query.version_stage, None);
-        assert_eq!(query.refresh_now, false);
+        assert!(!query.refresh_now);
     }
 
     #[test]
@@ -113,7 +113,7 @@ mod tests {
         assert_eq!(query.secret_id, secret_id);
         assert_eq!(query.version_id, None);
         assert_eq!(query.version_stage, None);
-        assert_eq!(query.refresh_now, true);
+        assert!(query.refresh_now);
     }
 
     #[test]
@@ -128,7 +128,7 @@ mod tests {
         assert_eq!(query.secret_id, secret_id);
         assert_eq!(query.version_id, None);
         assert_eq!(query.version_stage, None);
-        assert_eq!(query.refresh_now, false);
+        assert!(!query.refresh_now);
     }
 
     #[test]
@@ -160,7 +160,7 @@ mod tests {
         assert_eq!(query.secret_id, secret_id);
         assert_eq!(query.version_id, None);
         assert_eq!(query.version_stage, None);
-        assert_eq!(query.refresh_now, false);
+        assert!(!query.refresh_now);
     }
 
     #[test]

aws_secretsmanager_agent/src/server.rs

@@ -27,6 +27,20 @@
     max_conn: usize,
 }
 
+/// HTTP response relevant fields
+#[derive(Debug)]
+struct ResponseContent {
+    rsp_body: String,
+    content_type: ContentType,
+}
+
+/// Used to set Content-Type header
+#[derive(Debug)]
+enum ContentType {
+    Plain,
+    Json,
+}
+
 /// Handle incoming HTTP requests.
 ///
 /// Implements the HTTP handler. Each incomming request is handled in its own
@@ -108,11 +122,22 @@
 
         // Format the response.
         match result {
-            Ok(rsp_body) => Ok(Response::builder()
+            Ok(ResponseContent {
+                rsp_body,
+                content_type,
+            }) => Ok(Response::builder()
+                .header(
+                    "Content-Type",
+                    match content_type {
+                        ContentType::Plain => "text/plain",
+                        ContentType::Json => "application/json",
+                    },
+                )
                 .body(Full::new(Bytes::from(rsp_body)))
                 .unwrap()),
             Err(e) => Ok(Response::builder()
                 .status(e.0)
+                .header("Content-Type", "text/plain")
                 .body(Full::new(Bytes::from(e.1)))
                 .unwrap()),
         }
@@ -134,40 +159,49 @@
         &self,
         req: &Request<IncomingBody>,
         count: usize,
-    ) -> Result<String, HttpError> {
+    ) -> Result<ResponseContent, HttpError> {
         self.validate_max_conn(req, count)?; // Verify connection limits are not exceeded
         self.validate_token(req)?; // Check for a valid SSRF token
         self.validate_method(req)?; // Allow only GET requests
 
         match req.uri().path() {
-            "/ping" => Ok("healthy".into()), // Standard health check
+            "/ping" => Ok(ResponseContent {
+                rsp_body: "healthy".into(),
+                content_type: ContentType::Plain,
+            }), // Standard health check
 
             // Lambda extension style query
             "/secretsmanager/get" => {
                 let qry = GSVQuery::try_from_query(&req.uri().to_string())?;
-                Ok(self
-                    .cache_mgr
-                    .fetch(
-                        &qry.secret_id,
-                        qry.version_id.as_deref(),
-                        qry.version_stage.as_deref(),
-                        qry.refresh_now,
-                    )
-                    .await?)
+                Ok(ResponseContent {
+                    rsp_body: self
+                        .cache_mgr
+                        .fetch(
+                            &qry.secret_id,
+                            qry.version_id.as_deref(),
+                            qry.version_stage.as_deref(),
+                            qry.refresh_now,
+                        )
+                        .await?,
+                    content_type: ContentType::Json,
+                })
             }
 
             // Path style request
             path if path.starts_with(self.path_prefix.as_str()) => {
                 let qry = GSVQuery::try_from_path_query(&req.uri().to_string(), &self.path_prefix)?;
-                Ok(self
-                    .cache_mgr
-                    .fetch(
-                        &qry.secret_id,
-                        qry.version_id.as_deref(),
-                        qry.version_stage.as_deref(),
-                        qry.refresh_now,
-                    )
-                    .await?)
+                Ok(ResponseContent {
+                    rsp_body: self
+                        .cache_mgr
+                        .fetch(
+                            &qry.secret_id,
+                            qry.version_id.as_deref(),
+                            qry.version_stage.as_deref(),
+                            qry.refresh_now,
+                        )
+                        .await?,
+                    content_type: ContentType::Json,
+                })
             }
             _ => Err(HttpError(404, "Not found".into())),
         }

aws_secretsmanager_caching/src/lib.rs

@@ -68,7 +68,7 @@ impl SecretsManagerCachingClient {
     /// use aws_secretsmanager_caching::SecretsManagerCachingClient;
     /// use std::num::NonZeroUsize;
     /// use std::time::Duration;
-
+    ///
     /// let asm_client = SecretsManagerClient::from_conf(
     /// Config::builder()
     ///     .behavior_version_latest()
@@ -141,15 +141,15 @@ impl SecretsManagerCachingClient {
     /// use std::num::NonZeroUsize;
     /// use std::time::Duration;
     /// use aws_config::{BehaviorVersion, Region};
-
+    ///
     /// let config = aws_config::load_defaults(BehaviorVersion::latest())
     /// .await
     /// .into_builder()
     /// .region(Region::from_static("us-west-2"))
     /// .build();
-
+    ///
     /// let asm_builder = aws_sdk_secretsmanager::config::Builder::from(&config);
-
+    ///
     /// let client = SecretsManagerCachingClient::from_builder(
     /// asm_builder,
     /// NonZeroUsize::new(1000).unwrap(),
@@ -206,9 +206,9 @@ impl SecretsManagerCachingClient {
                 );
             }
 
-            return Ok(self
+            return self
                 .refresh_secret_value(secret_id, version_id, version_stage, None)
-                .await?);
+                .await;
         }
 
         let read_lock = self.store.read().await;
@@ -424,12 +424,12 @@ impl SecretsManagerCachingClient {
     }
 
     #[cfg(debug_assertions)]
-    fn increment_counter(&self, counter: &AtomicU32) -> () {
+    fn increment_counter(&self, counter: &AtomicU32) {
         counter.fetch_add(1, Ordering::SeqCst);
     }
 
     #[cfg(debug_assertions)]
-    fn reset_counter(&self, counter: &AtomicU32) -> () {
+    fn reset_counter(&self, counter: &AtomicU32) {
         counter.store(0, Ordering::SeqCst);
     }
 

aws_secretsmanager_caching/src/secret_store/memory_store/mod.rs

@@ -97,7 +97,6 @@ impl SecretStore for MemoryStore {
 }
 
 /// Write the secret value to the store
-
 #[cfg(test)]
 mod tests {