quantus-cli 2.1.1

//! Wormhole Library Functions
//!
//! This module provides library-friendly functions for wormhole proof generation
//! that can be used by external crates (like quantus-sdk) without requiring
//! a chain client connection.
//!
//! These functions handle the core cryptographic operations:
//! - Computing leaf hashes for ZK Merkle proof verification
//! - Computing wormhole addresses from secrets
//! - Generating ZK proofs from raw inputs

use qp_wormhole_circuit::{
	inputs::{CircuitInputs, PrivateCircuitInputs},
	nullifier::Nullifier,
	sensitive::Secret,
};
use qp_wormhole_inputs::PublicCircuitInputs;
use qp_zk_circuits_common::{
	utils::{digest_to_bytes, BytesDigest},
	zk_merkle::SIBLINGS_PER_LEVEL,
};
use std::{
	mem::size_of,
	path::Path,
	ptr,
	sync::atomic::{compiler_fence, Ordering},
};

/// Native asset id for QTU token
pub const NATIVE_ASSET_ID: u32 = 0;

/// Scale down factor for quantizing amounts (10^10 to go from 12 to 2 decimal places)
pub const SCALE_DOWN_FACTOR: u128 = 10_000_000_000;

/// Volume fee rate in basis points (must match on-chain `VolumeFeeRateBps`).
/// Runtime currently sets 4 bps (0.04%).
pub const VOLUME_FEE_BPS: u32 = 4;

/// Result type for wormhole library operations
pub type Result<T> = std::result::Result<T, WormholeLibError>;

/// Error type for wormhole library operations
#[derive(Debug, Clone)]
pub struct WormholeLibError {
	pub message: String,
}

impl std::fmt::Display for WormholeLibError {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "{}", self.message)
	}
}

impl std::error::Error for WormholeLibError {}

impl From<String> for WormholeLibError {
	fn from(message: String) -> Self {
		Self { message }
	}
}

fn zeroize_bytes(bytes: &mut [u8]) {
	for byte in bytes {
		unsafe { ptr::write_volatile(byte, 0) };
	}
	compiler_fence(Ordering::SeqCst);
}

fn zeroize_bytes_digest(digest: &mut BytesDigest) {
	let ptr = ptr::addr_of_mut!(*digest).cast::<u8>();
	for offset in 0..size_of::<BytesDigest>() {
		unsafe { ptr.add(offset).write_volatile(0) };
	}
	compiler_fence(Ordering::SeqCst);
}

/// Zeroize-on-drop wrapper for a copy of the secret digest, so every exit path
/// out of proof generation (including early `?` returns) wipes the copy.
struct ZeroizingDigest(BytesDigest);

impl Drop for ZeroizingDigest {
	fn drop(&mut self) {
		zeroize_bytes_digest(&mut self.0);
	}
}

/// Zeroize-on-drop wrapper for the assembled circuit inputs, which embed a copy
/// of the secret in `private.secret`.
struct ZeroizingCircuitInputs(CircuitInputs);

impl Drop for ZeroizingCircuitInputs {
	fn drop(&mut self) {
		// `Secret` zeroizes on drop; replace now so the embedded copy is wiped
		// even if the outer `CircuitInputs` value is moved elsewhere later.
		self.0.private.secret = Secret::from(
			BytesDigest::try_from([0u8; 32]).expect("all-zero digest is always valid"),
		);
	}
}

/// Input data for generating a wormhole proof.
/// All fields are raw bytes - no chain client required.
#[derive(Debug, Clone)]
pub struct ProofGenerationInput {
	/// 32-byte secret
	pub secret: [u8; 32],
	/// Transfer count (atomic counter per recipient)
	pub transfer_count: u64,
	/// Wormhole address (recipient/unspendable account) as 32 bytes
	pub wormhole_address: [u8; 32],
	/// Input amount (quantized, 2 decimals) - from ZK leaf data
	pub input_amount: u32,
	/// Block hash as 32 bytes
	pub block_hash: [u8; 32],
	/// Block number
	pub block_number: u32,
	/// Parent hash as 32 bytes
	pub parent_hash: [u8; 32],
	/// State root as 32 bytes (still needed for block hash computation)
	pub state_root: [u8; 32],
	/// Extrinsics root as 32 bytes
	pub extrinsics_root: [u8; 32],
	/// SCALE-encoded digest (variable length, padded to 110 bytes internally)
	pub digest: Vec<u8>,
	/// ZK tree root (from block header's zk_tree_root field)
	pub zk_tree_root: [u8; 32],
	/// ZK Merkle proof siblings at each level (3 siblings per level, in sorted order)
	pub zk_merkle_siblings: Vec<[[u8; 32]; SIBLINGS_PER_LEVEL]>,
	/// Position hints (0-3) for each level
	pub zk_merkle_positions: Vec<u8>,
	/// Exit account 1 as 32 bytes
	pub exit_account_1: [u8; 32],
	/// Exit account 2 as 32 bytes (use zeros for single output)
	pub exit_account_2: [u8; 32],
	/// Output amount 1 (quantized, 2 decimals)
	pub output_amount_1: u32,
	/// Output amount 2 (quantized, 2 decimals, 0 for single output)
	pub output_amount_2: u32,
	/// Volume fee in basis points
	pub volume_fee_bps: u32,
	/// Asset ID (0 for native token)
	pub asset_id: u32,
}

