Zero-Knowledge proofs in Rust using the RISC Zero Virtual Machine

Posted underTech Talks
Cover Image for Zero-Knowledge proofs in Rust using the RISC Zero Virtual Machine

The concept of zero-knowledge proofs (ZKPs) has been getting quite popular recently, especially in the blockchain space. Despite that, they haven’t seen a lot of adoption in real-world applications. There are certain limitations on computation and architectural requirements that make it not so straightforward to incorporate them into an existing project, but that also depends on the context. ZCash was the first project and probably the most popular to incorporate them, but not without a significant engineering effort and more importantly – they are a part of the core system, not something that was tacked on. Let’s explore ZKPs, particularly in the context of the Risc Zero virtual machine, by building a simple tic-tac-toe game. We are not going to go through all the code in detail, only certain parts that are relevant to the topic. But you are free to view the full codebase here: https://github.com/hackbg/zk-tic-tac-toe

ZKPs allow one party (the prover) to prove to another (the verifier) that a given statement is true, without revealing any additional information. The Risc Zero VM allows the two parties to agree upon a function which is executed within the VM. The prover will run the function inside the VM, giving it some input, and it will generate a public output and a “receipt” of that function’s correct execution. The prover can then send that receipt to the verifier who will check whether the function ran correctly and that it produced a specific output. In our case, we will have a “server” that will act as a prover and the two players will be the verifiers.

Risc0

We will be using the Risc Zero CLI tool for Cargo (Rust’s package manager). Running cargo risczero new tic-tac-toe generates an empty pre-configured project structure. The methods folder contains the functions that will run inside the Risc Zero VM.

In our case this will be the function that executes a single player move in the game. The CLI already generated the necessary boilerplate that will turn our guest function code into a RISC-V ELF binary that will be executed inside the VM. The binary is also hashed to produce an image ID that uniquely identifies that function. This is also auto-generated by the boilerplate code. The methods crate (Rust module) will then export those two as two constants {NAME}_ELF and {NAME}_ID that refer to the RISC-V binary and its hash respectively. The {NAME} placeholder is the name of the Rust crate in which the function is implemented, which is the guest crate located in methods/guest/Cargo.toml.

We will change the name property inside that file to “make_move” to better reflect its functionality. Thus the two constants will now be named MAKE_MOVE_ELF and MAKE_MOVE_ID. We will later use these in our host code. Next we will create a new crate (using cargo new game --vcs none) in the root directory which will define the game state and implementation. We will be needing these two dependencies in our new crate:

[dependencies]
risc0-zkvm = { version = "0.15.1", default-features = false, features = ["std"] }
serde = { version = "1.0", default-features = false }

We define the following structures for our game state:

const CELL_COUNT: usize = 3;

#[repr(C)]
#[derive(Serialize, Deserialize, Clone, Copy, Debug)]
pub struct TicTacToe {
    board: [[Cell; CELL_COUNT]; CELL_COUNT],
    previous: Player,
    state: State
}

#[repr(u8)]
#[derive(Serialize, Deserialize, Clone, Copy, PartialEq, Debug)]
enum Cell {
    Player1,
    Player2,
    Vacant
}

#[repr(u8)]
#[derive(Serialize, Deserialize, Clone, Copy, PartialEq, Debug)]
pub enum Player {
    A,
    B
}

#[derive(Serialize, Deserialize, Clone, Copy, PartialEq, Debug)]
pub enum State {
    InProgress,
    Stalemate,
    Winner(Player)
}

The TicTacToe struct holds a multidimensional array that represents the state of each cell on the board. The previous field is the player whose turn it was last and the state field is the state at which the game currently is. Notice the #[repr(...)] directives. We use these because we want to be able to represent the game entire state as a single contiguous byte array which we can then easily use to produce a hash of. The guest method will use this to produce a hash of the game state that was provided as an input before running the game logic. The players will then compare that hash to the previous one they got (produced during the previous move) to verify that the state wasn’t somehow altered.

