diff --git a/examples/http-gateway/canister/src/custom_assets/src/lib.rs b/examples/http-gateway/canister/src/custom_assets/src/lib.rs index 4d85e79..27846ab 100644 --- a/examples/http-gateway/canister/src/custom_assets/src/lib.rs +++ b/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 { - 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 { fn cert_corruption_requested(req: &HttpRequest) -> Option { 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 { } } +fn chunk_swap_requested(req: &HttpRequest) -> Option { + 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 { for (name, value) in headers.iter() { if name.to_lowercase().eq(&header_name.to_lowercase()) { diff --git a/packages/ic-http-gateway/src/response/response_handler.rs b/packages/ic-http-gateway/src/response/response_handler.rs index 65d81d3..86ceecf 100644 --- a/packages/ic-http-gateway/src/response/response_handler.rs +++ b/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 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>, + fetched_length: usize, ) -> Result { 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>, skip_verification: bool, ) -> Result, 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,7 +433,16 @@ 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]) @@ -412,10 +450,10 @@ mod tests { 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")); + } } diff --git a/packages/ic-http-gateway/tests/range_request_stream.rs b/packages/ic-http-gateway/tests/range_request_stream.rs index f3dbd6e..d27e3ad 100644 --- a/packages/ic-http-gateway/tests/range_request_stream.rs +++ b/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}"))