use crate::pocket_ic::common::rest::{
AutoProgressConfig, BlobCompression, BlobId, CanisterHttpRequest, CreateHttpGatewayResponse,
CreateInstanceResponse, ExtendedSubnetConfigSet, HttpGatewayBackend, HttpGatewayConfig,
HttpGatewayInfo, HttpsConfig, IcpConfig, IcpFeatures, InitialTime, InstanceConfig,
InstanceHttpGatewayConfig, InstanceId, MockCanisterHttpResponse,
MockFlexibleCanisterHttpResponse, RawAddCycles, RawCanisterCall, RawCanisterHttpRequest,
RawCanisterId, RawCanisterResult, RawCanisterSnapshotDownload, RawCanisterSnapshotId,
RawCanisterSnapshotUpload, RawCycles, RawEffectivePrincipal, RawIngressStatusArgs,
RawMessageId, RawMockCanisterHttpResponse, RawMockFlexibleCanisterHttpResponse, RawPrincipalId,
RawSenderInfo, RawSetStableMemory, RawStableMemory, RawSubnetId, RawTickConfigs, RawTime,
RawVerifyCanisterSigArg, SubnetId, Topology,
};
#[cfg(windows)]
use crate::pocket_ic::wsl_path;
use crate::pocket_ic::{ClientError, Operation, RequestContext, http::Exchange};
use crate::pocket_ic::{
CreateCanisterParams, CreateCanisterPlacement, IngressStatusResult, PocketIcBuilder,
PocketIcState, RejectResponse, StartServerParams, TickConfigs, Time, copy_dir, start_server,
};
use backon::{BackoffBuilder, ExponentialBuilder};
use candid::{
Principal, decode_args, encode_args,
utils::{ArgumentDecoder, ArgumentEncoder},
};
use ic_certification::{Certificate, Label, LookupResult};
use ic_management_canister_types::{
CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterLogRecord, CanisterSettings,
CanisterStatusResult, ChunkHash, DeleteCanisterSnapshotArgs, FetchCanisterLogsResult,
InstallChunkedCodeArgs, InstallCodeArgs, LoadCanisterSnapshotArgs,
ProvisionalCreateCanisterWithCyclesArgs, Snapshot, StoredChunksResult,
TakeCanisterSnapshotArgs, UpdateSettingsArgs, UpgradeFlags, UploadChunkArgs, UploadChunkResult,
WasmMemoryPersistence,
};
use ic_transport_types::Envelope;
use ic_transport_types::EnvelopeContent::ReadState;
use ic_transport_types::{ReadStateResponse, SubnetMetrics};
use reqwest::{StatusCode, Url};
use serde::{Serialize, de::DeserializeOwned};
use sha2::{Digest, Sha256};
use slog::Level;
use std::fs::read_dir;
use std::future::Future;
use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::time::Duration;
use tracing::{debug, instrument};
use tracing_appender::non_blocking::WorkerGuard;
use tracing_subscriber::EnvFilter;
mod snapshot;
const LOG_DIR_PATH_ENV_NAME: &str = "POCKET_IC_LOG_DIR";
const LOG_DIR_LEVELS_ENV_NAME: &str = "POCKET_IC_LOG_DIR_LEVELS";
const LOCALHOST: &str = "localhost";
const INSTALL_CHUNKED_CODE_THRESHOLD: usize = 2_000_000;
const INSTALL_CODE_CHUNK_SIZE: usize = 1 << 20;
enum HttpMethod {
Get,
Post,
}
pub struct PocketIc {
pub instance_id: InstanceId,
max_request_time_ms: Option<u64>,
http_gateway: Option<HttpGatewayInfo>,
server_url: Url,
reqwest_client: reqwest::Client,
owns_instance: bool,
pub(crate) cleanup_attempted: bool,
state_dir: Option<PocketIcState>,
_log_guard: Option<WorkerGuard>,
}
impl PocketIc {
pub async fn new() -> Result<Self, ClientError> {
PocketIcBuilder::new()
.with_application_subnet()
.build_async()
.await
}
pub fn new_from_existing_instance(
server_url: Url,
instance_id: InstanceId,
max_request_time_ms: Option<u64>,
) -> Self {
let test_driver_pid = std::process::id();
let log_guard = setup_tracing(test_driver_pid);
let reqwest_client = reqwest::Client::new();
debug!("instance_id={} Reusing existing instance.", instance_id);
Self {
instance_id,
max_request_time_ms,
http_gateway: None,
server_url,
reqwest_client,
owns_instance: false,
cleanup_attempted: false,
state_dir: None,
_log_guard: log_guard,
}
}
#[expect(
clippy::too_many_arguments,
reason = "retained upstream client method signature"
)]
pub(crate) async fn from_components(
subnet_config_set: impl Into<ExtendedSubnetConfigSet>,
server_url: Option<Url>,
server_binary: Option<PathBuf>,
max_request_time_ms: Option<u64>,
read_only_state_dir: Option<PathBuf>,
mut state_dir: Option<PocketIcState>,
icp_config: IcpConfig,
log_level: Option<Level>,
bitcoind_addr: Option<Vec<SocketAddr>>,
dogecoind_addr: Option<Vec<SocketAddr>>,
icp_features: IcpFeatures,
initial_time: Option<InitialTime>,
http_gateway_config: Option<InstanceHttpGatewayConfig>,
mainnet_nns_subnet_id: Option<bool>,
disable_ingress_validation: Option<bool>,
) -> Result<Self, ClientError> {
let server_url = if let Some(server_url) = server_url {
server_url
} else {
let (_, server_url) = start_server(StartServerParams {
server_binary,
reuse: true,
ttl: None,
hard_ttl: None,
})
.await?;
server_url
};
let subnet_config_set: ExtendedSubnetConfigSet = subnet_config_set.into();
if let Some(read_only_state_dir) = read_only_state_dir {
if let Some(ref state_dir) = state_dir {
let mut state_dir_contents = read_dir(state_dir.state_dir())?;
if state_dir_contents.next().is_some() {
return Err(ClientError::Protocol(
"PocketIC instance state must be empty if a read-only state is mounted."
.into(),
));
}
} else {
state_dir = Some(PocketIcState::new());
};
copy_dir(
read_only_state_dir,
state_dir
.as_ref()
.map(|state_dir| state_dir.state_dir())
.unwrap(),
)?;
};
let instance_config = InstanceConfig {
subnet_config_set,
http_gateway_config,
#[cfg(not(windows))]
state_dir: state_dir.as_ref().map(|state_dir| state_dir.state_dir()),
#[cfg(windows)]
state_dir: state_dir
.as_ref()
.map(|state_dir| wsl_path(&state_dir.state_dir(), "state directory").into()),
icp_config: Some(icp_config),
log_level: log_level.map(|l| l.to_string()),
bitcoind_addr,
dogecoind_addr,
icp_features: Some(icp_features),
incomplete_state: None,
initial_time,
mainnet_nns_subnet_id,
disable_ingress_validation,
};
let test_driver_pid = std::process::id();
let log_guard = setup_tracing(test_driver_pid);
let reqwest_client = reqwest::Client::new();
let url = server_url.join("instances").unwrap();
let mut exchange = Exchange::new(&url, max_request_time_ms);
let response: CreateInstanceResponse = exchange
.json(
&reqwest_client,
reqwest_client.post(url.clone()).json(&instance_config),
)
.await?;
let (instance_id, http_gateway_info) = match response {
CreateInstanceResponse::Created {
instance_id,
http_gateway_info,
..
} => (instance_id, http_gateway_info),
CreateInstanceResponse::Error { message } => {
return Err(ClientError::Server {
context: Box::new(exchange.context.clone()),
status: exchange.context.status.unwrap_or(StatusCode::OK),
message,
});
}
};
debug!("instance_id={} New instance created.", instance_id);
Ok(Self {
instance_id,
max_request_time_ms,
http_gateway: http_gateway_info,
server_url,
reqwest_client,
owns_instance: true,
cleanup_attempted: false,
state_dir,
_log_guard: log_guard,
})
}
pub async fn drop_and_take_state(&mut self) -> Result<Option<PocketIcState>, ClientError> {
self.do_drop().await?;
Ok(self.state_dir.take())
}
pub(crate) fn take_state_internal(&mut self) -> Option<PocketIcState> {
self.state_dir.take()
}
pub fn get_server_url(&self) -> Url {
self.server_url.clone()
}
pub async fn topology(&self) -> Result<Topology, ClientError> {
let endpoint = "read/topology";
self.get(endpoint).await
}
#[instrument(ret, skip(self, blob), fields(instance_id=self.instance_id, blob_len = %blob.len(), compression = ?compression))]
pub async fn upload_blob(
&self,
blob: Vec<u8>,
compression: BlobCompression,
) -> Result<BlobId, ClientError> {
let url = self.server_url.join("blobstore/").unwrap();
let mut request = self.reqwest_client.post(url.clone()).body(blob);
if let BlobCompression::Gzip = compression {
request = request.header(reqwest::header::CONTENT_ENCODING, "gzip");
}
let mut exchange = Exchange::new(&url, self.max_request_time_ms);
let (status, bytes) = exchange.bytes(&self.reqwest_client, request).await?;
if !status.is_success() {
return Err(exchange.server_error(status, &bytes));
}
let hash = hex::decode(bytes).map_err(|error| ClientError::Protocol(error.to_string()))?;
Ok(BlobId(hash))
}
#[instrument(skip(self, data), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), data_len = %data.len(), compression = ?compression))]
pub async fn set_stable_memory(
&self,
canister_id: CanisterId,
data: Vec<u8>,
compression: BlobCompression,
) -> Result<(), ClientError> {
let blob_id = self.upload_blob(data, compression).await?;
let endpoint = "update/set_stable_memory";
self.post::<(), _>(
endpoint,
RawSetStableMemory {
canister_id: canister_id.as_slice().to_vec(),
blob_id,
},
)
.await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
pub async fn get_stable_memory(&self, canister_id: CanisterId) -> Result<Vec<u8>, ClientError> {
let endpoint = "read/get_stable_memory";
let RawStableMemory { blob } = self
.post(
endpoint,
RawCanisterId {
canister_id: canister_id.as_slice().to_vec(),
},
)
.await?;
Ok(blob)
}
#[instrument(ret)]
pub async fn list_instances() -> Result<Vec<String>, ClientError> {
let (_, server_url) = start_server(StartServerParams {
reuse: true,
..Default::default()
})
.await?;
let url = server_url.join("instances").unwrap();
let client = reqwest::Client::new();
Exchange::new(&url, None)
.json(&client, client.get(url.clone()))
.await
}
#[instrument(skip_all, fields(instance_id=self.instance_id))]
pub async fn verify_canister_signature(
&self,
msg: Vec<u8>,
sig: Vec<u8>,
pubkey: Vec<u8>,
root_pubkey: Vec<u8>,
) -> Result<(), ClientError> {
let url = self.server_url.join("verify_signature").unwrap();
let result: Result<(), String> = Exchange::new(&url, self.max_request_time_ms)
.json(
&self.reqwest_client,
self.reqwest_client
.post(url.clone())
.json(&RawVerifyCanisterSigArg {
msg,
sig,
pubkey,
root_pubkey,
}),
)
.await?;
result.map_err(ClientError::Protocol)
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn tick(&self) -> Result<(), ClientError> {
self.tick_with_configs(TickConfigs::default()).await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn tick_with_configs(&self, configs: TickConfigs) -> Result<(), ClientError> {
let endpoint = "update/tick";
self.post::<(), RawTickConfigs>(endpoint, configs.into())
.await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn auto_progress(&self) -> Result<Url, ClientError> {
let endpoint = "auto_progress";
let auto_progress_config = AutoProgressConfig {
artificial_delay_ms: None,
};
self.post::<(), _>(endpoint, auto_progress_config).await?;
Ok(self.instance_url())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn auto_progress_enabled(&self) -> Result<bool, ClientError> {
self.get("auto_progress").await
}
pub(crate) fn instance_url(&self) -> Url {
self.server_url
.join("/instances/")
.unwrap()
.join(&format!("{}/", self.instance_id))
.unwrap()
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn stop_progress(&self) -> Result<(), ClientError> {
let endpoint = "stop_progress";
self.post::<(), _>(endpoint, "").await?;
Ok(())
}
pub fn url(&self) -> Option<Url> {
self.http_gateway
.as_ref()
.map(|res| Url::parse(&format!("http://{}:{}/", LOCALHOST, res.port)).unwrap())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn make_live(&mut self, listen_at: Option<u16>) -> Result<Url, ClientError> {
self.make_live_with_params(None, listen_at, None, None)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn make_live_with_params(
&mut self,
ip_addr: Option<IpAddr>,
listen_at: Option<u16>,
domains: Option<Vec<String>>,
https_config: Option<HttpsConfig>,
) -> Result<Url, ClientError> {
if let Some(url) = self.url() {
return Ok(url);
}
if !self.auto_progress_enabled().await? {
self.auto_progress().await?;
}
self.start_http_gateway(
ip_addr.map(|ip_addr| ip_addr.to_string()),
listen_at,
domains,
https_config,
)
.await
}
async fn start_http_gateway(
&mut self,
ip_addr: Option<String>,
port: Option<u16>,
domains: Option<Vec<String>>,
https_config: Option<HttpsConfig>,
) -> Result<Url, ClientError> {
let endpoint = self.server_url.join("http_gateway").unwrap();
let http_gateway_config = HttpGatewayConfig {
ip_addr,
port,
forward_to: HttpGatewayBackend::PocketIcInstance(self.instance_id),
domains: domains.clone(),
https_config: https_config.clone(),
domain_custom_provider_local_file: None,
};
let mut exchange = Exchange::new(&endpoint, self.max_request_time_ms);
let res: CreateHttpGatewayResponse = exchange
.json(
&self.reqwest_client,
self.reqwest_client
.post(endpoint.clone())
.json(&http_gateway_config),
)
.await?;
Ok(match res {
CreateHttpGatewayResponse::Created(info) => {
let port = info.port;
self.http_gateway = Some(info);
let proto = if https_config.is_some() {
"https"
} else {
"http"
};
Url::parse(&format!(
"{}://{}:{}/",
proto,
domains
.unwrap_or_default()
.into_iter()
.next()
.unwrap_or(LOCALHOST.to_string()),
port
))
.unwrap()
}
CreateHttpGatewayResponse::Error { message } => {
return Err(ClientError::Server {
context: Box::new(exchange.context.clone()),
status: exchange.context.status.unwrap_or(StatusCode::OK),
message,
});
}
})
}
async fn stop_http_gateway(&mut self) -> Result<(), ClientError> {
if let Some(info) = &self.http_gateway {
let url = self
.server_url
.join(&format!("http_gateway/{}/stop", info.instance_id))
.unwrap();
Exchange::new(&url, self.max_request_time_ms)
.json::<()>(&self.reqwest_client, self.reqwest_client.post(url.clone()))
.await?;
self.http_gateway = None;
}
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn stop_live(&mut self) -> Result<(), ClientError> {
self.stop_http_gateway().await?;
self.stop_progress().await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id))]
pub async fn root_key(&self) -> Result<Option<Vec<u8>>, ClientError> {
let Some(subnet_id) = self.topology().await?.get_nns() else {
return Ok(None);
};
let subnet_id: RawSubnetId = subnet_id.into();
let endpoint = "read/pub_key";
let res = self.post::<Vec<u8>, _>(endpoint, subnet_id).await?;
Ok(Some(res))
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn get_time(&self) -> Result<Time, ClientError> {
let endpoint = "read/get_time";
let result: RawTime = self.get(endpoint).await?;
Ok(Time::from_nanos_since_unix_epoch(result.nanos_since_epoch))
}
#[instrument(skip(self), fields(instance_id=self.instance_id, time = ?time))]
pub async fn set_time(&self, time: Time) -> Result<(), ClientError> {
let endpoint = "update/set_time";
self.post::<(), _>(
endpoint,
RawTime {
nanos_since_epoch: time.as_nanos_since_unix_epoch(),
},
)
.await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id, time = ?time))]
pub async fn set_certified_time(&self, time: Time) -> Result<(), ClientError> {
let endpoint = "update/set_certified_time";
self.post::<(), _>(
endpoint,
RawTime {
nanos_since_epoch: time.as_nanos_since_unix_epoch(),
},
)
.await?;
Ok(())
}
#[instrument(skip(self), fields(instance_id=self.instance_id, duration = ?duration))]
pub async fn advance_time(&self, duration: Duration) -> Result<(), ClientError> {
let now = self.get_time().await?;
self.set_time(now + duration).await?;
Ok(())
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
pub async fn get_controllers(
&self,
canister_id: CanisterId,
) -> Result<Vec<Principal>, ClientError> {
let endpoint = "read/get_controllers";
let result: Vec<RawPrincipalId> = self
.post(
endpoint,
RawCanisterId {
canister_id: canister_id.as_slice().to_vec(),
},
)
.await?;
Ok(result.into_iter().map(|p| p.into()).collect())
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
pub async fn cycle_balance(&self, canister_id: CanisterId) -> Result<u128, ClientError> {
let endpoint = "read/get_cycles";
let result: RawCycles = self
.post(
endpoint,
RawCanisterId {
canister_id: canister_id.as_slice().to_vec(),
},
)
.await?;
Ok(result.cycles)
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), amount = %amount))]
pub async fn add_cycles(
&self,
canister_id: CanisterId,
amount: u128,
) -> Result<u128, ClientError> {
let endpoint = "update/add_cycles";
let result: RawCycles = self
.post(
endpoint,
RawAddCycles {
canister_id: canister_id.as_slice().to_vec(),
amount,
},
)
.await?;
Ok(result.cycles)
}
pub async fn submit_call(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<RawMessageId, ClientError> {
self.submit_call_with_effective_principal(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
)
.await
}
async fn submit_call_with_effective_principal_and_sender_info_inner(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: Option<RawSenderInfo>,
) -> Result<RawMessageId, ClientError> {
let endpoint = "update/submit_ingress_message";
let raw_canister_call = RawCanisterCall {
sender: sender.as_slice().to_vec(),
canister_id: canister_id.as_slice().to_vec(),
method: method.to_string(),
payload,
effective_principal,
sender_info,
};
self.post::<Result<RawMessageId, RejectResponse>, _>(endpoint, raw_canister_call)
.await?
.map_err(ClientError::from)
}
pub async fn submit_call_with_effective_principal(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<RawMessageId, ClientError> {
self.submit_call_with_effective_principal_and_sender_info_inner(
canister_id,
effective_principal,
sender,
method,
payload,
None,
)
.await
}
pub async fn submit_call_with_effective_principal_and_sender_info(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<RawMessageId, ClientError> {
self.submit_call_with_effective_principal_and_sender_info_inner(
canister_id,
effective_principal,
sender,
method,
payload,
Some(sender_info),
)
.await
}
pub async fn submit_call_with_sender_info(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<RawMessageId, ClientError> {
self.submit_call_with_effective_principal_and_sender_info(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
sender_info,
)
.await
}
pub async fn await_call(&self, message_id: RawMessageId) -> Result<Vec<u8>, ClientError> {
let endpoint = "update/await_ingress_message";
let result: RawCanisterResult = self
.post(endpoint, &message_id)
.await
.map_err(|error| error.with_ingress(message_id))?;
Result::<Vec<u8>, RejectResponse>::from(result).map_err(ClientError::from)
}
pub async fn ingress_status(
&self,
raw_message_id: RawMessageId,
) -> Result<Option<Result<Vec<u8>, RejectResponse>>, ClientError> {
let status = self.ingress_status_as_caller(raw_message_id, None).await?;
Ok(match status {
IngressStatusResult::NotAvailable => None,
IngressStatusResult::Success(status) => Some(status),
IngressStatusResult::Forbidden(err) => return Err(ClientError::Protocol(err)),
})
}
pub async fn ingress_status_as(
&self,
raw_message_id: RawMessageId,
caller: Principal,
) -> Result<IngressStatusResult, ClientError> {
self.ingress_status_as_caller(raw_message_id, Some(caller))
.await
}
async fn ingress_status_as_caller(
&self,
raw_message_id: RawMessageId,
caller: Option<Principal>,
) -> Result<IngressStatusResult, ClientError> {
let args = RawIngressStatusArgs {
raw_message_id: raw_message_id.clone(),
raw_caller: caller.map(Into::into),
};
match self
.post::<Option<RawCanisterResult>, _>("read/ingress_status", args)
.await
{
Ok(None) => Ok(IngressStatusResult::NotAvailable),
Ok(Some(result)) => Ok(IngressStatusResult::Success(result.into())),
Err(ClientError::Server {
status: StatusCode::FORBIDDEN,
message,
..
}) => Ok(IngressStatusResult::Forbidden(message)),
Err(error) => Err(error.with_ingress(raw_message_id)),
}
}
pub async fn await_call_no_ticks(
&self,
message_id: RawMessageId,
) -> Result<Vec<u8>, ClientError> {
let mut retry_policy = ExponentialBuilder::new()
.with_min_delay(Duration::from_millis(10))
.with_max_delay(Duration::from_secs(1))
.with_factor(2.0)
.with_jitter()
.with_total_delay(Some(Duration::from_secs(15 * 60)))
.without_max_times()
.build();
Ok(loop {
if let Some(ingress_status) = self.ingress_status(message_id.clone()).await? {
break ingress_status.map_err(ClientError::from)?;
}
let delay = retry_policy.next().ok_or_else(|| ClientError::Timeout {
context: Box::new(RequestContext {
url: self.instance_url().to_string(),
operation: None,
ingress: Some(message_id.clone()),
may_have_executed: true,
status: None,
}),
})?;
tokio::time::sleep(delay).await;
})
}
#[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
pub async fn update_call(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<Vec<u8>, ClientError> {
self.update_call_with_effective_principal(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
)
.await
}
async fn update_call_with_effective_principal_and_sender_info_inner(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: Option<RawSenderInfo>,
) -> Result<Vec<u8>, ClientError> {
let message_id = self
.submit_call_with_effective_principal_and_sender_info_inner(
canister_id,
effective_principal,
sender,
method,
payload,
sender_info,
)
.await?;
self.await_call(message_id).await
}
pub async fn update_call_with_effective_principal_and_sender_info(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<Vec<u8>, ClientError> {
self.update_call_with_effective_principal_and_sender_info_inner(
canister_id,
effective_principal,
sender,
method,
payload,
Some(sender_info),
)
.await
}
pub async fn update_call_with_sender_info(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<Vec<u8>, ClientError> {
self.update_call_with_effective_principal_and_sender_info(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
sender_info,
)
.await
}
#[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
pub async fn query_call(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<Vec<u8>, ClientError> {
self.query_call_with_effective_principal(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
)
.await
}
#[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), effective_principal = %effective_principal.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
pub async fn query_call_with_effective_principal(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<Vec<u8>, ClientError> {
let endpoint = "read/query";
self.canister_call(
endpoint,
effective_principal,
canister_id,
sender,
method,
payload,
None,
)
.await
}
pub async fn query_call_with_effective_principal_and_sender_info(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<Vec<u8>, ClientError> {
let endpoint = "read/query";
self.canister_call(
endpoint,
effective_principal,
canister_id,
sender,
method,
payload,
Some(sender_info),
)
.await
}
pub async fn query_call_with_sender_info(
&self,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: RawSenderInfo,
) -> Result<Vec<u8>, ClientError> {
self.query_call_with_effective_principal_and_sender_info(
canister_id,
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
method,
payload,
sender_info,
)
.await
}
pub async fn fetch_canister_logs(
&self,
canister_id: CanisterId,
sender: Principal,
) -> Result<Vec<CanisterLogRecord>, ClientError> {
Ok(with_candid::<_, (FetchCanisterLogsResult,), _>(
(CanisterIdRecord { canister_id },),
|payload| async {
self.query_call_with_effective_principal(
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender,
"fetch_canister_logs",
payload,
)
.await
},
)
.await?
.0
.canister_log_records)
}
#[instrument(skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn canister_status(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<CanisterStatusResult, ClientError> {
Ok(
call_candid_as::<(CanisterIdRecord,), (CanisterStatusResult,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"canister_status",
(CanisterIdRecord { canister_id },),
)
.await?
.0,
)
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn create_canister(&self) -> Result<CanisterId, ClientError> {
self.create_canister_with_params(None, CreateCanisterParams::default())
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, settings = ?settings, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn create_canister_with_settings(
&self,
sender: Option<Principal>,
settings: Option<CanisterSettings>,
) -> Result<CanisterId, ClientError> {
self.create_canister_with_params(
sender,
CreateCanisterParams {
settings,
..Default::default()
},
)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string(), settings = ?settings, canister_id = %canister_id.to_string()))]
pub async fn create_canister_with_id(
&self,
sender: Option<Principal>,
settings: Option<CanisterSettings>,
canister_id: CanisterId,
) -> Result<CanisterId, ClientError> {
self.create_canister_with_params(
sender,
CreateCanisterParams {
settings,
placement: Some(CreateCanisterPlacement::CanisterId(canister_id)),
..Default::default()
},
)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string(), settings = ?settings, subnet_id = %subnet_id.to_string()))]
pub async fn create_canister_on_subnet(
&self,
sender: Option<Principal>,
settings: Option<CanisterSettings>,
subnet_id: SubnetId,
) -> Result<CanisterId, ClientError> {
self.create_canister_with_params(
sender,
CreateCanisterParams {
settings,
placement: Some(CreateCanisterPlacement::SubnetId(subnet_id)),
..Default::default()
},
)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn create_canister_with_params(
&self,
sender: Option<Principal>,
params: CreateCanisterParams,
) -> Result<CanisterId, ClientError> {
let cycles = params.cycles.unwrap_or(100_000_000_000_000);
let (effective_principal, specified_id) = match params.placement {
None => (RawEffectivePrincipal::None, None),
Some(CreateCanisterPlacement::SubnetId(subnet_id)) => (
RawEffectivePrincipal::SubnetId(subnet_id.as_slice().to_vec()),
None,
),
Some(CreateCanisterPlacement::CanisterId(canister_id)) => (
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
Some(canister_id),
),
};
Ok(call_candid_as::<_, (CanisterIdRecord,)>(
self,
Principal::management_canister(),
effective_principal,
sender.unwrap_or(Principal::anonymous()),
"provisional_create_canister_with_cycles",
(ProvisionalCreateCanisterWithCyclesArgs {
settings: params.settings,
amount: Some(cycles.into()),
specified_id,
sender_canister_version: None,
},),
)
.await?
.0
.canister_id)
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn upload_chunk(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
chunk: Vec<u8>,
) -> Result<Vec<u8>, ClientError> {
Ok(call_candid_as::<_, (UploadChunkResult,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"upload_chunk",
(UploadChunkArgs { canister_id, chunk },),
)
.await?
.0
.hash)
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn stored_chunks(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<Vec<Vec<u8>>, ClientError> {
Ok(call_candid_as::<_, (StoredChunksResult,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"stored_chunks",
(CanisterIdRecord { canister_id },),
)
.await?
.0
.into_iter()
.map(|chunk| chunk.hash)
.collect())
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn clear_chunk_store(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), ClientError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"clear_chunk_store",
(CanisterIdRecord { canister_id },),
)
.await
}
#[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))]
pub async fn install_chunked_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
mode: CanisterInstallMode,
store_canister_id: CanisterId,
chunk_hashes_list: Vec<Vec<u8>>,
wasm_module_hash: Vec<u8>,
arg: Vec<u8>,
) -> Result<(), ClientError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"install_chunked_code",
(InstallChunkedCodeArgs {
mode,
target_canister: canister_id,
store_canister: Some(store_canister_id),
chunk_hashes_list: chunk_hashes_list
.into_iter()
.map(|hash| ChunkHash { hash })
.collect(),
wasm_module_hash,
arg,
sender_canister_version: None,
},),
)
.await
}
async fn install_canister_helper(
&self,
mode: CanisterInstallMode,
canister_id: CanisterId,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), ClientError> {
if wasm_module.len() + arg.len() < INSTALL_CHUNKED_CODE_THRESHOLD {
call_candid_as::<(InstallCodeArgs,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"install_code",
(InstallCodeArgs {
mode,
canister_id,
wasm_module,
arg,
sender_canister_version: None,
},),
)
.await?
} else {
self.clear_chunk_store(canister_id, sender).await?;
let chunks: Vec<_> = wasm_module.chunks(INSTALL_CODE_CHUNK_SIZE).collect();
let mut hashes = vec![];
for chunk in chunks {
let hash = self
.upload_chunk(canister_id, sender, chunk.to_vec())
.await?;
hashes.push(hash);
}
let mut hasher = Sha256::new();
hasher.update(wasm_module);
let wasm_module_hash = hasher.finalize().to_vec();
self.install_chunked_canister(
canister_id,
sender,
mode,
canister_id,
hashes,
wasm_module_hash,
arg,
)
.await?
};
Ok(())
}
#[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn install_canister(
&self,
canister_id: CanisterId,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), ClientError> {
self.install_canister_helper(
CanisterInstallMode::Install,
canister_id,
wasm_module,
arg,
sender,
)
.await
}
#[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn upgrade_canister(
&self,
canister_id: CanisterId,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), ClientError> {
self.install_canister_helper(
CanisterInstallMode::Upgrade(None),
canister_id,
wasm_module,
arg,
sender,
)
.await
}
#[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn upgrade_eop_canister(
&self,
canister_id: CanisterId,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), ClientError> {
self.install_canister_helper(
CanisterInstallMode::Upgrade(Some(UpgradeFlags {
wasm_memory_persistence: Some(WasmMemoryPersistence::Keep),
skip_pre_upgrade: None,
})),
canister_id,
wasm_module,
arg,
sender,
)
.await
}
#[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn reinstall_canister(
&self,
canister_id: CanisterId,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), ClientError> {
self.install_canister_helper(
CanisterInstallMode::Reinstall,
canister_id,
wasm_module,
arg,
sender,
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn uninstall_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), ClientError> {
call_candid_as::<(CanisterIdRecord,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"uninstall_code",
(CanisterIdRecord { canister_id },),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn take_canister_snapshot(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
replace_snapshot: Option<Vec<u8>>,
) -> Result<Snapshot, ClientError> {
Ok(call_candid_as::<_, (Snapshot,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"take_canister_snapshot",
(TakeCanisterSnapshotArgs {
canister_id,
replace_snapshot,
uninstall_code: None,
sender_canister_version: None,
},),
)
.await?
.0)
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn load_canister_snapshot(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
snapshot_id: Vec<u8>,
) -> Result<(), ClientError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"load_canister_snapshot",
(LoadCanisterSnapshotArgs {
canister_id,
snapshot_id,
sender_canister_version: None,
},),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn list_canister_snapshots(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<Vec<Snapshot>, ClientError> {
Ok(call_candid_as::<_, (Vec<Snapshot>,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"list_canister_snapshots",
(CanisterIdRecord { canister_id },),
)
.await?
.0)
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn delete_canister_snapshot(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
snapshot_id: Vec<u8>,
) -> Result<(), ClientError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"delete_canister_snapshot",
(DeleteCanisterSnapshotArgs {
canister_id,
snapshot_id,
},),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn update_canister_settings(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
settings: CanisterSettings,
) -> Result<(), ClientError> {
call_candid_as::<_, ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"update_settings",
(UpdateSettingsArgs {
canister_id,
settings,
sender_canister_version: None,
},),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn set_controllers(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
new_controllers: Vec<Principal>,
) -> Result<(), ClientError> {
let settings = CanisterSettings {
controllers: Some(new_controllers),
..CanisterSettings::default()
};
call_candid_as::<(UpdateSettingsArgs,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"update_settings",
(UpdateSettingsArgs {
canister_id,
settings,
sender_canister_version: None,
},),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn start_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), ClientError> {
call_candid_as::<(CanisterIdRecord,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"start_canister",
(CanisterIdRecord { canister_id },),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn stop_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), ClientError> {
call_candid_as::<(CanisterIdRecord,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"stop_canister",
(CanisterIdRecord { canister_id },),
)
.await
}
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn delete_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), ClientError> {
call_candid_as::<(CanisterIdRecord,), ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"delete_canister",
(CanisterIdRecord { canister_id },),
)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
pub async fn canister_exists(&self, canister_id: CanisterId) -> Result<bool, ClientError> {
Ok(self.get_subnet(canister_id).await?.is_some())
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, subnet_id = %subnet_id.to_string()))]
pub async fn delete_subnet(&self, subnet_id: SubnetId) -> Result<(), ClientError> {
let endpoint = "update/delete_subnet";
self.post::<(), RawSubnetId>(
endpoint,
RawSubnetId {
subnet_id: subnet_id.as_slice().to_vec(),
},
)
.await?;
Ok(())
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
pub async fn get_subnet(
&self,
canister_id: CanisterId,
) -> Result<Option<SubnetId>, ClientError> {
let endpoint = "read/get_subnet";
let result: Option<RawSubnetId> = self
.post(
endpoint,
RawCanisterId {
canister_id: canister_id.as_slice().to_vec(),
},
)
.await?;
Ok(result.map(|RawSubnetId { subnet_id }| SubnetId::from_slice(&subnet_id)))
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id, subnet_id = %subnet_id.to_string()))]
pub async fn get_subnet_metrics(
&self,
subnet_id: Principal,
) -> Result<Option<SubnetMetrics>, ClientError> {
let path = vec![
"subnet".into(),
Label::from_bytes(subnet_id.as_slice()),
"metrics".into(),
];
let paths = vec![path.clone()];
let content = ReadState {
ingress_expiry: self.get_time().await?.as_nanos_since_unix_epoch() + 240_000_000_000,
sender: Principal::anonymous(),
paths,
};
let envelope = Envelope {
content: std::borrow::Cow::Borrowed(&content),
sender_pubkey: None,
sender_sig: None,
sender_delegation: None,
};
let mut serialized_bytes = Vec::new();
let mut serializer = serde_cbor::Serializer::new(&mut serialized_bytes);
serializer.self_describe()?;
envelope.serialize(&mut serializer)?;
let endpoint = format!("api/v2/subnet/{}/read_state", subnet_id.to_text());
let url = self.instance_url().join(&endpoint).unwrap();
let mut exchange = Exchange::new(&url, self.max_request_time_ms);
let (status, bytes) = exchange
.bytes(
&self.reqwest_client,
self.reqwest_client
.post(url.clone())
.header(reqwest::header::CONTENT_TYPE, "application/cbor")
.body(serialized_bytes),
)
.await?;
if !status.is_success() {
return Err(exchange.server_error(status, &bytes));
}
let read_state_response: ReadStateResponse = serde_cbor::from_slice(&bytes)?;
let cert: Certificate = serde_cbor::from_slice(&read_state_response.certificate)?;
let metrics = match cert.tree.lookup_path(path) {
LookupResult::Found(value) => value,
_ => return Ok(None),
};
serde_cbor::from_slice(metrics).map_err(ClientError::from)
}
pub async fn drop(mut self) -> Result<(), ClientError> {
self.do_drop().await?;
Ok(())
}
pub(crate) async fn do_drop(&mut self) -> Result<(), ClientError> {
self.cleanup_attempted = true;
if self.owns_instance {
let url = self.instance_url();
let mut exchange = Exchange::new(&url, self.max_request_time_ms);
let (status, bytes) = exchange
.bytes(
&self.reqwest_client,
self.reqwest_client.delete(url.clone()),
)
.await?;
if !status.is_success() {
return Err(exchange.server_error(status, &bytes));
}
self.owns_instance = false;
Ok(())
} else {
self.stop_http_gateway().await
}
}
async fn get<T: DeserializeOwned>(&self, endpoint: &str) -> Result<T, ClientError> {
self.request(HttpMethod::Get, endpoint, ()).await
}
async fn post<T: DeserializeOwned, B: Serialize>(
&self,
endpoint: &str,
body: B,
) -> Result<T, ClientError> {
self.request(HttpMethod::Post, endpoint, body).await
}
async fn request<T: DeserializeOwned, B: Serialize>(
&self,
method: HttpMethod,
endpoint: &str,
body: B,
) -> Result<T, ClientError> {
self.request_limited(method, endpoint, body, None).await
}
async fn request_limited<T: DeserializeOwned, B: Serialize>(
&self,
method: HttpMethod,
endpoint: &str,
body: B,
response_limit: Option<usize>,
) -> Result<T, ClientError> {
let url = self.instance_url().join(endpoint).unwrap();
let mut exchange = Exchange::new(&url, self.max_request_time_ms);
exchange.response_limit = response_limit;
loop {
let builder = if let Some(operation) = &exchange.context.operation {
self.reqwest_client.get(
self.server_url
.join(&format!(
"/read_graph/{}/{}",
operation.state_label, operation.op_id
))
.unwrap(),
)
} else {
match method {
HttpMethod::Get => self.reqwest_client.get(url.clone()),
HttpMethod::Post => self.reqwest_client.post(url.clone()).json(&body),
}
};
let (status, bytes) = exchange.bytes(&self.reqwest_client, builder).await?;
match status {
StatusCode::OK => {
let value = exchange.decode(&bytes)?;
if let Some(operation) = &exchange.context.operation {
let request = self.reqwest_client.delete(
self.server_url
.join(&format!(
"/prune_graph/{}/{}",
operation.state_label, operation.op_id
))
.unwrap(),
);
tokio::spawn(async move {
let _ = request.send().await;
});
}
return Ok(value);
}
StatusCode::ACCEPTED | StatusCode::CONFLICT => {
let receipt: crate::pocket_ic::common::rest::StartedOrBusyResponse =
exchange.decode(&bytes)?;
if exchange.context.operation.is_none() {
if status == StatusCode::ACCEPTED {
exchange.started(Operation {
state_label: receipt.state_label,
op_id: receipt.op_id,
});
} else {
exchange.context.may_have_executed = false;
}
}
}
StatusCode::NOT_FOUND if exchange.context.operation.is_some() => {}
_ => return Err(exchange.server_error(status, &bytes)),
}
exchange.pause().await?;
}
}
async fn canister_call(
&self,
endpoint: &str,
effective_principal: RawEffectivePrincipal,
canister_id: CanisterId,
sender: Principal,
method: &str,
payload: Vec<u8>,
sender_info: Option<RawSenderInfo>,
) -> Result<Vec<u8>, ClientError> {
let raw_canister_call = RawCanisterCall {
sender: sender.as_slice().to_vec(),
canister_id: canister_id.as_slice().to_vec(),
method: method.to_string(),
payload,
effective_principal,
sender_info,
};
let result: RawCanisterResult = self.post(endpoint, raw_canister_call).await?;
Result::<Vec<u8>, RejectResponse>::from(result).map_err(ClientError::from)
}
pub async fn update_call_with_effective_principal(
&self,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<Vec<u8>, ClientError> {
self.update_call_with_effective_principal_and_sender_info_inner(
canister_id,
effective_principal,
sender,
method,
payload,
None,
)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn get_canister_http(&self) -> Result<Vec<CanisterHttpRequest>, ClientError> {
let endpoint = "read/get_canister_http";
let res: Vec<RawCanisterHttpRequest> = self.get(endpoint).await?;
Ok(res.into_iter().map(|r| r.into()).collect())
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn mock_canister_http_response(
&self,
mock_canister_http_response: MockCanisterHttpResponse,
) -> Result<(), ClientError> {
let endpoint = "update/mock_canister_http";
let raw_mock_canister_http_response: RawMockCanisterHttpResponse =
mock_canister_http_response.into();
self.post(endpoint, raw_mock_canister_http_response).await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn mock_flexible_canister_http_response(
&self,
mock_flexible_canister_http_response: MockFlexibleCanisterHttpResponse,
) -> Result<(), ClientError> {
let endpoint = "update/mock_flexible_canister_http";
let raw_mock_flexible_canister_http_response: RawMockFlexibleCanisterHttpResponse =
mock_flexible_canister_http_response.into();
self.post(endpoint, raw_mock_flexible_canister_http_response)
.await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn canister_snapshot_download(
&self,
canister_id: CanisterId,
sender: Principal,
snapshot_id: Vec<u8>,
snapshot_dir: PathBuf,
) -> Result<(), ClientError> {
let endpoint = "update/canister_snapshot_download";
#[cfg(not(windows))]
let snapshot_dir = snapshot_dir;
#[cfg(windows)]
let snapshot_dir = wsl_path(&snapshot_dir, "snapshot directory").into();
let raw_canister_snapshot_download = RawCanisterSnapshotDownload {
sender: sender.into(),
canister_id: canister_id.into(),
snapshot_id,
snapshot_dir,
};
self.post(endpoint, raw_canister_snapshot_download).await
}
#[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
pub async fn canister_snapshot_upload(
&self,
canister_id: CanisterId,
sender: Principal,
replace_snapshot: Option<Vec<u8>>,
snapshot_dir: PathBuf,
) -> Result<Vec<u8>, ClientError> {
let endpoint = "update/canister_snapshot_upload";
let replace_snapshot =
replace_snapshot.map(|snapshot_id| RawCanisterSnapshotId { snapshot_id });
#[cfg(not(windows))]
let snapshot_dir = snapshot_dir;
#[cfg(windows)]
let snapshot_dir = wsl_path(&snapshot_dir, "snapshot directory").into();
let raw_canister_snapshot_upload = RawCanisterSnapshotUpload {
sender: sender.into(),
canister_id: canister_id.into(),
replace_snapshot,
snapshot_dir,
};
Ok(self
.post::<RawCanisterSnapshotId, _>(endpoint, raw_canister_snapshot_upload)
.await?
.snapshot_id)
}
}
pub async fn call_candid_as<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
sender: Principal,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
with_candid(input, |payload| async {
env.update_call_with_effective_principal(
canister_id,
effective_principal,
sender,
method,
payload,
)
.await
})
.await
}
pub async fn call_candid<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
effective_principal: RawEffectivePrincipal,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
call_candid_as(
env,
canister_id,
effective_principal,
Principal::anonymous(),
method,
input,
)
.await
}
pub async fn query_candid<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
query_candid_as(env, canister_id, Principal::anonymous(), method, input).await
}
pub async fn query_candid_as<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
sender: Principal,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
with_candid(input, |bytes| async {
env.query_call(canister_id, sender, method, bytes).await
})
.await
}
pub async fn update_candid<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
update_candid_as(env, canister_id, Principal::anonymous(), method, input).await
}
pub async fn update_candid_as<Input, Output>(
env: &PocketIc,
canister_id: CanisterId,
sender: Principal,
method: &str,
input: Input,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
{
with_candid(input, |bytes| async {
env.update_call(canister_id, sender, method, bytes).await
})
.await
}
pub async fn with_candid<Input, Output, Fut>(
input: Input,
f: impl FnOnce(Vec<u8>) -> Fut,
) -> Result<Output, ClientError>
where
Input: ArgumentEncoder,
Output: for<'a> ArgumentDecoder<'a>,
Fut: Future<Output = Result<Vec<u8>, ClientError>>,
{
let input = encode_args(input)?;
let output = f(input).await?;
Ok(decode_args(&output)?)
}
fn setup_tracing(pid: u32) -> Option<WorkerGuard> {
use tracing_subscriber::prelude::*;
match std::env::var(LOG_DIR_PATH_ENV_NAME).map(std::path::PathBuf::from) {
Ok(p) => {
std::fs::create_dir_all(&p).expect("Could not create directory");
let file_name = format!("pocket_ic_client_{pid}.log");
let appender = tracing_appender::rolling::never(&p, file_name);
let (non_blocking_appender, guard) = tracing_appender::non_blocking(appender);
let log_dir_filter: EnvFilter =
tracing_subscriber::EnvFilter::try_from_env(LOG_DIR_LEVELS_ENV_NAME)
.unwrap_or_else(|_| "trace".parse().unwrap());
let layers = vec![
tracing_subscriber::fmt::layer()
.with_writer(non_blocking_appender)
.with_ansi(false)
.with_filter(log_dir_filter)
.boxed(),
];
let _ = tracing_subscriber::registry().with(layers).try_init();
Some(guard)
}
_ => None,
}
}
pub async fn get_default_effective_canister_id(
pocket_ic_url: String,
) -> Result<Principal, ClientError> {
let url = Url::parse(&format!(
"{}/_/topology",
pocket_ic_url.trim_end_matches('/')
))
.map_err(|e| ClientError::Protocol(e.to_string()))?;
let client = reqwest::Client::new();
let mut exchange = Exchange::new(&url, Some(crate::pocket_ic::DEFAULT_MAX_REQUEST_TIME_MS));
loop {
let (status, bytes) = exchange.bytes(&client, client.get(url.clone())).await?;
if status == StatusCode::CONFLICT {
exchange.pause().await?;
continue;
}
if !status.is_success() {
return Err(exchange.server_error(status, &bytes));
}
let topology: Topology = exchange.decode(&bytes)?;
return Principal::try_from_slice(&topology.default_effective_canister_id.canister_id)
.map_err(|e| ClientError::Protocol(e.to_string()));
}
}