Add support for custom headers in PUT requests (#1059)

* Add custom headers to put_object_single

Signed-off-by: Alessandro Passaro <[email protected]>

* Add custom headers to put_object

Signed-off-by: Alessandro Passaro <[email protected]>

---------

Signed-off-by: Alessandro Passaro <[email protected]>

mountpoint-s3-client/src/object_client.rs

@@ -316,6 +316,8 @@ pub struct PutObjectParams {
     /// If `server_side_encryption` has a valid value of aws:kms or aws:kms:dsse, this value may be used to specify AWS KMS key ID to be used
     /// when creating new S3 object
     pub ssekms_key_id: Option<String>,
+    /// Custom headers to add to the request
+    pub custom_headers: Vec<(String, String)>,
 }
 
 impl PutObjectParams {
@@ -347,6 +349,12 @@ impl PutObjectParams {
         self.ssekms_key_id = value;
         self
     }
+
+    /// Add a custom header to the request.
+    pub fn add_custom_header(mut self, name: String, value: String) -> Self {
+        self.custom_headers.push((name, value));
+        self
+    }
 }
 
 /// How CRC32c checksums are used for parts of a multi-part PutObject request
@@ -383,6 +391,8 @@ pub struct PutObjectSingleParams {
     /// If `server_side_encryption` has a valid value of aws:kms or aws:kms:dsse, this value may be used to specify AWS KMS key ID to be used
     /// when creating new S3 object
     pub ssekms_key_id: Option<String>,
+    /// Custom headers to add to the request
+    pub custom_headers: Vec<(String, String)>,
 }
 
 impl PutObjectSingleParams {
@@ -414,6 +424,12 @@ impl PutObjectSingleParams {
         self.ssekms_key_id = value;
         self
     }
+
+    /// Add a custom header to the request.
+    pub fn add_custom_header(mut self, name: String, value: String) -> Self {
+        self.custom_headers.push((name, value));
+        self
+    }
 }
 
 /// A checksum used by the object client for integrity checks on uploads.

mountpoint-s3-client/src/s3_crt_client/put_object.rs

@@ -42,6 +42,13 @@ impl S3CrtClient {
         };
         message.set_checksum_config(checksum_config);
 
+        for (name, value) in &params.custom_headers {
+            message
+                .inner
+                .add_header(&Header::new(name, value))
+                .map_err(S3RequestError::construction_failure)?;
+        }
+
         let review_callback = ReviewCallbackBox::default();
         let callback = review_callback.clone();
 
@@ -128,6 +135,12 @@ impl S3CrtClient {
                     .set_checksum_header(checksum)
                     .map_err(S3RequestError::construction_failure)?;
             }
+            for (name, value) in &params.custom_headers {
+                message
+                    .inner
+                    .add_header(&Header::new(name, value))
+                    .map_err(S3RequestError::construction_failure)?;
+            }
 
             let body_input_stream =
                 InputStream::new_from_slice(&self.inner.allocator, slice).map_err(S3RequestError::CrtError)?;

mountpoint-s3-client/tests/put_object.rs

@@ -624,3 +624,30 @@ async fn test_concurrent_put_objects(throughput_target_gbps: f64, max_concurrent
     // Cancel all put_object requests.
     drop(req_vec);
 }
+
+#[tokio::test]
+async fn test_put_object_header() {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_object_header");
+    let client = get_test_client();
+    let key = format!("{prefix}hello");
+
+    let content_type = "application/json";
+    let params = PutObjectParams::new().add_custom_header("Content-Type".to_owned(), content_type.to_owned());
+    let mut request = client
+        .put_object(&bucket, &key, &params)
+        .await
+        .expect("put_object should succeed");
+    request
+        .write(b"{ \"key\": \"value\" }")
+        .await
+        .expect("write should succeed");
+    request
+        .complete()
+        .await
+        .expect("the upload should complete successfully");
+
+    let sdk_client = get_test_sdk_client().await;
+    let output = sdk_client.head_object().bucket(bucket).key(key).send().await.unwrap();
+
+    assert_eq!(Some(content_type), output.content_type());
+}

mountpoint-s3-client/tests/put_object_single.rs

@@ -236,3 +236,22 @@ async fn test_put_object_sse(sse_type: Option<&str>, kms_key_id: Option<String>)
     let put_object_result = test_put_object_single(&client, &bucket, &key, request_params.clone()).await;
     check_sse(&bucket, &key, sse_type, &kms_key_id, put_object_result).await;
 }
+
+#[tokio::test]
+async fn test_put_object_header() {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_object_header");
+    let client = get_test_client();
+    let key = format!("{prefix}hello");
+
+    let content_type = "application/json";
+    let params = PutObjectSingleParams::new().add_custom_header("Content-Type".to_owned(), content_type.to_owned());
+    client
+        .put_object_single(&bucket, &key, &params, b"{ \"key\": \"value\" }")
+        .await
+        .expect("put_object should succeed");
+
+    let sdk_client = get_test_sdk_client().await;
+    let output = sdk_client.head_object().bucket(bucket).key(key).send().await.unwrap();
+
+    assert_eq!(Some(content_type), output.content_type());
+}