We use #[repr(u8)] on the Cell and Player enums which will cause their memory representation to be a single byte. Knowing this, the Rust compiler does a great job at making the State enum also be a single byte in memory even though it contains a variant with a payload. Finally, #[repr(C)] on the TicTacToe struct will use the C ABI which is well-defined, unlike the Rust one. This means that the memory representation of the struct will be all of its fields in order of definition. Since we made sure that all the fields have an alignment of 1 (the enums being a single byte and thus the board array also becoming an array of bytes which has an alignment of 1 as well) we know that no padding will be inserted by the compiler and can easily cast the struct to a raw byte array where each byte means something. We can now implement a method that does the cast:

impl TicTacToe {
    pub fn as_bytes(&self) -> [
        u8;
        (CELL_COUNT * CELL_COUNT) +
        mem::size_of::<Player>() +
        mem::size_of::<State>()
    ] {
        // Assert that the struct contains no padding.
        assert_eq!(mem::align_of::<TicTacToe>(), 1);

        unsafe { mem::transmute(*self) }
    }
}

The Rust compiler actually verifies at compile time that the length of the array that we are returning indeed matches the size of our struct. If we tried, for example, adding “+ 1” to the return type, we’d get a compile error. Next we will implement the make_move function that has the following signature:

pub fn make_move(&mut self, point: Point) -> Result<(), MoveError> { ... }

We skip the implementation as it simply implements the game logic. A method that displays the game is also needed, we call it print_board. Now we can move ahead to implementing the guest method, i.e the function that runs inside the Risk Zero VM. We will just need a struct that represents the output of that. We also define it in the game crate as it is used by the host as well:

#[derive(Serialize, Deserialize, Debug)]
pub struct VmResponse {
    pub game: TicTacToe,
    pub prev_state_hash: Digest
}

The game field contains the game state after the move has been made and prev_state_hash is the hash of the state before that move was made. The contents of our guest method (in methods/guest/src/main.rs) look like this:

#![no_main]

use risc0_zkvm::{
    guest::env,
    sha::{Impl, Sha256}
};
use game::{VmResponse, TicTacToe, Point};

risc0_zkvm::guest::entry!(main);

pub fn main() {
    let mut game: TicTacToe = env::read();
    let point: Point = env::read();

    let prev_state_hash = *Impl::hash_bytes(&game.as_bytes());

    game.make_move(point).unwrap();

    env::commit(&VmResponse {
        game,
        prev_state_hash
    });
}

The method reads the current game state and the position of the cell where the current player wants to place their symbol in. It first hashes the state before mutating it by running the game logic and includes it in the response after doing so along with the hash. Now we can move onto implementing the host code which will also be the entry point of our program. The game loop will consist of displaying the state, gathering input from the current player, executing the move inside the VM and then sending the receipt of the execution to both players for them to verify. We will only go over the code that deals with the executing the VM, producing the receipt end verifying it specifically. For starters, we define these two structs that will simulate the server and clients:

struct Server {
    game: TicTacToe
}

struct Client {
    game_state: State,
    state_hash: Digest
}

The client has the last state change hash which it compares with the one from the output of the receipt and then hashes the new state. It keeps a State instance (which it again gets from the VM output) to check whether the game has finished. The initial state is always known, which is an empty board where the next player is A and the game state is State::InProgress. So we don’t need to compute that in the VM, we simply have a convenience method on TicTacToe that we use when creating an instance of the client:

pub fn initial_hash() -> Digest {
    let bytes = Self::new().as_bytes();

    *Impl::hash_bytes(&bytes)
}

Then creating a new client looks like this:

pub fn new() -> Self {
    Self {
        state_hash: TicTacToe::initial_hash(),
        game_state: State::InProgress
    }
}

Now when the next move happens and the server sends us the receipt, we verify that the new state was derived from that initial state. The new state is then hashed and used to verify the move after that and so on… We don’t really need to verify whether the server executed a different move than the one we wanted since, the game state returned in the receipt can show us that visually. Maybe the state that needs to be verified would be different in a real-world scenario but that would also depend on the context and architecture, so we simplify tremendously in our example. Anyways, the player verification method looks like this and does exactly what was described above:

