|
| 1 | +pub mod types; |
| 2 | +use candid::Principal; |
| 3 | +use ic_cdk::management_canister::{VetKDCurve, VetKDKeyId, VetKDPublicKeyArgs}; |
| 4 | +use ic_cdk::{init, query, update}; |
| 5 | +use ic_stable_structures::{ |
| 6 | + memory_manager::{MemoryId, MemoryManager, VirtualMemory}, |
| 7 | + Cell as StableCell, DefaultMemoryImpl, StableBTreeMap, |
| 8 | +}; |
| 9 | +use serde_bytes::ByteBuf; |
| 10 | +use std::cell::RefCell; |
| 11 | +use types::Signature; |
| 12 | + |
| 13 | +type Memory = VirtualMemory<DefaultMemoryImpl>; |
| 14 | + |
| 15 | +type VetKeyPublicKey = ByteBuf; |
| 16 | +type RawSignature = ByteBuf; |
| 17 | +type RawMessage = String; |
| 18 | +type Timestamp = u64; |
| 19 | + |
| 20 | +thread_local! { |
| 21 | + static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> = |
| 22 | + RefCell::new(MemoryManager::init(DefaultMemoryImpl::default())); |
| 23 | + |
| 24 | + static SIGNATURES: RefCell<StableBTreeMap<(Principal, Timestamp), Signature, Memory>> = RefCell::new(StableBTreeMap::init( |
| 25 | + MEMORY_MANAGER.with_borrow(|m| m.get(MemoryId::new(3))), |
| 26 | + )); |
| 27 | + |
| 28 | + static KEY_NAME: RefCell<StableCell<String, Memory>> = |
| 29 | + RefCell::new(StableCell::init( |
| 30 | + MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(2))), |
| 31 | + String::new(), |
| 32 | + ) |
| 33 | + .expect("failed to initialize key name")); |
| 34 | +} |
| 35 | + |
| 36 | +#[init] |
| 37 | +fn init(key_name_string: String) { |
| 38 | + KEY_NAME.with_borrow_mut(|key_name| { |
| 39 | + key_name |
| 40 | + .set(key_name_string) |
| 41 | + .expect("failed to set key name"); |
| 42 | + }); |
| 43 | +} |
| 44 | + |
| 45 | +#[update] |
| 46 | +async fn sign_message(message: RawMessage) -> RawSignature { |
| 47 | + let signer = ic_cdk::api::msg_caller(); |
| 48 | + let signature_bytes = ic_vetkeys::management_canister::sign_with_bls( |
| 49 | + message.as_bytes().to_vec(), |
| 50 | + context(&signer), |
| 51 | + key_id(), |
| 52 | + ) |
| 53 | + .await |
| 54 | + .expect("ic_vetkeys' sign_with_bls failed"); |
| 55 | + |
| 56 | + SIGNATURES.with_borrow_mut(|sigs| { |
| 57 | + let timestamp = ic_cdk::api::time(); |
| 58 | + let sig = Signature { |
| 59 | + message, |
| 60 | + signature: signature_bytes.clone(), |
| 61 | + timestamp, |
| 62 | + }; |
| 63 | + |
| 64 | + // In rare cases a user may call `sign_message` in quick succession |
| 65 | + // so that multiple signature requests are in a single consensus round, |
| 66 | + // which leads to ic_cdk::api::time() returning the same value for all |
| 67 | + // of the requests. In that case, we just keep increasing the timestamp |
| 68 | + // used in the map key until we hit a slot this is available. |
| 69 | + let mut timestamp_for_mapkey = timestamp; |
| 70 | + while sigs.get(&(signer, timestamp_for_mapkey)).is_some() { |
| 71 | + timestamp_for_mapkey += 1; |
| 72 | + } |
| 73 | + |
| 74 | + assert!(sigs.insert((signer, timestamp_for_mapkey), sig).is_none()); |
| 75 | + }); |
| 76 | + |
| 77 | + ByteBuf::from(signature_bytes) |
| 78 | +} |
| 79 | + |
| 80 | +#[query] |
| 81 | +fn get_my_signatures() -> Vec<Signature> { |
| 82 | + let me = ic_cdk::api::msg_caller(); |
| 83 | + SIGNATURES.with_borrow(|signer_and_timestamp_to_sig| { |
| 84 | + signer_and_timestamp_to_sig |
| 85 | + .range((me, 0)..) |
| 86 | + .take_while(|((signer, _ts), _sig)| signer == &me) |
| 87 | + .map(|((_, _), sig)| sig) |
| 88 | + .collect() |
| 89 | + }) |
| 90 | +} |
| 91 | + |
| 92 | +#[update] |
| 93 | +async fn get_my_verification_key() -> VetKeyPublicKey { |
| 94 | + let request = VetKDPublicKeyArgs { |
| 95 | + canister_id: None, |
| 96 | + context: context(&ic_cdk::api::msg_caller()), |
| 97 | + key_id: key_id(), |
| 98 | + }; |
| 99 | + let result = ic_cdk::management_canister::vetkd_public_key(&request) |
| 100 | + .await |
| 101 | + .expect("call to vetkd_public_key failed"); |
| 102 | + |
| 103 | + VetKeyPublicKey::from(result.public_key) |
| 104 | +} |
| 105 | + |
| 106 | +fn context(signer: &Principal) -> Vec<u8> { |
| 107 | + // A domain separator is not strictly necessary in this dapp, but having one is considered a good practice. |
| 108 | + const DOMAIN_SEPARATOR: [u8; 22] = *b"basic_bls_signing_dapp"; |
| 109 | + const DOMAIN_SEPARATOR_LENGTH: u8 = DOMAIN_SEPARATOR.len() as u8; |
| 110 | + [DOMAIN_SEPARATOR_LENGTH] |
| 111 | + .into_iter() |
| 112 | + .chain(DOMAIN_SEPARATOR) |
| 113 | + .chain(signer.as_ref().iter().cloned()) |
| 114 | + .collect() |
| 115 | +} |
| 116 | + |
| 117 | +fn key_id() -> VetKDKeyId { |
| 118 | + VetKDKeyId { |
| 119 | + curve: VetKDCurve::Bls12_381_G2, |
| 120 | + name: KEY_NAME.with_borrow(|key_name| key_name.get().clone()), |
| 121 | + } |
| 122 | +} |
| 123 | + |
| 124 | +// In the following, we register a custom getrandom implementation because |
| 125 | +// otherwise getrandom (which is a dependency of some other dependencies) fails to compile. |
| 126 | +// This is necessary because getrandom by default fails to compile for the |
| 127 | +// wasm32-unknown-unknown target (which is required for deploying a canister). |
| 128 | +// Our custom implementation always fails, which is sufficient here because |
| 129 | +// the used RNGs are _manually_ seeded rather than by the system. |
| 130 | +#[cfg(all( |
| 131 | + target_arch = "wasm32", |
| 132 | + target_vendor = "unknown", |
| 133 | + target_os = "unknown" |
| 134 | +))] |
| 135 | +getrandom::register_custom_getrandom!(always_fail); |
| 136 | +#[cfg(all( |
| 137 | + target_arch = "wasm32", |
| 138 | + target_vendor = "unknown", |
| 139 | + target_os = "unknown" |
| 140 | +))] |
| 141 | +fn always_fail(_buf: &mut [u8]) -> Result<(), getrandom::Error> { |
| 142 | + Err(getrandom::Error::UNSUPPORTED) |
| 143 | +} |
| 144 | + |
| 145 | +ic_cdk::export_candid!(); |
0 commit comments