Skip to content

Commit

Permalink
feat(ic-http-gateway): add more tests and checks for long asset handling
Browse files Browse the repository at this point in the history
  • Loading branch information
przydatek committed Oct 30, 2024
1 parent 432c56b commit 4af91b5
Show file tree
Hide file tree
Showing 3 changed files with 203 additions and 26 deletions.
60 changes: 57 additions & 3 deletions examples/http-gateway/canister/src/custom_assets/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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
}

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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()) {
Expand Down
87 changes: 70 additions & 17 deletions packages/ic-http-gateway/src/response/response_handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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),
)))
},
Expand Down Expand Up @@ -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,
Expand All @@ -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])
Expand All @@ -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])
Expand All @@ -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"));
}
}
82 changes: 76 additions & 6 deletions packages/ic-http-gateway/tests/range_request_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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}"))
Expand All @@ -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,
Expand All @@ -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)]
Expand Down Expand Up @@ -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}"))
Expand Down

0 comments on commit 4af91b5

Please sign in to comment.