/// Output from proof generation
#[derive(Debug, Clone)]
pub struct ProofGenerationOutput {
	/// Generated proof as bytes
	pub proof_bytes: Vec<u8>,
	/// Nullifier as 32 bytes (available for callers who need it)
	#[allow(dead_code)]
	pub nullifier: [u8; 32],
}

/// Compute the unspendable wormhole account from a secret.
///
/// # Arguments
/// * `secret` - 32-byte secret
///
/// # Returns
/// 32-byte wormhole account address
pub fn compute_wormhole_address(secret: &[u8; 32]) -> Result<[u8; 32]> {
	let secret_digest: BytesDigest = (*secret)
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid secret: {:?}", e)))?;

	let unspendable =
		qp_wormhole_circuit::unspendable_account::UnspendableAccount::from_secret(secret_digest);

	Ok(*digest_to_bytes(unspendable.account_id))
}

/// Compute the nullifier from secret and transfer count.
///
/// # Arguments
/// * `secret` - 32-byte secret
/// * `transfer_count` - Transfer counter
///
/// # Returns
/// 32-byte nullifier
#[allow(dead_code)]
pub fn compute_nullifier(secret: &[u8; 32], transfer_count: u64) -> Result<[u8; 32]> {
	let secret_digest: BytesDigest = (*secret)
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid secret: {:?}", e)))?;

	let nullifier = Nullifier::from_preimage(secret_digest, transfer_count);
	Ok(*digest_to_bytes(nullifier.hash))
}

/// Quantize a funding amount from 12 decimal places to 2 decimal places.
///
/// # Arguments
/// * `amount` - Amount in planck (12 decimals)
///
/// # Returns
/// Quantized amount (2 decimals) as u32
pub fn quantize_amount(amount: u128) -> Result<u32> {
	let quantized = amount / SCALE_DOWN_FACTOR;
	if quantized > u32::MAX as u128 {
		return Err(WormholeLibError::from(format!(
			"Quantized amount {} exceeds u32::MAX",
			quantized
		)));
	}
	Ok(quantized as u32)
}

/// Compute output amount after fee deduction.
///
/// output = input * (10000 - fee_bps) / 10000
pub fn compute_output_amount(input_amount: u32, fee_bps: u32) -> u32 {
	((input_amount as u64) * (10000 - fee_bps as u64) / 10000) as u32
}

/// Generate a wormhole proof from raw inputs.
///
/// This function takes all necessary data as raw bytes and generates a ZK proof.
/// It does not require a chain client - all data must be pre-fetched.
///
/// The leaf prover is built fresh via [`qp_wormhole_prover::build_fresh`]; there
/// is no leaf `prover.bin` artifact. The path arguments are retained only for
/// API compatibility with existing callers and are ignored.
///
/// # Arguments
/// * `input` - All input data for proof generation (including ZK Merkle proof). Borrowed mutably:
///   `input.secret` is zeroized before this function returns, on success and on every error path.
///   Callers that retry must rebuild the input with a fresh secret.
/// * `prover_bin_path` - Ignored (legacy; leaf prover is built in-process)
/// * `common_bin_path` - Ignored (legacy; leaf prover is built in-process)
///
/// # Returns
/// Proof bytes and nullifier
pub fn generate_proof(
	input: &mut ProofGenerationInput,
	prover_bin_path: &Path,
	common_bin_path: &Path,
) -> Result<ProofGenerationOutput> {
	// Leaf prover is built from the canonical circuit config (no longer loads prover.bin).
	// Paths are kept for API compatibility with callers that still pass bin locations.
	let _ = (prover_bin_path, common_bin_path);

	let result = generate_proof_inner(input);
	// Wipe the caller-visible secret unconditionally, on success and on every error path.
	zeroize_bytes(&mut input.secret);
	result
}

