dfinity · lwshang · Jul 6, 2022 · Jul 5, 2022 · Jul 6, 2022 · Jul 6, 2022
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.3] - 2022-07-06
+### Added
+- Support for setting custom HTTP headers on asset creation 
+
 ## [0.2.2] - 2022-05-12
 ### Fixed
 - Parse and produce ETag headers with quotes around the hash

@@ -1,6 +1,6 @@
 [package]
 name = "ic-certified-assets"
-version = "0.2.2"
+version = "0.2.3"
 edition = "2021"
 authors = ["DFINITY Stiftung <sdk@dfinity.org>"]
 description = "Rust support for asset certification."

@@ -7,6 +7,7 @@ type CreateAssetArguments = record {
   key: Key;
   content_type: text;
   max_age: opt nat64;
+  headers: opt HeaderField;
 };
 
 // Add or change content for an asset, by content encoding

@@ -41,6 +41,7 @@ pub struct Asset {
     pub content_type: String,
     pub encodings: HashMap<String, AssetEncoding>,
     pub max_age: Option<u64>,
+    pub headers: Option<HashMap<String, String>>,
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
@@ -129,6 +130,7 @@ impl State {
                     content_type: arg.content_type,
                     encodings: HashMap::new(),
                     max_age: arg.max_age,
+                    headers: arg.headers,
                 },
             );
         }
@@ -788,6 +790,11 @@ fn build_ok(
     if let Some(max_age) = asset.max_age {
         headers.push(("Cache-Control".to_string(), format!("max-age={}", max_age)));
     }
+    if let Some(arg_headers) = asset.headers.as_ref() {
+        for (k, v) in arg_headers {
+            headers.push((k.to_owned(), v.to_owned()));
+        }
+    }
 
     let streaming_strategy = create_token(asset, enc_name, enc, key, chunk_index)
         .map(|token| StreamingStrategy::Callback { callback, token });

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use crate::state_machine::{StableState, State, BATCH_EXPIRY_NANOS};
 use crate::types::{
     BatchId, BatchOperation, CommitBatchArguments, CreateAssetArguments, CreateChunkArg,
@@ -24,6 +26,7 @@ struct AssetBuilder {
     content_type: String,
     max_age: Option<u64>,
     encodings: Vec<(String, Vec<ByteBuf>)>,
+    headers: Option<HashMap<String, String>>,
 }
 
 impl AssetBuilder {
@@ -33,6 +36,7 @@ impl AssetBuilder {
             content_type: content_type.as_ref().to_string(),
             max_age: None,
             encodings: vec![],
+            headers: None,
         }
     }
 
@@ -51,6 +55,12 @@ impl AssetBuilder {
         ));
         self
     }
+
+    fn with_header(mut self, header_key: &str, header_value: &str) -> Self {
+        let hm = self.headers.get_or_insert(HashMap::new());
+        hm.insert(header_key.to_string(), header_value.to_string());
+        self
+    }
 }
 
 struct RequestBuilder {
@@ -96,6 +106,7 @@ fn create_assets(state: &mut State, time_now: u64, assets: Vec<AssetBuilder>) ->
             key: asset.name.clone(),
             content_type: asset.content_type,
             max_age: asset.max_age,
+            headers: asset.headers,
         }));
 
         for (enc, chunks) in asset.encodings {
@@ -525,3 +536,73 @@ fn check_url_decode() {
     );
     assert_eq!(url_decode("/%e6"), Ok("/æ".to_string()));
 }
+
+#[test]
+fn supports_custom_http_headers() {
+    let mut state = State::default();
+    let time_now = 100_000_000_000;
+
+    const BODY: &[u8] = b"<!DOCTYPE html><html></html>";
+
+    create_assets(
+        &mut state,
+        time_now,
+        vec![
+            AssetBuilder::new("/contents.html", "text/html")
+                .with_encoding("identity", vec![BODY])
+                .with_header("Access-Control-Allow-Origin", "*"),
+            AssetBuilder::new("/max-age.html", "text/html")
+                .with_max_age(604800)
+                .with_encoding("identity", vec![BODY])
+                .with_header("X-Content-Type-Options", "nosniff"),
+        ],
+    );
+
+    let response = state.http_request(
+        RequestBuilder::get("/contents.html")
+            .with_header("Accept-Encoding", "gzip,identity")
+            .build(),
+        &[],
+        unused_callback(),
+    );
+
+    assert_eq!(response.status_code, 200);
+    assert_eq!(response.body.as_ref(), BODY);
+    assert!(
+        lookup_header(&response, "Access-Control-Allow-Origin").is_some(),
+        "Missing Access-Control-Allow-Origin header in response: {:#?}",
+        response,
+    );
+    assert!(
+        lookup_header(&response, "Access-Control-Allow-Origin") == Some("*"),
+        "Incorrect value for Access-Control-Allow-Origin header in response: {:#?}",
+        response,
+    );
+
+    let response = state.http_request(
+        RequestBuilder::get("/max-age.html")
+            .with_header("Accept-Encoding", "gzip,identity")
+            .build(),
+        &[],
+        unused_callback(),
+    );
+
+    assert_eq!(response.status_code, 200);
+    assert_eq!(response.body.as_ref(), BODY);
+    assert_eq!(
+        lookup_header(&response, "Cache-Control"),
+        Some("max-age=604800"),
+        "No matching Cache-Control header in response: {:#?}",
+        response,
+    );
+    assert!(
+        lookup_header(&response, "X-Content-Type-Options").is_some(),
+        "Missing X-Content-Type-Options header in response: {:#?}",
+        response,
+    );
+    assert!(
+        lookup_header(&response, "X-Content-Type-Options") == Some("nosniff"),
+        "Incorrect value for X-Content-Type-Options header in response: {:#?}",
+        response,
+    );
+}
@@ -1,5 +1,7 @@
 //! This module defines types shared by the certified assets state machine and the canister
 //! endpoints.
+use std::collections::HashMap;
+
 use crate::rc_bytes::RcBytes;
 use candid::{CandidType, Deserialize, Func, Nat};
 use serde_bytes::ByteBuf;
@@ -15,6 +17,7 @@ pub struct CreateAssetArguments {
     pub key: Key,
     pub content_type: String,
     pub max_age: Option<u64>,
+    pub headers: Option<HashMap<String, String>>,
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]