[zk-token-sdk] Add option to create proof context state in the proof verification program

[samkim-crypto]

Problem

Programs such as token-2022 depends on the zkp proof verification by the ZkTokenProof program. A token-2022 instruction such as the confidential transfer requires that the token-2022 instruction be accompanied by a ZkTokenProof program instruction that certifies that the token-2022 instruction data is valid.

This setting can be quite limiting for certain use caes such as when a PDA must make a confidential transfer. A PDA cannot own an ElGamal private key, which is needed to create a zkp and submit a proof instruction.

For example, applications like a private escrow as specified in the SPL issue is not currently possible.

Summary of Changes

Extend the proof program so that users can optionally create a "zkp context state account".

• Re-organize the zero-knowledge proof data into two components: the zk proof and the proof context. A zk proof component are needed to verify the proof, but can be thrown out after it is verified. The proof context is needed by applications such as token-2022 to check that a zk proof instruction was generated w.r.t. the correct pubkey, ciphertext, etc.
• In the ZkTokenProof program, add the option to store the proof context information in an account. This account can only be created if the zero-knowledge proof verifies.
• In the ZkTokenProof program, add a CloseContextState instruction that closes the context state account by an authority.

Token-22 Invocation by PDAs

The token-22 program should be updated so that for instructions such as a confidential transfer is accepted by the processor if one of the following conditions are met:

1. a ZkTokenProof::Transfer instruction is attached in the same transaction (as before)
2. a ZkTokenProof TransferProofContext account is specified in the set of instruction addresses

These changes allow a ZkTokenProof instruction to be processed in a separate transaction before a token-2022 instruction. This enables applications that require PDAs to invoke the token-2022 program by supplying a proof context state account.

For example, consider the private escrow as specified in the SPL issue.

• A private escrow instigator (Alice) can create a token-2022 confidential extension account and set its owner to an escrow PDA.
• Use a ZkTokenProof::Transfer instruction to create a proof context state intended for the transfer of the tokens in the confidential extension account to the swap recipient (Bob) beforehand.
• When Bob satisfies the requirements by the escrow program, the escrow PDA can invoke the token-22 program to make the transfer of Alice's tokens to Bob. Instead of creating the ZkTokenProof instruction itself, it will specify the proof context account that Alice previously created for the PDA.

Processing Large Zkps by Chunks

The option to create proof context states also enables large zkps that can't fit into a single transaction to be processed in multiple instructions. For example, a confidential transfer proof is a combination of range proof, equality proof, and validity proof. Currently, these individual proofs are not exposed in the ZkTokenProof program, but if the program is updated to do so, then users can submit and create context accounts for smaller ZkTokenProof::RangeProof, ZkTokenProof::EqualityProof, and ZkTokenProof::ValidityProof instructions. Then the token-22 could check the existence of these context state accounts and then approve the transfer.

This may not be the fastest or the most UI friendly, but it could be one way to use the token-22 instructions before the transaction size limit is increased.

[samkim-crypto]

These proof program changes break downstream projects in SPL token. If the changes look reasonable, then I can temporarily disable proof verification in SPL and figure out a way to get the ci tests to accept.

[joncinque]

The guts are looking great! Just some bits to clarify / test a bit more

[samkim-crypto]

The changes should be ready for another review.

One issue that came up is that the CloseContextState instruction does not really depend on the generic type T: ProofContext. Currently, the instruction decodes ProofContextState<T> according to the proof type in order to get access to context_state_authority (to check that the authority signed).

pub struct ProofContextState<T: ZkProofContext> {
    /// The zero-knowledge proof type
    pub proof_type: PodProofType,
    /// The proof context authority that can close the account
    pub context_state_authority: Pubkey,
    /// The proof context data
    pub proof_context: T,
}

However, context_state_authority does exists for all ProofContextState independent of T. I wonder if it is possible/better to not include the proof type in the CloseContextState account. Then, instead of decoding the account data to ProofContextState, just assume that account data at certain bytes positions (i.e. [1..PUBKEY_SIZE+1]) encodes the context_state_authority pubkey. This will be less elegant, but this could allow a set of instructions that are encodable as a single byte as before:

