Add Content-Type header to Agent responses

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/server.rs

@@ -27,6 +27,20 @@ pub struct Server {
     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 @@ impl Server {
 
         // 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 @@ impl Server {
         &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

@@ -51,7 +51,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()
@@ -118,15 +118,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(),
@@ -164,9 +164,9 @@ impl SecretsManagerCachingClient {
         refresh_now: bool,
     ) -> Result<GetSecretValueOutputDef, Box<dyn Error>> {
         if refresh_now {
-            return Ok(self
+            return self
                 .refresh_secret_value(secret_id, version_id, version_stage, None)
-                .await?);
+                .await;
         }
 
         let read_lock = self.store.read().await;