vetkd_public_key : (record {
  canister_id : opt canister_id;
  context : blob;
  key_id : record { curve : vetkd_curve; name : text };
}) -> (record { public_key : blob })

vetkd_derive_key : (record {
  input : blob;
  context : blob;
  transport_public_key : blob;
  key_id : record { curve : vetkd_curve; name : text };
}) -> (record { encrypted_key : blob })

[dependencies]
candid = "0.10"
ic-cdk = "0.19"
serde = { version = "1", features = ["derive"] }
serde_bytes = "0.11"

# High-level library (recommended) — source: https://github.com/dfinity/vetkeys
ic-vetkeys = "0.6"
ic-stable-structures = "0.7"

use candid::Principal;
use ic_cdk::update;
use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemory};
use ic_stable_structures::DefaultMemoryImpl;
use ic_vetkeys::key_manager::KeyManager;
use ic_vetkeys::types::{AccessRights, VetKDCurve, VetKDKeyId};

// KeyManager is generic over an AccessControl type — AccessRights is the default.
// It uses stable memory for persistent storage of access control state.
thread_local! {
    static MEMORY_MANAGER: std::cell::RefCell<MemoryManager<DefaultMemoryImpl>> =
        std::cell::RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));

    static KEY_MANAGER: std::cell::RefCell<Option<KeyManager<AccessRights>>> =
        std::cell::RefCell::new(None);
}

#[ic_cdk::init]
fn init() {
    let key_id = VetKDKeyId {
        curve: VetKDCurve::Bls12381G2,
        name: "key_1".to_string(), // "test_key_1" for local + mainnet testing
    };
    MEMORY_MANAGER.with(|mm| {
        let mm = mm.borrow();
        KEY_MANAGER.with(|km| {
            *km.borrow_mut() = Some(KeyManager::init(
                "my_app_v1",              // domain separator
                key_id,
                mm.get(MemoryId::new(0)), // config memory
                mm.get(MemoryId::new(1)), // access control memory
                mm.get(MemoryId::new(2)), // shared keys memory
            ));
        });
    });
}

#[update]
async fn get_encrypted_vetkey(subkey_id: Vec<u8>, transport_public_key: Vec<u8>) -> Vec<u8> {
    let caller = ic_cdk::caller(); // Capture BEFORE await
    let future = KEY_MANAGER.with(|km| {
        let km = km.borrow();
        let km = km.as_ref().expect("not initialized");
        km.get_encrypted_vetkey(caller, subkey_id, transport_public_key)
            .expect("access denied")
    });
    future.await
}

#[update]
async fn get_vetkey_verification_key() -> Vec<u8> {
    let future = KEY_MANAGER.with(|km| {
        let km = km.borrow();
        let km = km.as_ref().expect("not initialized");
        km.get_vetkey_verification_key()
    });
    future.await
}

use candid::{CandidType, Deserialize, Principal};
use ic_cdk::update;

#[derive(CandidType, Deserialize)]
struct VetKdKeyId {
    curve: VetKdCurve,
    name: String,
}

#[derive(CandidType, Deserialize)]
enum VetKdCurve {
    #[serde(rename = "bls12_381_g2")]
    Bls12381G2,
}

#[derive(CandidType)]
struct VetKdPublicKeyRequest {
    canister_id: Option<Principal>,
    context: Vec<u8>,
    key_id: VetKdKeyId,
}

#[derive(CandidType, Deserialize)]
struct VetKdPublicKeyResponse {
    public_key: Vec<u8>,
}

#[derive(CandidType)]
struct VetKdDeriveKeyRequest {
    input: Vec<u8>,
    context: Vec<u8>,
    transport_public_key: Vec<u8>,
    key_id: VetKdKeyId,
}

#[derive(CandidType, Deserialize)]
struct VetKdDeriveKeyResponse {
    encrypted_key: Vec<u8>,
}

const CONTEXT: &[u8] = b"my_app_v1";

fn key_id() -> VetKdKeyId {
    VetKdKeyId {
        curve: VetKdCurve::Bls12381G2,
        // Key names: "test_key_1" for local + mainnet testing, "key_1" for production
        name: "key_1".to_string(),
    }
}

#[update]
async fn vetkd_public_key() -> Vec<u8> {
    let request = VetKdPublicKeyRequest {
        canister_id: None, // defaults to this canister
        context: CONTEXT.to_vec(),
        key_id: key_id(),
    };

    // vetkd_public_key does not require cycles (unlike vetkd_derive_key).
    let (response,): (VetKdPublicKeyResponse,) = ic_cdk::api::call::call(
        Principal::management_canister(), // aaaaa-aa
        "vetkd_public_key",
        (request,),
    )
    .await
    .expect("vetkd_public_key call failed");

    response.public_key
}