pub enum ProofInstruction {
   CloseContextState,
   VerifyCloseAccount,
   VerifyWithdraw,
   VerifyWithdrawWithheldTokens,
   ...
   VerifyPubkeyValidity,
}

[joncinque]

The changes look great overall! Regarding:

I wonder if it is possible/better to not include the proof type in the CloseContextState account. Then, instead of decoding the account data to ProofContextState, just assume that account data at certain bytes positions (i.e. [1..PUBKEY_SIZE+1]) encodes the context_state_authority pubkey

This would be pretty slick. How about just decoding as a different type? Ie:

struct ProoflessContextState {
    /// The zero-knowledge proof type
    pub proof_type: PodProofType,
    /// The proof context authority that can close the account
    pub context_state_authority: Pubkey,
}

Then you can decode the struct from the data, make sure there's a valid proof type, and assume that context_state_authority is correct. To be honest, I don't feel too strongly about it, but this would be closer to how other programs handle "close", which has just one "close" instruction to close any account managed by the program.

Maybe someone else will have a stronger opinion than me. I'll let you / them decide.

Side note: feel free to go with a better name! ProoflessContextState was the first thing that came to mind

[samkim-crypto]

Great! I was able to remove the proof type from CloseContextState. This simplified the set of instructions and also shaved off 1 byte from the instruction encodings.

I ended up creating CreateContextStateMeta struct like you suggested. I think this is really the cleanest way to organize functions that are generic-independent.

[joncinque]

Looking really close! Just one last question

[joncinque]

I may be missing something obvious, but what's the reason for having this bool and also the proof type? you could probably just move the proof type to here, and then make sure that it's not uninitialized during closing.

[samkim-crypto]

Yes, that should definitely work and it will save a byte. I just thought about what is the most natural in terms of the literal interpretation of a "proof context". ProofType is context information for a proof, so it is natural inside proof_context: T. On the other hand, is_initialized is state information for the state account.

But thinking about it now, maybe it is worth using proof_type for is_initialized since that gives the most efficiency...

[codecov]

Codecov Report

Merging #29996 (275763c) into master (65cd552) will decrease coverage by 0.3%.
The diff coverage is 28.6%.

@@            Coverage Diff            @@
##           master   #29996     +/-   ##
=========================================
- Coverage    81.7%    81.4%   -0.3%     
=========================================
  Files         723      724      +1     
  Lines      201802   202471    +669     
=========================================
- Hits       164891   164854     -37     
- Misses      36911    37617    +706     

[samkim-crypto]

The downstream CI tests finally passes. I apologize for the long delay in this PR. If anyone can have a quick sanity check to approve this, then it would be great 🙏

[joncinque]

Just one more question, then this should be good to go! Thanks for your patience

[joncinque]

Now that I'm reading back through this, do we need ZkProofContext at all? I think we can generalize and simplify this a bit if we have:

pub trait ZkProofData<T: Pod> {
    const PROOF_TYPE: ProofType;
    fn context_data(&self) -> &T;
    #[cfg(not(target_os = "solana"))]
    fn verify_proof(&self) -> Result<(), ProofError>;
}

Alternatively, you can remove the proof_type parameter from ProofContextState::encode since it can be fetched directly from T, but I would lean towards the first option to keep things simpler.

[samkim-crypto]

Yes this is a good point. I was initially a little torn with this because there should technically be one ZkProofContext type for each ZkProofData type. Even if ZkProofContext is an empty market trait, having it as an associated type in ZkProofData naturally enforces this. For example, implementing CloseProofContext inside TransferProofData would not really make sense.

But in the end, I decided to go with removing ZkProofContext because it does simplify the instruction implementations a bit. Thanks for the suggestion!

[joncinque]

Looks great! Thanks for all of your patience during that long back and forth