use super::{
RegistryFetchError,
proto::{
LargeValueChunkKeys, RegistryErrorCode, RegistryGetLatestVersionResponse,
RegistryGetValueRequest, RegistryGetValueResponse, UInt64Value,
registry_get_value_response,
},
wire::{RegistryChunk, RegistryGetChunkRequest, RegistryValueContent},
};
pub(super) use crate::hex::hex_bytes;
use candid::{Decode, Encode, Principal};
use ic_agent::Agent;
use prost::Message;
use sha2::{Digest, Sha256};
pub(super) async fn get_latest_version(
agent: &Agent,
registry_canister: &Principal,
) -> Result<u64, RegistryFetchError> {
let bytes = agent
.query(registry_canister, "get_latest_version")
.with_arg(Vec::<u8>::new())
.call()
.await
.map_err(|err| RegistryFetchError::AgentCall {
method: "get_latest_version",
reason: err.to_string(),
})?;
let response = decode_message::<RegistryGetLatestVersionResponse>(
"RegistryGetLatestVersionResponse",
&bytes,
)?;
Ok(response.version)
}
pub(super) async fn get_registry_value(
agent: &Agent,
registry_canister: &Principal,
key: &str,
version: u64,
) -> Result<Vec<u8>, RegistryFetchError> {
let request = RegistryGetValueRequest {
version: Some(UInt64Value { value: version }),
key: key.as_bytes().to_vec(),
};
let mut arg = Vec::new();
request
.encode(&mut arg)
.map_err(|err| RegistryFetchError::ProtobufEncode {
message: "RegistryGetValueRequest",
reason: err.to_string(),
})?;
let bytes = agent
.query(registry_canister, "get_value")
.with_arg(arg)
.call()
.await
.map_err(|err| RegistryFetchError::AgentCall {
method: "get_value",
reason: err.to_string(),
})?;
let response = decode_message::<RegistryGetValueResponse>("RegistryGetValueResponse", &bytes)?;
match registry_value_content_from_response(key, response)? {
RegistryValueContent::Value(value) => Ok(value),
RegistryValueContent::LargeValueChunkKeys(keys) => {
get_large_registry_value(agent, registry_canister, &keys).await
}
}
}
pub(super) fn registry_value_content_from_response(
key: &str,
response: RegistryGetValueResponse,
) -> Result<RegistryValueContent, RegistryFetchError> {
if let Some(error) = response.error {
return Err(RegistryFetchError::RegistryValue {
key: key.to_string(),
code: registry_error_code(error.code).to_string(),
reason: error.reason,
});
}
match response.content {
Some(registry_get_value_response::Content::Value(value)) => {
Ok(RegistryValueContent::Value(value))
}
Some(registry_get_value_response::Content::LargeValueChunkKeys(keys)) => {
Ok(RegistryValueContent::LargeValueChunkKeys(keys))
}
None => Err(RegistryFetchError::MissingValue {
key: key.to_string(),
}),
}
}
async fn get_large_registry_value(
agent: &Agent,
registry_canister: &Principal,
keys: &LargeValueChunkKeys,
) -> Result<Vec<u8>, RegistryFetchError> {
let mut value = Vec::new();
for chunk_sha256 in &keys.chunk_content_sha256s {
let chunk_content = get_registry_chunk(agent, registry_canister, chunk_sha256).await?;
append_validated_chunk(&mut value, chunk_sha256, chunk_content)?;
}
Ok(value)
}
async fn get_registry_chunk(
agent: &Agent,
registry_canister: &Principal,
content_sha256: &[u8],
) -> Result<Vec<u8>, RegistryFetchError> {
let request = RegistryGetChunkRequest {
content_sha256: Some(content_sha256.to_vec()),
};
let arg = Encode!(&request).map_err(|err| RegistryFetchError::CandidEncode {
message: "RegistryGetChunkRequest",
reason: err.to_string(),
})?;
let bytes = agent
.query(registry_canister, "get_chunk")
.with_arg(arg)
.call()
.await
.map_err(|err| RegistryFetchError::AgentCall {
method: "get_chunk",
reason: err.to_string(),
})?;
let result = Decode!(&bytes, Result<RegistryChunk, String>).map_err(|err| {
RegistryFetchError::CandidDecode {
message: "Result<RegistryChunk, String>",
reason: err.to_string(),
}
})?;
match result {
Ok(chunk) => chunk
.content
.ok_or_else(|| RegistryFetchError::MissingChunkContent {
sha256: hex_bytes(content_sha256),
}),
Err(reason) => Err(RegistryFetchError::RegistryChunkRejected {
sha256: hex_bytes(content_sha256),
reason,
}),
}
}
pub(super) fn append_validated_chunk(
value: &mut Vec<u8>,
expected_sha256: &[u8],
chunk_content: Vec<u8>,
) -> Result<(), RegistryFetchError> {
let actual_sha256 = sha256_digest(&chunk_content);
if actual_sha256.as_slice() != expected_sha256 {
return Err(RegistryFetchError::ChunkHashMismatch {
sha256: hex_bytes(expected_sha256),
actual_sha256: hex_bytes(&actual_sha256),
});
}
value.extend(chunk_content);
Ok(())
}
pub(super) fn sha256_digest(bytes: &[u8]) -> [u8; 32] {
Sha256::digest(bytes).into()
}
pub(super) fn decode_message<T>(
message: &'static str,
bytes: &[u8],
) -> Result<T, RegistryFetchError>
where
T: Message + Default,
{
T::decode(bytes).map_err(|err| RegistryFetchError::ProtobufDecode {
message,
reason: err.to_string(),
})
}
fn registry_error_code(code: i32) -> &'static str {
match RegistryErrorCode::try_from(code).ok() {
Some(RegistryErrorCode::MalformedMessage) => "malformed_message",
Some(RegistryErrorCode::KeyNotPresent) => "key_not_present",
Some(RegistryErrorCode::KeyAlreadyPresent) => "key_already_present",
Some(RegistryErrorCode::VersionNotLatest) => "version_not_latest",
Some(RegistryErrorCode::VersionBeyondLatest) => "version_beyond_latest",
Some(RegistryErrorCode::Authorization) => "authorization",
Some(RegistryErrorCode::InternalError) => "internal_error",
None => "unknown",
}
}