lib.rs

// Copyright 2024 RISC Zero, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This application demonstrates how to send an off-chain proof request
// to the Bonsai proving service and publish the received proofs directly
// to your deployed app contract.

use alloy_primitives::{U256, Address, FixedBytes, utils::Keccak256};
use alloy_sol_types::{sol, SolInterface, SolValue};
use anyhow::{Context, Result};
//use clap::Parser;
//use ethers::prelude::*;
use methods::{SHA256_VERIFIER_ELF, KECCAK_VERIFIER_ELF};
//use risc0_ethereum_contracts::groth16;
use risc0_zkvm::{sha::Digestible, default_prover, ExecutorEnv, ProverOpts, VerifierContext, Groth16ReceiptVerifierParameters};
use godot::prelude::*;
use godot::engine::Engine;
use godot::classes::{Object, Node, Control};
use sha2::*;
use hex::*;
use zeroize::*;

#[gdextension]
unsafe impl ExtensionLibrary for RISC0 {
    fn on_level_init(level: InitLevel) {
        if level == InitLevel::Scene {
            // The StringName identifies your singleton and can be
            // used later to access it.
            Engine::singleton().register_singleton(
                StringName::from("RISC0"),
                RISC0::new_alloc().upcast(),
            );
        }
    }

    fn on_level_deinit(level: InitLevel) {
        if level == InitLevel::Scene {
            // Get the `Engine` instance and `StringName` for your singleton.
            let mut engine = Engine::singleton();
            let singleton_name = StringName::from("RISC0");

            // We need to retrieve the pointer to the singleton object,
            // as it has to be freed manually - unregistering singleton 
            // doesn't do it automatically.
            let singleton = engine
                .get_singleton(singleton_name.clone())
                .expect("cannot retrieve the singleton");

            // Unregistering singleton and freeing the object itself is needed 
            // to avoid memory leaks and warnings, especially for hot reloading.
            engine.unregister_singleton(singleton_name);
            singleton.free();
        }
    }
}


#[derive(GodotClass)]
#[class(init, base=Node)]
pub struct RISC0 {
    base: Base<Node>,
}

#[godot_api]
impl RISC0 {

#[func]
fn prove(mut node: Gd<Control>, _address: GString, secret: PackedByteArray, nullifier: PackedByteArray, _path: PackedStringArray, _hash_function: GString) {
    
    let hash_function = _hash_function.to_string();

    let address: Address = _address.to_string().parse().unwrap();

    let path: Vec<String> = _path.to_vec().iter().map(|e| String::from(e)).collect();

    let mut secret: FixedBytes<32> = packed_bytes_to_bytes32(secret);

    let mut nullifier: FixedBytes<32> = packed_bytes_to_bytes32(nullifier);

    let mut input_bytes = (address, secret, nullifier, path).abi_encode();

    secret.zeroize();

    nullifier.zeroize();
 
    let env = ExecutorEnv::builder().write_slice(&input_bytes).build().unwrap();

    input_bytes.zeroize();

    let prover = default_prover();

    let guest: &[u8];

    if hash_function == "SHA256" {
        guest = SHA256_VERIFIER_ELF;
    }
    else {
        guest = KECCAK_VERIFIER_ELF;
    }

    let prove_info = prover
        .prove_with_ctx(
            env, 
            &VerifierContext::default(),
            guest,
            &ProverOpts::groth16())
        .unwrap();
    
    let receipt = prove_info.receipt;

    // let _output = receipt.journal.bytes.clone();

    // Encode the seal with the selector.

    let seal = receipt.inner.groth16().unwrap().seal.clone();

    let verifier_parameters_digest = Groth16ReceiptVerifierParameters::default().digest();

    let selector = &verifier_parameters_digest.as_bytes()[..4];

    let mut selector_seal = Vec::with_capacity(selector.len() + seal.len());

    selector_seal.extend_from_slice(selector);

    selector_seal.extend_from_slice(&seal);

    let proof: PackedByteArray = selector_seal.iter().map(|e| *e as u8).collect();

    node.call("load_proof".into(), &[proof.to_variant()]);

}



#[func]
fn compute_merkle_root_from_path(_secret: PackedByteArray, _nullifier: PackedByteArray, packed_path: PackedStringArray, _hash_function: GString) -> PackedByteArray {

    let hash_function = _hash_function.to_string();

    // Hash the secret and nullifier, then hash the hashes together.   
    let mut secret = get_hash(_secret.to_vec(), hash_function.clone());
    let mut nullifier = get_hash(_nullifier.to_vec(), hash_function.clone());
 
    let secret_hash = hash_together(secret.to_vec().try_into().unwrap(), nullifier.to_vec().try_into().unwrap(), hash_function.clone());
   
    secret.zeroize();
    nullifier.zeroize();

    let _path: Vec<String> = packed_path.to_vec().iter().map(|e| String::from(e)).collect();
    let path: Vec<[u8; 32]> = _path
        .iter()
        .map(|x| hex::decode(x).unwrap().try_into().unwrap())
        .collect();


    // The merkle path is a flat array consisting of hashpairs.
    // Each hashpair will be concatenated and hashed, and the
    // produced hash must be present in the following hashpair.
    // The verification starts by checking whether the secret
    // hash is present in the first hashpair.
    let mut target_hash = secret_hash;

    let mut i = 0;
    while i < 62 {
        let hashpair = &path[i..i+2];
        assert!(hashpair.contains(&target_hash), "Hash not found in hashpair");
        target_hash = hash_together(hashpair[0], hashpair[1], hash_function.clone());
        i += 2;
    }

    // The root is computed from the final hashpair, after the integrity
    // of every hashpair has been validated.  This root must now pass
    // the validation check against the on-chain root history.
    let root: PackedByteArray = target_hash.iter().map(|e| *e as u8).collect();

    root

}


#[func]
fn keccak(data: PackedByteArray) -> PackedByteArray {
    let mut data = data.to_vec();
    let _hash = get_keccak_hash(data.clone());
    let hash: PackedByteArray = _hash.iter().map(|e| *e as u8).collect();
    data.zeroize();
    
    hash
}



}


fn packed_bytes_to_bytes32(_bytes: PackedByteArray) -> FixedBytes<32> {
    let bytes = _bytes.to_vec().try_into().unwrap();

    let bytes32: FixedBytes<32> = FixedBytes::new(bytes);

    bytes32

}



fn hash_together(first_hash: [u8; 32], second_hash: [u8; 32], hash_function: String) -> [u8; 32] {

    let mut first_hash = first_hash.to_vec();
    let mut second_hash = second_hash.to_vec();
    first_hash.append(&mut second_hash);
    let hash = get_hash(first_hash.clone(), hash_function);
    first_hash.zeroize();
    second_hash.zeroize();
    
    hash
}


fn get_hash(mut data: Vec<u8>, hash_function: String) -> [u8; 32] {
    let hash: [u8; 32];
    if hash_function == "SHA256" {
        hash = get_sha_hash(data.clone());
    }
    else {
        hash = get_keccak_hash(data.clone());
    }
    data.zeroize();

    hash

}

fn get_sha_hash(mut data: Vec<u8>) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(&data);
    let _hash = hasher.finalize();
    let hash: [u8; 32] = _hash.to_vec().try_into().unwrap();
    data.zeroize();
    
    hash
}

fn get_keccak_hash(mut data: Vec<u8>) -> [u8; 32] {
    let mut hasher = Keccak256::new();
    hasher.update(&data);
    let _hash = hasher.finalize();
    let hash: [u8; 32] = _hash.to_vec().try_into().unwrap();
    data.zeroize();
    
    hash
}