#[update]
async fn vetkd_derive_key(transport_public_key: Vec<u8>) -> Vec<u8> {
    let caller = ic_cdk::caller(); // MUST capture before await

    let request = VetKdDeriveKeyRequest {
        input: caller.as_slice().to_vec(), // derive key specific to this caller
        context: CONTEXT.to_vec(),
        transport_public_key,
        key_id: key_id(),
    };

    // key_1 costs ~26B cycles, test_key_1 costs ~10B cycles.
    let (response,): (VetKdDeriveKeyResponse,) = ic_cdk::api::call::call_with_payment128(
        Principal::management_canister(),
        "vetkd_derive_key",
        (request,),
        26_000_000_000, // cycles for key_1 (use 10_000_000_000 for test_key_1)
    )
    .await
    .expect("vetkd_derive_key call failed");

    response.encrypted_key
}

[package]
name = "my-vetkd-app"
version = "0.1.0"

[dependencies]
core = "2.0.0"

import Blob "mo:core/Blob";
import Principal "mo:core/Principal";
import Text "mo:core/Text";

persistent actor {

  type VetKdCurve = { #bls12_381_g2 };

  type VetKdKeyId = {
    curve : VetKdCurve;
    name : Text;
  };

  type VetKdPublicKeyRequest = {
    canister_id : ?Principal;
    context : Blob;
    key_id : VetKdKeyId;
  };

  type VetKdPublicKeyResponse = {
    public_key : Blob;
  };

  type VetKdDeriveKeyRequest = {
    input : Blob;
    context : Blob;
    transport_public_key : Blob;
    key_id : VetKdKeyId;
  };

  type VetKdDeriveKeyResponse = {
    encrypted_key : Blob;
  };

  let managementCanister : actor {
    vetkd_public_key : VetKdPublicKeyRequest -> async VetKdPublicKeyResponse;
    vetkd_derive_key : VetKdDeriveKeyRequest -> async VetKdDeriveKeyResponse;
  } = actor "aaaaa-aa";

  let context : Blob = Text.encodeUtf8("my_app_v1");

  // Key names: "test_key_1" for local + mainnet testing, "key_1" for production
  func keyId() : VetKdKeyId {
    { curve = #bls12_381_g2; name = "key_1" }
  };

  public shared func getPublicKey() : async Blob {
    // vetkd_public_key does not require cycles (unlike vetkd_derive_key).
    let response = await managementCanister.vetkd_public_key({
      canister_id = null;
      context;
      key_id = keyId();
    });
    response.public_key
  };

  public shared ({ caller }) func deriveKey(transportPublicKey : Blob) : async Blob {
    // caller is captured here, before the await. vetkd_derive_key requires cycles.
    let response = await (with cycles = 26_000_000_000) managementCanister.vetkd_derive_key({
      input = Principal.toBlob(caller);
      context;
      transport_public_key = transportPublicKey;
      key_id = keyId();
    });
    response.encrypted_key
  };
};

import { TransportSecretKey, DerivedPublicKey, EncryptedVetKey } from "@dfinity/vetkeys";

// 1. Generate a transport secret key (BLS12-381)
const seed = crypto.getRandomValues(new Uint8Array(32));
const transportSecretKey = TransportSecretKey.fromSeed(seed);
const transportPublicKey = transportSecretKey.publicKey();

// 2. Request encrypted vetkey and verification key from your canister
const [encryptedKeyBytes, verificationKeyBytes] = await Promise.all([
  backendActor.get_encrypted_vetkey(subkeyId, transportPublicKey),
  backendActor.get_vetkey_verification_key(),
]);

// 3. Deserialize and decrypt
const verificationKey = DerivedPublicKey.deserialize(new Uint8Array(verificationKeyBytes));
const encryptedVetKey = EncryptedVetKey.deserialize(new Uint8Array(encryptedKeyBytes));
const vetKey = encryptedVetKey.decryptAndVerify(
  transportSecretKey,
  verificationKey,
  new Uint8Array(subkeyId),
);

// 4. Derive a symmetric key for AES-GCM
const aesKeyMaterial = vetKey.toDerivedKeyMaterial();
const aesKey = await crypto.subtle.importKey(
  "raw",
  aesKeyMaterial.data.slice(0, 32), // 256-bit AES key
  { name: "AES-GCM" },
  false,
  ["encrypt", "decrypt"],
);

// 5. Encrypt
const iv = crypto.getRandomValues(new Uint8Array(12));
const ciphertext = await crypto.subtle.encrypt(
  { name: "AES-GCM", iv },
  aesKey,
  new TextEncoder().encode("secret message"),
);

// 6. Decrypt
const plaintext = await crypto.subtle.decrypt(
  { name: "AES-GCM", iv },
  aesKey,
  ciphertext,
);

use ic_vetkeys::{MasterPublicKey, DerivedPublicKey};

// Start from the known mainnet master public key for key_1
let master_key = MasterPublicKey::for_mainnet_key("key_1")
    .expect("unknown key name");

// Derive the canister-level key
let canister_key = master_key.derive_canister_key(canister_id.as_slice());

// Derive a sub-key for a specific context/identity
let derived_key: DerivedPublicKey = canister_key.derive_sub_key(b"my_app_v1");

// Use derived_key for IBE encryption — no canister call needed

import { MasterPublicKey, DerivedPublicKey } from "@dfinity/vetkeys";

// Start from the known mainnet master public key
const masterKey = MasterPublicKey.productionKey();

// Derive the canister-level key
const canisterKey = masterKey.deriveCanisterKey(canisterId);

// Derive a sub-key for a specific context/identity
const derivedKey: DerivedPublicKey = canisterKey.deriveSubKey(
  new TextEncoder().encode("my_app_v1"),
);

// Use derivedKey for IBE encryption — no canister call needed

import {
  TransportSecretKey, DerivedPublicKey, EncryptedVetKey,
  IbeCiphertext, IbeIdentity, IbeSeed,
} from "@dfinity/vetkeys";

// --- Encrypt (sender side, no canister call needed) ---

// Derive the recipient's public key offline (see "Offline Public Key Derivation" above)
const recipientIdentity = IbeIdentity.fromBytes(recipientPrincipalBytes);
const seed = IbeSeed.random();
const plaintext = new TextEncoder().encode("secret message");

const ciphertext = IbeCiphertext.encrypt(derivedPublicKey, recipientIdentity, plaintext, seed);
const serialized = ciphertext.serialize(); // store or transmit this

// --- Decrypt (recipient side, requires canister call to get vetKey) ---

// 1. Get the vetKey (same flow as the Frontend section above)
const transportSecretKey = TransportSecretKey.fromSeed(crypto.getRandomValues(new Uint8Array(32)));
const [encryptedKeyBytes, verificationKeyBytes] = await Promise.all([
  backendActor.get_encrypted_vetkey(subkeyId, transportSecretKey.publicKey()),
  backendActor.get_vetkey_verification_key(),
]);
const verificationKey = DerivedPublicKey.deserialize(new Uint8Array(verificationKeyBytes));
const encryptedVetKey = EncryptedVetKey.deserialize(new Uint8Array(encryptedKeyBytes));
const vetKey = encryptedVetKey.decryptAndVerify(
  transportSecretKey, verificationKey, new Uint8Array(subkeyId),
);

// 2. Decrypt the IBE ciphertext
const deserialized = IbeCiphertext.deserialize(serialized);
const decrypted = deserialized.decrypt(vetKey);
// decrypted is Uint8Array containing "secret message"

use ic_vetkeys::{
    DerivedPublicKey, IbeCiphertext, IbeIdentity, IbeSeed, VetKey,
};

// --- Encrypt ---
let identity = IbeIdentity::from_bytes(recipient_principal.as_slice());
let seed = IbeSeed::new(&mut rand::rng());
let plaintext = b"secret message";

let ciphertext = IbeCiphertext::encrypt(
    &derived_public_key,
    &identity,
    plaintext,
    &seed,
);
let serialized = ciphertext.serialize();

// --- Decrypt (after obtaining the VetKey) ---
let deserialized = IbeCiphertext::deserialize(&serialized)
    .expect("invalid ciphertext");
let decrypted = deserialized.decrypt(&vet_key)
    .expect("decryption failed");
// decrypted == b"secret message"

# Start the local network (provisions test_key_1 and key_1 automatically)
icp network start -d

# Deploy your canister
icp deploy backend

# Test public key retrieval
icp canister call backend getPublicKey '()'
# Returns: (blob "...") -- the vetKD public key

# For derive_key, you need a transport public key (generated by frontend)
# Test with a dummy 48-byte blob:
icp canister call backend deriveKey '(blob "\00\01\02\03\04\05\06\07\08\09\0a\0b\0c\0d\0e\0f\10\11\12\13\14\15\16\17\18\19\1a\1b\1c\1d\1e\1f\20\21\22\23\24\25\26\27\28\29\2a\2b\2c\2d\2e\2f")'

# Deploy to mainnet
icp deploy backend -e ic

# Use test_key_1 for initial testing, key_1 for production
# Make sure your canister code references the correct key name

# 1. Verify public key is returned (non-empty blob)
icp canister call backend getPublicKey '()'
# Expected: (blob "\ab\cd\ef...")  -- 48+ bytes of BLS public key data

# 2. Verify derive_key returns encrypted key (non-empty blob)
icp canister call backend deriveKey '(blob "\00\01...")'
# Expected: (blob "\12\34\56...")  -- encrypted key material

# 3. Verify determinism: same (caller, context, input) and same transport key produce same encrypted_key
# Call deriveKey twice with the same identity and transport key
# Expected: identical encrypted_key blobs both times

# 4. Verify isolation: different callers get different keys
icp identity new test-user-1 --storage-mode=plaintext
icp identity new test-user-2 --storage-mode=plaintext
icp identity default test-user-1
icp canister call backend deriveKey '(blob "\00\01...")'
# Note the output
icp identity default test-user-2
icp canister call backend deriveKey '(blob "\00\01...")'
# Expected: DIFFERENT encrypted_key (different caller = different derived key)

# 5. Frontend integration test
# Open the frontend, trigger encryption/decryption
# Verify: encrypted data can be decrypted by the same user
# Verify: encrypted data CANNOT be decrypted by a different user