fn generate_proof_inner(input: &ProofGenerationInput) -> Result<ProofGenerationOutput> {
	// Perform every fallible conversion before the secret is copied anywhere, so an
	// early `?` return can never skip zeroization of a secret copy.
	let parent_hash = input
		.parent_hash
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid parent hash: {:?}", e)))?;
	let state_root = input
		.state_root
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid state root: {:?}", e)))?;
	let extrinsics_root = input
		.extrinsics_root
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid extrinsics root: {:?}", e)))?;
	let exit_account_1 = input
		.exit_account_1
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid exit account 1: {:?}", e)))?;
	let exit_account_2 = input
		.exit_account_2
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid exit account 2: {:?}", e)))?;
	let block_hash = input
		.block_hash
		.as_slice()
		.try_into()
		.map_err(|e| WormholeLibError::from(format!("Invalid block hash: {:?}", e)))?;

	// Convert secret to BytesDigest; the guard wipes this copy on every exit path.
	let secret_digest = ZeroizingDigest(
		input
			.secret
			.try_into()
			.map_err(|e| WormholeLibError::from(format!("Invalid secret: {:?}", e)))?,
	);

	// Compute nullifier
	let nullifier = Nullifier::from_preimage(secret_digest.0, input.transfer_count);
	let nullifier_bytes = digest_to_bytes(nullifier.hash);

	// Compute unspendable account
	let unspendable =
		qp_wormhole_circuit::unspendable_account::UnspendableAccount::from_secret(secret_digest.0);
	let unspendable_bytes = digest_to_bytes(unspendable.account_id);

	// Verify the wormhole address matches what we computed from the secret
	if *unspendable_bytes != input.wormhole_address {
		return Err(WormholeLibError::from(
			"Wormhole address doesn't match the computed unspendable account from secret"
				.to_string(),
		));
	}

	// Prepare digest (padded to 110 bytes)
	const DIGEST_LOGS_SIZE: usize = 110;
	let mut digest_padded = [0u8; DIGEST_LOGS_SIZE];
	let copy_len = input.digest.len().min(DIGEST_LOGS_SIZE);
	digest_padded[..copy_len].copy_from_slice(&input.digest[..copy_len]);

	// Build circuit inputs with ZK Merkle proof; the guard wipes the embedded
	// secret copy on every exit path.
	let circuit_inputs = ZeroizingCircuitInputs(CircuitInputs {
		public: PublicCircuitInputs {
			asset_id: input.asset_id,
			output_amount_1: input.output_amount_1,
			output_amount_2: input.output_amount_2,
			volume_fee_bps: input.volume_fee_bps,
			nullifier: nullifier_bytes,
			exit_account_1,
			exit_account_2,
			block_hash,
			block_number: input.block_number,
		},
		private: PrivateCircuitInputs {
			secret: Secret::from(secret_digest.0),
			transfer_count: input.transfer_count,
			unspendable_account: unspendable_bytes,
			parent_hash,
			state_root,
			extrinsics_root,
			digest: digest_padded,
			input_amount: input.input_amount,
			zk_tree_root: input.zk_tree_root,
			zk_merkle_siblings: input.zk_merkle_siblings.clone(),
			zk_merkle_positions: input.zk_merkle_positions.clone(),
		},
	});
	drop(secret_digest);
	zeroize_bytes(&mut digest_padded);

	let prover = qp_wormhole_prover::build_fresh();

	let prover_with_inputs = prover
		.commit(&circuit_inputs.0)
		.map_err(|e| WormholeLibError::from(format!("Failed to commit inputs: {}", e)))?;

	let proof = prover_with_inputs
		.prove()
		.map_err(|e| WormholeLibError::from(format!("Proof generation failed: {}", e)))?;

	Ok(ProofGenerationOutput { proof_bytes: proof.to_bytes(), nullifier: *nullifier_bytes })
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn test_quantize_amount() {
		// 1 QTU = 10^12 planck -> should quantize to 100 (1.00 with 2 decimals)
		let result = quantize_amount(1_000_000_000_000).unwrap();
		assert_eq!(result, 100);

		// 0.01 QTU = 10^10 planck -> should quantize to 1
		let result = quantize_amount(10_000_000_000).unwrap();
		assert_eq!(result, 1);
	}

	#[test]
	fn test_compute_output_amount() {
		// 100 input with 10 bps fee -> 99.9 -> 99
		let result = compute_output_amount(100, 10);
		assert_eq!(result, 99);

		// 1000 input with 10 bps fee -> 999
		let result = compute_output_amount(1000, 10);
		assert_eq!(result, 999);
	}

	#[test]
	fn test_compute_wormhole_address() {
		// Just verify it doesn't panic and returns 32 bytes
		let secret = [42u8; 32];
		let address = compute_wormhole_address(&secret).unwrap();
		assert_eq!(address.len(), 32);
		// Should be deterministic
		let address2 = compute_wormhole_address(&secret).unwrap();
		assert_eq!(address, address2);
	}

	fn decode_32(hex_str: &str) -> [u8; 32] {
		let bytes = hex::decode(hex_str).expect("valid hex fixture");
		bytes.try_into().expect("fixture is 32 bytes")
	}

	/// #160105: generate_proof must clear the caller-owned secret after use.
	#[test]
	fn secret_is_zeroized_after_successful_wormhole_proof_generation() {
		let secret = decode_32("4c8587bd422e01d961acdc75e7d66f6761b7af7c9b1864a492f369c9d6724f05");
		let transfer_count = 4u64;
		let wormhole_address = compute_wormhole_address(&secret).expect("secret derives address");

		let mut input = ProofGenerationInput {
			secret,
			transfer_count,
			wormhole_address,
			input_amount: 100,
			block_hash: [0u8; 32],
			block_number: 0,
			parent_hash: [0u8; 32],
			state_root: decode_32(
				"ae6e4ff0dca1ef5ede9dccc84365cecfab4e431c6f3086216bc3b819cdf0a893",
			),
			extrinsics_root: [0u8; 32],
			digest: vec![
				8, 6, 112, 111, 119, 95, 128, 233, 182, 183, 107, 158, 1, 115, 19, 219, 126, 253,
				86, 30, 208, 176, 70, 21, 45, 180, 229, 9, 62, 91, 4, 6, 53, 245, 52, 48, 38, 123,
				225, 5, 112, 111, 119, 95, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
				0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 18, 79, 226,
			],
			zk_tree_root: [0u8; 32],
			zk_merkle_siblings: vec![],
			zk_merkle_positions: vec![],
			exit_account_1: [0u8; 32],
			exit_account_2: [0u8; 32],
			output_amount_1: 0,
			output_amount_2: 0,
			volume_fee_bps: VOLUME_FEE_BPS,
			asset_id: NATIVE_ASSET_ID,
		};

		let output = generate_proof(
			&mut input,
			Path::new("ignored-prover.bin"),
			Path::new("ignored-common.bin"),
		)
		.expect("real wormhole proof generation succeeds");

		assert!(!output.proof_bytes.is_empty(), "the real prover produced a proof");
		assert_eq!(
			input.secret, [0u8; 32],
			"generate_proof must zeroize the caller-owned secret before returning"
		);
	}

	/// The secret must also be wiped on error paths, e.g. a wormhole address that
	/// does not match the secret.
	#[test]
	fn secret_is_zeroized_when_proof_generation_fails_early() {
		let secret = decode_32("4c8587bd422e01d961acdc75e7d66f6761b7af7c9b1864a492f369c9d6724f05");

		let mut input = ProofGenerationInput {
			secret,
			transfer_count: 0,
			// Deliberately not the address derived from `secret`.
			wormhole_address: [0xAAu8; 32],
			input_amount: 100,
			block_hash: [0u8; 32],
			block_number: 0,
			parent_hash: [0u8; 32],
			state_root: [0u8; 32],
			extrinsics_root: [0u8; 32],
			digest: vec![],
			zk_tree_root: [0u8; 32],
			zk_merkle_siblings: vec![],
			zk_merkle_positions: vec![],
			exit_account_1: [0u8; 32],
			exit_account_2: [0u8; 32],
			output_amount_1: 0,
			output_amount_2: 0,
			volume_fee_bps: VOLUME_FEE_BPS,
			asset_id: NATIVE_ASSET_ID,
		};

		let err = generate_proof(
			&mut input,
			Path::new("ignored-prover.bin"),
			Path::new("ignored-common.bin"),
		)
		.expect_err("mismatched wormhole address must be rejected");

		assert!(
			err.message.contains("doesn't match"),
			"expected address-mismatch error, got: {}",
			err.message
		);
		assert_eq!(
			input.secret, [0u8; 32],
			"generate_proof must zeroize the caller-owned secret on error paths too"
		);
	}
}