# Write-Up of ZK Hack Mini Puzzle #1
*March 07, 2022*

This [puzzle](https://github.com/kobigurk/zkhack-there-is-something-in-the-AIR) was called "There's something in the AIR", and was based around a STARK Prover and Verifier built with [Winterfell](https://github.com/novifinancial/winterfell) that used a broken AIR. There were two possible solutions that exploited two different sets of missing constraints within the AIR.

The puzzle involves the following:

- **private key:** 4 field elements (```priv_key```)
- **public key:** ```hash(priv_key, [0, 0, 0, 0])```
- **access set:** a merkle tree whose leaves are the public keys in the set. Membership in the access set can be verified by computing the merkle root of the access set from a member's public key and the corresponding merkle path.
- **topic:** a string representing the topic on which members of the access set can vote
- **nullifier:** ```hash(priv_key, hash(topic))``` representing a vote on a topic by a member of the access set. For the Semaphore protocol to be correct, "a combination (private key, topic) should always produce the same nullifier".
- **signal:** ```(STARK proof, nullifier)```
- **Rescue Prime**: the hash function used to compute the public keys, nullifiers, and the merkle tree representing the access set. The specification is [here](https://eprint.iacr.org/2020/1143.pdf).

The goal of the puzzle is to generate a different nullifier on the same topic.

## Our Approach

The nullifier is generated from the private key and the topic, so broadly this gives us 2 approaches to try:

1. Compute the nullifier differently (but use the same inputs).
2. Replace the private key with a different input.

To figure out which of these options will be more promising, let's assess the constraints and see what we have to work with.

(As it turns out, both approaches will lead to solutions!)

## The AIR Constraints

The AIR constraints are made up of boundary constraints (defined in [```air::get_assertions```](https://github.com/kobigurk/zkhack-there-is-something-in-the-AIR/blob/30fcf5312fc0fa78c57ecfed58c0ec7af002453c/src/air/mod.rs#L120)) and transition constraints (defined in [```air::evaluate_transition```](https://github.com/kobigurk/zkhack-there-is-something-in-the-AIR/blob/30fcf5312fc0fa78c57ecfed58c0ec7af002453c/src/air/mod.rs#L147)). However, what we actually care about is identifying everything that's *unconstrained*, so instead of thinking about the types of constraints separately, let's draw them out all together and list all of the constraints explicitly.

We'll think of our trace in 3 chunks:

1. The first application of Rescue Prime hash function (steps 0-7).
2. Repeating applications of the Rescue Prime hash (every subsequent set of 8 steps, starting at step 8). There are 2 possible versions of this, depending on the index bit (which depends on the key index and how many hashes we've accumulated).
- index bit = 0:
- index bit = 1
3. The final row. (This will be some step 7mod8, so it will overlap with the last row of one of the two preceding chunks, but it's useful to look at it separately.)

All of the constraints we have to worry about are pictured above.

## Interlude: a bit about Rescue Prime and Rescue-XLIX

For this puzzle you only need a high level understanding of the Rescue Prime hash function, since it's already been implemented in Winterfell (with some minor modifications). The important points to note are:

- **It uses a sponge construction, like Keccak.** There's some brief general background on sponge functions with pictures [here](https://keccak.team/sponge_duplex.html). Rescue Prime is obtained by applying a permutation called Rescue-XLIX to a state array and then retaining only a set portion of the resulting elements.
- **The state input is an array of 12 elements, the first four of which make up the "capacity" portion and the rest of which are the inputs to be hashed.** (This is an adjustment from the spec, where the capacity portion is the last 4 elements of the state.) Changing any of the elements in the input array will affect the hash output, whether they're in the capacity portion or not.
- **The capacity should be set to the number of elements to be hashed, followed by zeros** (another change from the original spec). For example, when hashing 8 elements, the capacity portion of the state input should be ```[8, 0, 0, 0]```. (Note: I've used integers here, but every value should be provided as a field element.)
- **A Rescue-XLIX permutation consists of 7 rounds.**
- **The output of the Rescue Prime hash is in the middle** of the permuted state array that results from Rescue-XLIX. This portion is defined as ```Rescue::DIGEST_RANGE``` in Winterfell, so the result is found at ```state[Rescue::DIGEST_RANGE]``` or ```state[4..8]```.

----


----

For more details on Rescue Prime, you can review the [specification](https://eprint.iacr.org/2020/1143.pdf). You can find Winterfell's functions for computing Rescue Prime rounds and permutations [here](https://github.com/novifinancial/winterfell/blob/main/crypto/src/hash/rescue/rp64_256/mod.rs). They'll come in handy for our solutions.

## Solution 1: Compute the nullifier differently by changing the capacity in Rescue.

As our picture of the constraints makes clear, the 4 elements of the capacity portion in the input to the Rescue Prime function that generates the nullifier are never constrained. This means we can set any or all of them to any field element(s) we want, which will cause the Rescue Prime hash computation to output different nullifiers.

To take advantage of this for generating a different nullifier, we need to make two changes:
1. **Change the initial state of the trace in the [```SemaphoreProver```](https://github.com/kobigurk/zkhack-there-is-something-in-the-AIR/blob/30fcf5312fc0fa78c57ecfed58c0ec7af002453c/src/prover.rs#L35)** so that a different proof is generated.
2. **Change the way the nullifier is generated in [```AccessSet::make_signal```](https://github.com/kobigurk/zkhack-there-is-something-in-the-AIR/blob/30fcf5312fc0fa78c57ecfed58c0ec7af002453c/src/lib.rs#L127)** so that it matches the new nullifier generated in the proof.

For the change to the ```SemaphoreProver```, we'll initialize the trace with different values in the nullifier section of the trace. Specifically, we'll change the first 4 elements, which is the capacity portion of the input to Rescue Prime. We can use any values we want.

impl SemaphoreProver {
    /// Builds an execution trace for the computation required to generate a signal.
    pub fn build_trace(
        &self,
        priv_key: &PrivKey,
        key_index: usize,
        topic: Digest,
        merkle_path: &[Digest],
    ) -> TraceTable<Felt> {
    
    ... unchanged ...
    
    trace.fill(
        |state| {
            // initialize first state of the computation

            // -- merkle path section of the trace --
            ... unchanged ...

            // -- nullifier section of the trace --
            
            // BEGIN CHANGE FOR HACK
            
            // Change the capacity portion of the state input to 
            // the Rescue Prime hash function
            state[12] = Felt::new(2);
            state[13] = Felt::new(4);
            state[14] = Felt::new(6);
            state[15] = Felt::new(8);
            
            // END CHANGE FOR HACK
            
            // the rest of the nullifier input state is unchanged:
            state[16] = priv_key[0];
            state[17] = priv_key[1];
            state[18] = priv_key[2];
            state[19] = priv_key[3];
            state[20] = topic[0];
            state[21] = topic[1];
            state[22] = topic[2];
            state[23] = topic[3];

            // -- index bits column --
            ... unchanged ...
        }
        ... unchanged ...
    )
    ... unchanged ...
}

To update the nullifier generated during ```make_signal```, we want to change the ```PrivKey::get_nullifier``` function so that it computes a Rescue Prime hash using the same input that our ```SemaphoreProver``` is now using.

To do this, we'll create an array of elements that matches the new nullifier section of the trace initialization in our prover, then apply a Rescue permutation:

/// Creates a nullifier for the provided topic against this private key.
///
/// A nullifier is computed simply as hash(key, topic).
pub fn get_nullifier(&self, topic: Digest) -> Digest {

    // BEGIN CHANGE FOR HACK

    // use the new capacity portion that we're using in the prover
    let mut elements = vec![Felt::new(2), Felt::new(4), Felt::new(6), Felt::new(8)];
    
    // add the elements of our private key
    elements.extend_from_slice(&self.elements());
    
    // add the elements of our topic
    elements.extend_from_slice(topic.as_elements());

    // run the rescue hash function
    let mut hash: [Felt; Rescue::STATE_WIDTH] = elements.try_into().unwrap();
    Rescue::apply_permutation(&mut hash);

    // return the result of the Rescue Prime hash
    Digest::new(hash[Rescue::DIGEST_RANGE].try_into().unwrap())
    
    // END CHANGE FOR HACK
}

With these changes, everything now runs successfully!

Not only have we generated a different nullifier, but we have a method which will let us generate $p^4$ unique nullifiers where $p$ is our field modulus, since we can initialize the Rescue capacity with any permutation of 4 field elements in our field.

## Solution 2: Replace the private key with a different input by pre-accumulating the merkle path verification.

This method is fun, but a bit more involved to implement.

### The Plan

First, recall our constraint diagram and notice the following:

- The only constraint on the "private key" input to the nullifier is that it needs to match the input in the middle part of the section of the trace used for merkle path verification.
- The constraints on the merkle path section of the trace are:
- The result at the end of the computation must equal the root of our access set
- The transitions for steps 0mod8 to 7mod8 must always be valid rounds of a Rescue-XLIX permutation.
- The transition from steps 7mod8 to 0mod8 must either move the digest resulting from the previous Rescue permutation to the right by 4 slots or keep it in the same place, depending on the value of the index bit.

**Importantly, there's no constraint on the number of Rescue permutations we perform, and there's no constraint on the length of our trace.**

As long as we perform valid Rescue permutations and end our hashing with a result matching the root of our access set, there's no reason we need to initialize the trace with our private key. That means we could potentially pass a different input to the hash that generates our nullifier.

**The solution is to perform some of the Merkle path verification *before we start our proof*.** Then we can start our trace with the hash of a partially-accumulated Merkle path and finish verifying the rest of the Merkle path during the proof.

----

----

----

There's one more detail to consider. I said there was no constraint on the length of our trace, but that was a bit of a lie... There's no explicit constraint in the AIR, but the length of the execution trace must be a power of two for the STARK proof.

That restricts us a bit - it means this method will only work for computing either 1 or 2 Rescue permutations (8 steps or 16 steps). Therefore we can only make this method work if we use the prover to compute either the last stage or the last two stages of the Merkle path verification.

**For simplicity, let's hash only the final node of the Merkle path inside our proof**, using the accumulated hash verification to that point as an input, and making an 8 step execution trace.

### The Code

We'll start by rewriting the SemaphoreProver's ```build_trace``` function to include the following changes:
1. Accept two ```Digests``` to be hashed, one of which will be the accumulated Merkle path and the other of which will be the final node of our key's Merkle path. We no longer need the ```priv_key``` or the full ```merkle_path```.
2. Set the trace length to run for just one Rescue Prime hash (8 steps).
3. Initialize the trace with the two digests in the merkle section of the trace and place the first of those digests into the nullifier section of the trace (to pass the equality comparison constraint).
4. The transition function becomes very simple. We only need to execute the rescue rounds, since the trace will end after 1 full cycle.

impl SemaphoreProver {
    /// Builds an execution trace for the computation required to generate a signal.
    pub fn build_trace(&self, topic: Digest, parts_to_hash: &[Digest]) -> TraceTable<Felt> {
        // allocate memory to hold the trace table
        // run just one hash cycle to compute the nullifier and the last leg of the merkle path verification
        let trace_length = HASH_CYCLE_LEN;
        let mut trace = TraceTable::new(TRACE_WIDTH, trace_length);
        let hash_part_1: [Felt; 4] = parts_to_hash[0].into();
        let hash_part_2: [Felt; 4] = parts_to_hash[1].into();

        // let priv_key = priv_key.elements();
        let topic: [Felt; 4] = topic.into();

        trace.fill(
            |state| {
                // initialize first state of the computation

                // -- merkle path section of the trace --
                state[0] = Felt::new(8);
                state[1] = Felt::ZERO;
                state[2] = Felt::ZERO;
                state[3] = Felt::ZERO;
                state[4] = hash_part_1[0];
                state[5] = hash_part_1[1];
                state[6] = hash_part_1[2];
                state[7] = hash_part_1[3];
                state[8] = hash_part_2[0];
                state[9] = hash_part_2[1];
                state[10] = hash_part_2[2];
                state[11] = hash_part_2[3];

                // -- nullifier section of the trace --
                state[12] = Felt::new(8);
                state[13] = Felt::ZERO;
                state[14] = Felt::ZERO;
                state[15] = Felt::ZERO;
                state[16] = hash_part_1[0];
                state[17] = hash_part_1[1];
                state[18] = hash_part_1[2];
                state[19] = hash_part_1[3];
                state[20] = topic[0];
                state[21] = topic[1];
                state[22] = topic[2];
                state[23] = topic[3];

                // -- index bits column --
                state[24] = Felt::ZERO;
            },
            |step, state| {
                // execute the transition function for all steps

                // determine where in the cycle we are
                let cycle_pos = step % HASH_CYCLE_LEN;

                // compute one round of Rescue hash in columns [0..12] for private key
                // hashing and Merkle branch verification
                apply_rescue_round(&mut state[..12], cycle_pos);
                // compute one round of Rescue hash in columns [12..24] for nullifier
                // computation
                apply_rescue_round(&mut state[12..24], cycle_pos);

                // ends after HASH_CYCLE_LEN steps, so nothing else is needed
            },
        );

        // set index bit at the second step to one; this still results in a valid execution trace
        // because actual index bits are inserted into the trace after step 7, but it ensures
        // that there are no repeating patterns in the index bit column, and thus the degree
        // of the index bit constraint is stable.
        trace.set(24, 1, FieldElement::ONE);

        trace
    }
}

Next, we need add a helper function to accumulate the Merkle path. We'll accumulate as many nodes as are provided, so when calling this we'll pass only the portion of the path that we want to pre-accumulate.

/// Accumulate the specified partial merkle path for the index until we run out of path notes, then
/// return the accumulated hash.
/// Note: this code was liberally borrowed from the Winterfell Rescue module with minor edits
/// https://github.com/novifinancial/winterfell/blob/main/crypto/src/hash/rescue/rp64_256/mod.rs
fn hash_merkle_path(index: usize, merkle_path: &[Digest]) -> Digest {
    let r = index & 1;
    let mut acc_hash = Rescue::merge(&[merkle_path[r], merkle_path[1 - r]]);

    let mut index = (index + 2usize.pow((merkle_path.len() - 1) as u32)) >> 1;

    for &path_node in merkle_path.iter().skip(2) {
        acc_hash = if index & 1 == 0 {
            Rescue::merge(&[acc_hash, path_node])
        } else {
            Rescue::merge(&[path_node, acc_hash])
        };
        index >>= 1;
    }

    acc_hash
}

Finally, we'll update the ```make_signal``` function in the ```AccessSet``` to generate our false nullifier and build a valid proof for it.

impl AccessSet {

...UNCHANGED...

    /// Returns a signal of the user with specified private key on the specified topic.
    ///
    /// The signal includes a unique nullifier for the combination of (priv_key, topic), as well
    /// as the proof that the public key for the provided private key exists in this access set.
    pub fn make_signal(&self, priv_key: &PrivKey, topic: &str) -> Signal {
        // compute hash of the topic
        let topic = Rescue::hash(topic.as_bytes());

        // get the index of the key in the Merkle tree
        let pub_key = PubKey::new(priv_key);
        let key_idx = self
            .0
            .leaves()
            .iter()
            .position(|&v| v == pub_key.0)
            .expect("public key for the provided private key could not be found");

        // get the path to the key from the Merkle tree
        let key_path = self
            .0
            .prove(key_idx)
            .expect("failed to build a Merkle path for key index");

        // BEGIN CHANGES FOR HACK
        
        // accumulate the hash of the merkle path, minus the final node
        let last_merkle_idx = key_path.len() - 1;
        let (partial_path, last_leg) = key_path.split_at(last_merkle_idx);
        let acc_hash = hash_merkle_path(key_idx, partial_path);
        let final_index = key_idx >> last_merkle_idx;
        
        // use the index to determine the order of the final hash
        let parts_to_hash: [Digest; 2] = if (final_index & 1) == 0 {
            [acc_hash, last_leg[0]]
        } else {
            [last_leg[0], acc_hash]
        };

        // generate a false nullifier for this topic using the digest section of the merkle path
        // validation that is constrained to equal our nullifier's key
        let cheating_key: Digest = parts_to_hash[0];
        let nullifier = Rescue::merge(&[cheating_key, topic]);

        // build the proof, hashing the final leg of our merkle path verification
        // and hashing our false nullifier from cheating key and the topic
        let prover = SemaphoreProver::default();
        let trace = prover.build_trace(topic, &parts_to_hash);
        
        // END CHANGES FOR HACK
        
        let proof = prover.prove(trace).expect("failed to generate proof");

        // return the signal
        Signal { nullifier, proof }
    }
    
...UNCHANGED...

}

Try running it, and once again everything should pass, verifying our false nullifier!

With this approach, we can get 2 additional nullifiers, by hashing either the final node or the final 2 nodes during our proof generation.

## Fixing the AIR

To fix the AIR, we need to address both attack vectors we've uncovered, which we can do in the following ways.

1. **Constrain the capacity of the nullifier section of the trace initialization to [8, 0, 0, 0].** This will ensure that when Rescue is applied to valid nullifier inputs it always computes the same nullifier.
2. **Constrain the number of Rescue Prime hashes**. This will ensure that we can't pre-accumulate the merkle path verification and replace the private key with a different input.

### Fix 1: Constraining the capacity for the nullifier
This fix is straightforward. In fact, it even was documented in the code, giving a hint to the puzzle for the eagle-eyed. 👀

We'll add 4 new boundary constraints for the capacity initialization in the ```SemaphoreProver::get_assertions``` function.

/// Returns a set of assertions for Semaphore AIR. These assertions are transformed by the
/// prover and verifier into boundary constraints.
///
/// These assertions enforce that:
/// - The trace terminates with Merkle tree root in columns [4, 5, 6, 7].
/// - Columns [16, 16, 18, 19] at step 7 contain value of the nullifier.
/// - Topic was inserted into columns [20, 21, 22, 23] at the first step.
/// - Columns [0, 1, 2, 3] are reset to zeros every 8 steps to (8, 0, 0, 0).
/// - Columns [12, 13, 14, 15] are set to (8, 0, 0, 0) at the first step.
fn get_assertions(&self) -> Vec<Assertion<Felt>> {
    let last_step = self.trace_length() - 1;
    vec![
        Assertion::single(4, last_step, self.tree_root[0]),
        Assertion::single(5, last_step, self.tree_root[1]),
        Assertion::single(6, last_step, self.tree_root[2]),
        Assertion::single(7, last_step, self.tree_root[3]),
        
        // BEGIN CHANGE TO FIX AIR CONSTRAINTS
        
        Assertion::single(12, 0, Felt::new(8)),
        Assertion::single(13, 0, Felt::ZERO),
        Assertion::single(14, 0, Felt::ZERO),
        Assertion::single(15, 0, Felt::ZERO),
        
        // END CHANGE TO FIX AIR CONSTRAINTS
        
        Assertion::single(16, 7, self.nullifier[0]),
        Assertion::single(17, 7, self.nullifier[1]),
        Assertion::single(18, 7, self.nullifier[2]),
        Assertion::single(19, 7, self.nullifier[3]),
        Assertion::single(20, 0, self.topic[0]),
        Assertion::single(21, 0, self.topic[1]),
        Assertion::single(22, 0, self.topic[2]),
        Assertion::single(23, 0, self.topic[3]),
    ]
}

### Fix 2: Constrain the number of Rescue Prime hashes
This fix is a bit more complicated, so I'll just sketch it. We need to force the prover to do enough Rescue Prime hashes to convince us they must have hashed all the way from their private key to the root of the access set.

**Here's a sketch of how to prevent the short Merkle path hack:**

1. **Add an extra register ```s``` to the execution trace** which will hold the current step at all steps.
2. **Provide the path length of the Merkle path as an extra public input** for the verifier. It will be $depth + 1$ where $depth = (log_2(k))$ and $k$ is the number of public keys. The size of the access set is already publicly known, so it's fine to add this input. The path length is $depth + 1$ because we need proof that the signaling member holds a private key, which means we need to start by generating a public key from a private key.
3. **Add a boundary constraint to enforce the value in the final row of the step register** as (path_len * 8 - 1). We multiply by 8 because that's how many steps it takes to compute each hash.
4. We also need to **add a few more constraints for our new step register**:
* Add a boundary constraint that the first value in the step register must be 0.
* Add a transition constraint that the values must always increment by 1 for every step. We can use $s_i + 1 - s_{i+1}= 0$

This should be enough to prevent our hack in the future!