use super::{RawReplyError, RawReplyLimits, RawReplyMethod, ReplyMode, charge};
use crate::pic::{PocketIcOperationError, try_pocket_ic_operation};
use ic_host_artifacts::wasm_text::{CompileError, CompileLimits, compile};
use ic_platform::{
candid::contract::encode_one,
management_canister_types::{HttpMethod, HttpRequestArgs},
pocket_ic::{
PocketIc,
common::rest::{CanisterHttpReply, CanisterHttpResponse, MockCanisterHttpResponse},
},
principal::ExternalPrincipal as Principal,
};
use std::{collections::BTreeMap, fmt::Write};
use thiserror::Error;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct HeldRequest {
server: String,
instance: usize,
fixture: String,
slot: usize,
subnet: Principal,
request: u64,
}
impl HeldRequest {
#[must_use]
pub const fn slot(&self) -> usize {
self.slot
}
}
#[derive(Debug, Error)]
pub enum HeldReplyError {
#[error(transparent)]
Catalog(#[from] RawReplyError),
#[error("held reply slot {slot} is not an update")]
Query {
slot: usize,
},
#[error("held replies require manual PocketIC progress")]
AutomaticProgress,
#[error("held reply controller belongs to a different PocketIC instance")]
WrongInstance,
#[error("unknown or already released held request")]
UnknownRequest,
#[error("cannot encode HTTP barrier: {0}")]
Management(String),
#[error(transparent)]
Compile(#[from] CompileError),
#[error(transparent)]
Operation(#[from] PocketIcOperationError),
}
#[derive(Debug)]
pub struct HeldReplyFixture {
wasm: Vec<u8>,
server: String,
instance: usize,
prefix: String,
released: Vec<bool>,
}
impl HeldReplyFixture {
pub fn new(
pic: &PocketIc,
canister: Principal,
methods: &[RawReplyMethod<'_>],
limits: RawReplyLimits,
call_cycles: u128,
module_bytes: usize,
) -> Result<Self, HeldReplyError> {
manual(pic)?;
let prefix = format!("https://ic-testkit.invalid/held/{canister}/");
let wasm = module(methods, limits, &prefix, call_cycles, module_bytes)?;
Ok(Self {
wasm,
server: pic.get_server_url().to_string(),
instance: pic.instance_id(),
prefix,
released: vec![false; methods.len()],
})
}
#[must_use]
pub fn wasm(&self) -> &[u8] {
&self.wasm
}
pub fn pending(&self, pic: &PocketIc) -> Result<Vec<HeldRequest>, HeldReplyError> {
self.admit(pic)?;
Ok(try_pocket_ic_operation(|| pic.get_canister_http())?
.into_iter()
.filter_map(|request| {
let slot = request
.url
.strip_prefix(&self.prefix)?
.parse::<usize>()
.ok()?;
(slot < self.released.len()
&& !self.released[slot]
&& request.body == (slot as u64).to_le_bytes())
.then_some(HeldRequest {
server: self.server.clone(),
instance: self.instance,
fixture: self.prefix.clone(),
slot,
subnet: request.subnet_id,
request: request.request_id,
})
})
.collect())
}
pub fn release(&mut self, pic: &PocketIc, request: &HeldRequest) -> Result<(), HeldReplyError> {
if !self.pending(pic)?.contains(request) {
return Err(HeldReplyError::UnknownRequest);
}
self.released[request.slot] = true;
try_pocket_ic_operation(|| {
pic.mock_canister_http_response(MockCanisterHttpResponse {
subnet_id: request.subnet,
request_id: request.request,
response: CanisterHttpResponse::CanisterHttpReply(CanisterHttpReply {
status: 200,
headers: vec![],
body: vec![],
}),
additional_responses: vec![],
});
})?;
Ok(())
}
fn admit(&self, pic: &PocketIc) -> Result<(), HeldReplyError> {
if pic.instance_id() != self.instance || pic.get_server_url().as_str() != self.server {
return Err(HeldReplyError::WrongInstance);
}
manual(pic)
}
}
fn manual(pic: &PocketIc) -> Result<(), HeldReplyError> {
if try_pocket_ic_operation(|| pic.auto_progress_enabled())? {
return Err(HeldReplyError::AutomaticProgress);
}
Ok(())
}
fn module(
methods: &[RawReplyMethod<'_>],
limits: RawReplyLimits,
prefix: &str,
cycles: u128,
module_bytes: usize,
) -> Result<Vec<u8>, HeldReplyError> {
let exports = admit_catalog(methods, limits)?;
let mut source = String::from(
r#"(module
(import "ic0" "msg_reply_data_append" (func $append (param i32 i32)))
(import "ic0" "msg_reply" (func $reply))
(import "ic0" "msg_reject" (func $reject (param i32 i32)))
(import "ic0" "call_new" (func $new (param i32 i32 i32 i32 i32 i32 i32 i32)))
(import "ic0" "call_data_append" (func $data (param i32 i32)))
(import "ic0" "call_cycles_add128" (func $cycles (param i64 i64)))
(import "ic0" "call_perform" (func $perform (result i32)))
"#,
);
let mut offset = 0;
let failure = segment(&mut source, &mut offset, b"held reply barrier failed")?;
let exhausted = segment(&mut source, &mut offset, b"held reply slots exhausted")?;
let http = segment(&mut source, &mut offset, b"http_request")?;
writeln!(
source,
"(func $failed (param i32) (call $reject (i32.const {failure}) (i32.const 25)))"
)
.unwrap();
for (slot, method) in methods.iter().enumerate() {
let request = encode_one(HttpRequestArgs {
url: format!("{prefix}{slot}"),
max_response_bytes: Some(256),
method: HttpMethod::GET,
headers: vec![],
body: Some((slot as u64).to_le_bytes().to_vec()),
transform: None,
is_replicated: Some(true),
pricing_version: None,
})
.map_err(|error| HeldReplyError::Management(error.to_string()))?;
let payload = segment(&mut source, &mut offset, &request)?;
let reply = segment(&mut source, &mut offset, method.reply)?;
writeln!(
source,
r"
(global $used{slot} (mut i32) (i32.const 0))
(func $reply{slot} (param i32)
(call $append (i32.const {reply}) (i32.const {reply_len})) (call $reply))
(func $send{slot}
(call $new (i32.const 0) (i32.const 0) (i32.const {http}) (i32.const 12)
(i32.const {callback}) (i32.const 0) (i32.const 0) (i32.const 0))
(call $cycles (i64.const {high}) (i64.const {low}))
(call $data (i32.const {payload}) (i32.const {payload_len}))
(if (call $perform) (then unreachable)))
",
reply_len = method.reply.len(),
callback = slot + 1,
high = u64::try_from(cycles >> 64).unwrap().cast_signed(),
low = u64::try_from(cycles & u128::from(u64::MAX))
.unwrap()
.cast_signed(),
payload_len = request.len()
)
.unwrap();
}
for (name, slots) in exports {
let export = escaped(format!("canister_update {name}").as_bytes());
writeln!(source, "(func (export \"{export}\")").unwrap();
for slot in slots {
writeln!(source, "(if (i32.eqz (global.get $used{slot})) (then (global.set $used{slot} (i32.const 1)) (call $send{slot}) (return)))").unwrap();
}
writeln!(
source,
"(call $reject (i32.const {exhausted}) (i32.const 26)))"
)
.unwrap();
}
let pages = offset.div_ceil(65_536).max(1);
writeln!(source, "(memory (export \"memory\") {pages} {pages})\n(table {} funcref)\n(elem (i32.const 0) $failed", methods.len() + 1).unwrap();
for slot in 0..methods.len() {
write!(source, " $reply{slot}").unwrap();
}
source.push_str("))");
Ok(compile(
&source,
CompileLimits {
source_bytes: source.len(),
module_bytes,
},
)?)
}
fn admit_catalog<'a>(
methods: &[RawReplyMethod<'a>],
limits: RawReplyLimits,
) -> Result<BTreeMap<&'a str, Vec<usize>>, HeldReplyError> {
charge(0, methods.len(), limits.methods, "requests")?;
u32::try_from(methods.len()).map_err(|_| RawReplyError::Wasm32Limits)?;
let mut names = 0;
let mut replies = 0;
let mut exports = BTreeMap::<&str, Vec<usize>>::new();
for (slot, method) in methods.iter().enumerate() {
if method.mode != ReplyMode::Update {
return Err(HeldReplyError::Query { slot });
}
if method.name.is_empty() || method.name.chars().any(char::is_control) {
return Err(RawReplyError::InvalidName { index: slot }.into());
}
names = charge(
names,
method.name.len(),
limits.method_name_bytes,
"method names",
)?;
replies = charge(
replies,
method.reply.len(),
limits.reply_bytes,
"reply bytes",
)?;
exports.entry(method.name).or_default().push(slot);
}
Ok(exports)
}
fn segment(source: &mut String, offset: &mut usize, bytes: &[u8]) -> Result<usize, RawReplyError> {
let start = *offset;
*offset = offset
.checked_add(bytes.len())
.filter(|end| i32::try_from(*end).is_ok())
.ok_or(RawReplyError::Wasm32Limits)?;
writeln!(source, "(data (i32.const {start}) \"{}\")", escaped(bytes)).unwrap();
Ok(start)
}
fn escaped(bytes: &[u8]) -> String {
let mut result = String::new();
for byte in bytes {
write!(result, "\\{byte:02x}").unwrap();
}
result
}
#[cfg(test)]
mod tests;