feat(ic-http-gateway): add more tests and checks for long asset handl…

…ing (#28)

- Add detection of out-of-order chunks in streaming responses, and
corresponding tests.
- Add sanity checks for `Content-Range` values returned with chunked
responses.

examples/http-gateway/canister/src/custom_assets/src/lib.rs

@@ -3,7 +3,9 @@ use ic_cdk::{
     api::{data_certificate, set_certified_data},
     *,
 };
-use ic_http_certification::{HeaderField, HttpRequest, HttpResponse, HttpResponseBuilder};
+use ic_http_certification::{
+    HeaderField, HttpRequest, HttpRequestBuilder, HttpResponse, HttpResponseBuilder,
+};
 use rand_chacha::rand_core::{RngCore, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 use std::cell::RefCell;
@@ -18,6 +20,13 @@ fn post_upgrade() {
     init();
 }
 
+// In addition to serving configured assets, the canister supports various
+// corruption scenarios for testing purposes.  Specifically, the caller
+// can use one of the following custom HTTP headers to make the canister
+// "misbehave" in various ways:
+// - "Test-CorruptChunkAtIndex"
+// - "Test-CorruptCertificateAtIndex"
+// - "Test-SwapChunkAtIndexWithNext"
 #[query]
 fn http_request(req: HttpRequest) -> HttpResponse {
     let mut response = serve_asset(&req);
@@ -58,6 +67,37 @@ fn http_request(req: HttpRequest) -> HttpResponse {
                 .build();
         }
     }
+    if let Some(chunk_to_swap) = chunk_swap_requested(&req) {
+        let current_chunk = current_chunk_index(&response);
+        if current_chunk == chunk_to_swap {
+            // Create a request for the next chunk.
+            let next_chunk_req = HttpRequestBuilder::new()
+                .with_method(req.method())
+                .with_url(req.url())
+                .with_body(req.body())
+                .with_certificate_version(req.certificate_version().unwrap_or(2))
+                .with_headers({
+                    let mut headers = req.headers().to_owned();
+                    let mut range_updated = false;
+                    let new_range_value =
+                        format!("bytes={}-", (chunk_to_swap + 1) * ASSET_CHUNK_SIZE);
+                    for (key, value) in headers.iter_mut() {
+                        if key == "Range" {
+                            value.clear();
+                            value.push_str(&new_range_value);
+                            range_updated = true;
+                        }
+                    }
+                    if !range_updated {
+                        // The request had no Range-header, insert one.
+                        headers.push(("Range".to_string(), new_range_value));
+                    }
+                    headers.to_vec()
+                })
+                .build();
+            response = serve_asset(&next_chunk_req);
+        }
+    }
     response
 }
 
@@ -71,7 +111,7 @@ fn current_chunk_index(resp: &HttpResponse) -> usize {
 }
 
 fn chunk_corruption_requested(req: &HttpRequest) -> Option<usize> {
-    if let Some(corrupted_chunk_index) = get_header_value(req.headers(), "Test-CorruptedChunkIndex")
+    if let Some(corrupted_chunk_index) = get_header_value(req.headers(), "Test-CorruptChunkAtIndex")
     {
         Some(
             corrupted_chunk_index
@@ -85,7 +125,7 @@ fn chunk_corruption_requested(req: &HttpRequest) -> Option<usize> {
 
 fn cert_corruption_requested(req: &HttpRequest) -> Option<usize> {
     if let Some(corrupted_cert_chunk_index) =
-        get_header_value(req.headers(), "Test-CorruptedCertificate")
+        get_header_value(req.headers(), "Test-CorruptCertificateAtIndex")
     {
         Some(
             corrupted_cert_chunk_index
@@ -97,6 +137,20 @@ fn cert_corruption_requested(req: &HttpRequest) -> Option<usize> {
     }
 }
 
+fn chunk_swap_requested(req: &HttpRequest) -> Option<usize> {
+    if let Some(chunk_index_to_swap) =
+        get_header_value(req.headers(), "Test-SwapChunkAtIndexWithNext")
+    {
+        Some(
+            chunk_index_to_swap
+                .parse()
+                .expect("invalid index of chunk to swap"),
+        )
+    } else {
+        None
+    }
+}
+
 fn get_header_value(headers: &[HeaderField], header_name: &str) -> Option<String> {
     for (name, value) in headers.iter() {
         if name.to_lowercase().eq(&header_name.to_lowercase()) {

packages/ic-http-gateway/src/response/response_handler.rs

@@ -156,7 +156,7 @@ pub async fn get_206_stream_response_body_and_total_length(
         .expect("missing streamed chunk body")
         .to_bytes()
         .to_vec();
-    let stream_state = get_stream_state(
+    let stream_state = get_initial_stream_state(
         http_request,
         canister_id,
         response_headers,
@@ -201,12 +201,19 @@ fn parse_content_range_header_str(str_value: &str) -> Result<ContentRangeValues,
         .as_str()
         .parse()
         .map_err(|_| AgentError::InvalidHttpResponse("malformed size".to_string()))?;
-    // TODO: add sanity checks for the parsed values
-    Ok(ContentRangeValues {
+    let range_values = ContentRangeValues {
         range_begin,
         range_end,
         total_length,
-    })
+    };
+    if range_begin > range_end || range_begin >= total_length || range_end >= total_length {
+        Err(AgentError::InvalidHttpResponse(format!(
+            "inconsistent Content-Range header {:?}",
+            range_values
+        )))
+    } else {
+        Ok(range_values)
+    }
 }
 
 fn get_content_range_header_str(
@@ -224,18 +231,33 @@ fn get_content_range_header_str(
 
 fn get_content_range_values(
     response_headers: &Vec<HeaderField<'static>>,
+    fetched_length: usize,
 ) -> Result<ContentRangeValues, AgentError> {
     let str_value = get_content_range_header_str(response_headers)?;
-    parse_content_range_header_str(&str_value)
+    let range_values = parse_content_range_header_str(&str_value)?;
+    // ic_cdk::println!("--- range values: {:?}", range_values);
+    if range_values.range_begin > fetched_length {
+        return Err(AgentError::InvalidHttpResponse(format!(
+            "chunk out-of-order: range_begin={} is larger than expected begin={} ",
+            range_values.range_begin, fetched_length
+        )));
+    }
+    if range_values.range_end < fetched_length {
+        return Err(AgentError::InvalidHttpResponse(format!(
+            "chunk out-of-order: range_end={} is smaller than length fetched so far={} ",
+            range_values.range_begin, fetched_length
+        )));
+    }
+    Ok(range_values)
 }
 
-fn get_stream_state<'a>(
+fn get_initial_stream_state<'a>(
     http_request: HttpRequest<'a>,
     canister_id: Principal,
     response_headers: &Vec<HeaderField<'static>>,
     skip_verification: bool,
 ) -> Result<StreamState<'a>, AgentError> {
-    let range_values = get_content_range_values(response_headers)?;
+    let range_values = get_content_range_values(response_headers, 0)?;
 
     Ok(StreamState {
         http_request,
@@ -299,10 +321,14 @@ fn create_206_stream(
                 Ok((response,)) => response,
                 Err(e) => return Err(e),
             };
-            let range_values = get_content_range_values(&agent_response.headers)?;
+            let range_values =
+                get_content_range_values(&agent_response.headers, stream_state.fetched_length)?;
+            let new_bytes_begin = stream_state
+                .fetched_length
+                .saturating_sub(range_values.range_begin);
             let chunk_length = range_values
                 .range_end
-                .saturating_sub(range_values.range_begin)
+                .saturating_sub(stream_state.fetched_length)
                 + 1;
             let current_fetched_length = stream_state.fetched_length + chunk_length;
             // Verify the chunk from the range response.
@@ -345,7 +371,10 @@ fn create_206_stream(
                 None
             };
             Ok(Some((
-                (agent_response.body, maybe_new_state.clone()),
+                (
+                    agent_response.body[new_bytes_begin..].to_vec(),
+                    maybe_new_state.clone(),
+                ),
                 (agent, maybe_new_state),
             )))
         },
@@ -404,18 +433,27 @@ mod tests {
     }
 
     #[test]
-    fn should_get_stream_state() {
+    fn should_fail_parse_content_range_header_str_on_inconsistent_input() {
+        let inconsistent_inputs = ["bytes 100-200/190", "bytes 200-150/400", "bytes 100-110/40"];
+        for input in inconsistent_inputs {
+            let result = parse_content_range_header_str(input);
+            assert_matches!(result, Err(e) if format!("{}", e).contains("inconsistent Content-Range header"));
+        }
+    }
+
+    #[test]
+    fn should_get_initial_stream_state() {
         let http_request = HttpRequest::get("http://example.com/some_file")
             .with_headers(vec![("Xyz".to_string(), "some value".to_string())])
             .with_body(vec![42])
             .build();
         let canister_id = Principal::from_slice(&[1, 2, 3, 4]);
         let response_headers = vec![HeaderField(
             Cow::from("Content-Range"),
-            Cow::from("bytes 2-4/10"), // fetched 3 bytes, total length is 10
+            Cow::from("bytes 0-2/10"), // fetched 3 bytes, total length is 10
         )];
         let skip_verification = false;
-        let state = get_stream_state(
+        let state = get_initial_stream_state(
             http_request.clone(),
             canister_id,
             &response_headers,
@@ -430,7 +468,7 @@ mod tests {
     }
 
     #[test]
-    fn should_fail_get_stream_state_without_content_range_header() {
+    fn should_fail_get_initial_stream_state_without_content_range_header() {
         let http_request = HttpRequest::get("http://example.com/some_file")
             .with_headers(vec![("Xyz".to_string(), "some value".to_string())])
             .with_body(vec![42])
@@ -440,12 +478,12 @@ mod tests {
             Cow::from("other header"),
             Cow::from("other value"),
         )];
-        let result = get_stream_state(http_request, canister_id, &response_headers, false);
+        let result = get_initial_stream_state(http_request, canister_id, &response_headers, false);
         assert_matches!(result, Err(e) if format!("{}", e).contains("missing Content-Range header"));
     }
 
     #[test]
-    fn should_fail_get_stream_state_with_malformed_content_range_header() {
+    fn should_fail_get_initial_stream_state_with_malformed_content_range_header() {
         let http_request = HttpRequest::get("http://example.com/some_file")
             .with_headers(vec![("Xyz".to_string(), "some value".to_string())])
             .with_body(vec![42])
@@ -455,7 +493,22 @@ mod tests {
             Cow::from("Content-Range"),
             Cow::from("bytes 42/10"),
         )];
-        let result = get_stream_state(http_request, canister_id, &response_headers, false);
+        let result = get_initial_stream_state(http_request, canister_id, &response_headers, false);
         assert_matches!(result, Err(e) if format!("{}", e).contains("malformed Content-Range header"));
     }
+
+    #[test]
+    fn should_fail_get_initial_stream_state_with_inconsistent_content_range_header() {
+        let http_request = HttpRequest::get("http://example.com/some_file")
+            .with_headers(vec![("Xyz".to_string(), "some value".to_string())])
+            .with_body(vec![42])
+            .build();
+        let canister_id = Principal::from_slice(&[1, 2, 3, 4]);
+        let response_headers = vec![HeaderField(
+            Cow::from("Content-Range"),
+            Cow::from("bytes 40-100/90"),
+        )];
+        let result = get_initial_stream_state(http_request, canister_id, &response_headers, false);
+        assert_matches!(result, Err(e) if format!("{}", e).contains("inconsistent Content-Range header"));
+    }
 }

packages/ic-http-gateway/tests/range_request_stream.rs

@@ -198,7 +198,7 @@ fn test_corrupted_long_asset_request_fails(
                 canister_id,
                 canister_request: Request::builder()
                     .header(
-                        "Test-CorruptedChunkIndex",
+                        "Test-CorruptChunkAtIndex",
                         corrupted_chunk_index.to_string(),
                     )
                     .uri(format!("/{asset_name}"))
@@ -218,10 +218,9 @@ fn test_corrupted_long_asset_request_fails(
         if corrupted_chunk_index == 0 {
             // If the first chunk is corrupted, the status indicates the failure
             // and the full body contains the error message.
-            let body = body_result.expect("failed getting full body").to_bytes();
-            assert_eq!(
-                body,
-                "Response verification failed: Invalid response hashes"
+            assert_matches!(body_result,
+                Ok(body) if format!("{:?}", body).contains(
+                "Response verification failed: Invalid response hashes")
             );
         } else {
             // If the first chunk is ok, but some other chunk is corrupted, the response has 200-status,
@@ -235,6 +234,77 @@ fn test_corrupted_long_asset_request_fails(
     });
 }
 
+#[rstest]
+#[case(TWO_CHUNKS_ASSET_NAME, 0)]
+#[case(SIX_CHUNKS_ASSET_NAME, 3)]
+fn test_long_asset_with_chunks_out_of_order_fails(
+    #[case] asset_name: &str,
+    #[case] chunk_to_swap: usize,
+) {
+    let rt = tokio::runtime::Runtime::new().unwrap();
+    let wasm_bytes = rt.block_on(async { utils::load_custom_assets_wasm().await });
+
+    let pic = PocketIcBuilder::new()
+        .with_nns_subnet()
+        .with_application_subnet()
+        .build();
+
+    let canister_id = pic.create_canister();
+    pic.add_cycles(canister_id, 2_000_000_000_000);
+    pic.install_canister(canister_id, wasm_bytes, vec![], None);
+
+    let url = pic.auto_progress();
+
+    let agent = Agent::builder().with_url(url).build().unwrap();
+    rt.block_on(async {
+        agent.fetch_root_key().await.unwrap();
+    });
+
+    let http_gateway = HttpGatewayClient::builder()
+        .with_agent(agent)
+        .build()
+        .unwrap();
+
+    let response = rt.block_on(async {
+        http_gateway
+            .request(HttpGatewayRequestArgs {
+                canister_id,
+                canister_request: Request::builder()
+                    .header("Test-SwapChunkAtIndexWithNext", chunk_to_swap.to_string())
+                    .uri(format!("/{asset_name}"))
+                    .body(vec![])
+                    .unwrap(),
+            })
+            .send()
+            .await
+    });
+    let expected_status = match chunk_to_swap {
+        0 => 500,
+        _ => 200,
+    };
+    assert_eq!(response.canister_response.status(), expected_status);
+    rt.block_on(async {
+        let body_result = response.canister_response.into_body().collect().await;
+        if chunk_to_swap == 0 {
+            // If the first chunk is corrupted, the status indicates the failure
+            // and the full body contains the error message.
+            assert_matches!(body_result,
+                Ok(body) if format!("{:?}", body).contains(&format!(
+                    "chunk out-of-order: range_begin={}",
+                    ASSET_CHUNK_SIZE*(chunk_to_swap+1)))
+            );
+        } else {
+            // If the first chunk is ok, but some other chunk is corrupted, the response has 200-status,
+            // but fetching the full body fails with an error for the corrupted chunk.
+            assert_matches!(body_result,
+                Err(e) if e.to_string().contains(&format!(
+                    "chunk out-of-order: range_begin={}",
+                    ASSET_CHUNK_SIZE*(chunk_to_swap+1)))
+            );
+        }
+    });
+}
+
 #[rstest]
 #[case(TWO_CHUNKS_ASSET_NAME, 0)]
 #[case(TWO_CHUNKS_ASSET_NAME, 1)]
@@ -274,7 +344,7 @@ fn test_corrupted_chunk_certificate_for_long_asset_request_fails(
                 canister_id,
                 canister_request: Request::builder()
                     .header(
-                        "Test-CorruptedCertificate",
+                        "Test-CorruptCertificateAtIndex",
                         corrupted_chunk_index.to_string(),
                     )
                     .uri(format!("/{asset_name}"))