pub fn verify_receipt(&mut self, receipt: &SessionReceipt) {
    assert_eq!(self.game_state, State::InProgress, "Game has already ended!");

    receipt.verify(MAKE_MOVE_ID)
        .expect("receipt verification failed");

    let resp: VmResponse = from_slice(&receipt.journal).unwrap();
    assert_eq!(self.state_hash, resp.prev_state_hash, "Game state hash mismatch!");

    self.game_state = resp.game.state();
    self.state_hash = *Impl::hash_bytes(&resp.game.as_bytes());
}

We use that MAKE_MOVE_ID constant that we talked about before to make sure that this was indeed the method that was executed. Notice that we also assert that the game was still on-going in the previous turn. Conversely, the logic that produces that receipt (executed by the server) looks like this:

pub fn execute_move(&self, point: Point) -> Result<SessionReceipt> {
    let env = ExecutorEnv::builder()
        .add_input(&to_vec(&self.game)?)
        .add_input(&to_vec(&point)?)
        .build();

    let mut executor = Executor::from_elf(env, MAKE_MOVE_ELF)?;
    let session = executor.run()?;

    session.prove()
}

It passes the current game state and the coordinates of the cell that the current player wants to occupy as input to the VM. Then executes the binary that we created (MAKE_MOVE_ELF) and finally generates a proof that it then passes on to the clients. With this, the whole thing comes together and we can finally play the game.

While the integration with Risc Zero is relatively straightforward, there are some considerations and limitations to take into account. Apart from the architectural decisions that need to be taken into account that we mentioned in the beginning, there is also the consideration of what data needs to be verifiable and what functionality needs to be executed inside the VM.

These, of course, strongly depend on the particular scenario. There is also the constrain the Risc Zero requires that no-std is used. This means that you can’t use the standard library, but only the core one. Although, there is experimental support for std which we actually use in our example, so that might not be a problem in the future.

It also seems that the utilization of hashes is a necessity in most cases which apart from being a computationally intensive task, also has the requirement that whatever data is being hashed has to be representable as a byte array. In our case, that’s simple as our game state is really small but for more complex programs that might become a problem, though converting to JSON could be used if necessary.

A more significant limitation is that the program size that is to run inside the VM needs to be kept to a minimum, the recommended size being no larger than a megabyte. Furthermore, generating the proof has an extreme computational cost. For reference, generating the receipt for a single turn in our tic-tac-toe game takes ~10 seconds on a Ryzen 7 2700X CPU. There is the option to offload this to the GPU for CUDA-enabled GPUs and Metal on MacOS, but obviously specialized hardware is required for those options to be enabled.

These limitations are also the reason why we can’t just embed the VM in a smart contract, using CosmWasm for example (since it’s also Rust-based), and have it prove stuff for us. While the the smart contract program size is usually small, the computational cost to generate a proof itself will probably exceed the block limit. There is also the fact that the Risc Zero VM is a sandbox, but the smart contract which itself runs in a WASM VM needs to have access to particular foreign functions and syscalls. For these reasons, smart contracts can only really be used to verify receipts which can still be a useful feature.

Conclusion

In conclusion, while the concept of ZKPs and the scenarios in which they can be applied are very convincing, the process of actually incorporating them is not without its hurdles and needs to be well thought out. It’s important to understand how they can be useful and their limitations in order to make an informed decision on whether you should use them in your next project or not. Hopefully, this article was helpful in that regard.


Asparuh Kamenov
About the author:

Asparuh Kamenov


More Stories

Cover Image for Real-World Assets and the Future of DeFi

Real-World Assets and the Future of DeFi

Real-world assets are blockchain-based digital tokens that represent physical and traditional financial assets. The first wave of decentralized finance (DeFi) […]

Read more
Cover Image for Navigating Upgradable Solidity Smart Contracts: A Practical Guide

Navigating Upgradable Solidity Smart Contracts: A Practical Guide

In the ever-evolving landscape of blockchain technology, the concept of upgradeable smart contracts has emerged as a crucial tool for […]

Read more

Have a project in mind?

We have the expertise! Drop us a line and lets talk!