From 8b15d7bc283b0623fecd4f332c3c58821252ebd0 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Tue, 21 Sep 2021 19:58:14 -0700
Subject: [PATCH] Add CloseAccountData proof program plumbing

---
 programs/zk-token-proof/src/lib.rs | 32 +++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/programs/zk-token-proof/src/lib.rs b/programs/zk-token-proof/src/lib.rs
index d1d187fbde..2d73f3fcce 100644
--- a/programs/zk-token-proof/src/lib.rs
+++ b/programs/zk-token-proof/src/lib.rs
@@ -1,6 +1,7 @@
 #![forbid(unsafe_code)]
 
 use {
+    bytemuck::Pod,
     solana_sdk::{
         ic_msg, instruction::InstructionError, process_instruction::InvokeContext, pubkey::Pubkey,
     },
@@ -8,6 +9,21 @@ use {
     std::result::Result,
 };
 
+fn verify<T: Pod + Verifiable>(
+    input: &[u8],
+    invoke_context: &mut dyn InvokeContext,
+) -> Result<(), InstructionError> {
+    let proof = ProofInstruction::decode_data::<T>(input).ok_or_else(|| {
+        ic_msg!(invoke_context, "invalid proof data");
+        InstructionError::InvalidInstructionData
+    })?;
+
+    proof.verify().map_err(|err| {
+        ic_msg!(invoke_context, "proof verification failed: {:?}", err);
+        InstructionError::InvalidInstructionData
+    })
+}
+
 pub fn process_instruction(
     program_id: &Pubkey,
     input: &[u8],
@@ -16,15 +32,13 @@ pub fn process_instruction(
     match ProofInstruction::decode_type(program_id, input)
         .ok_or(InstructionError::InvalidInstructionData)?
     {
-        ProofInstruction::VerifyUpdateAccountPkData => {
-            ic_msg!(invoke_context, "VerifyUpdateAccountPkData");
-            let proof = ProofInstruction::decode_data::<UpdateAccountPkData>(input)
-                .ok_or(InstructionError::InvalidInstructionData)?;
-
-            proof.verify().map_err(|err| {
-                ic_msg!(invoke_context, "proof verification failed: {:?}", err);
-                InstructionError::InvalidInstructionData
-            })
+        ProofInstruction::VerifyUpdateAccountPk => {
+            ic_msg!(invoke_context, "VerifyUpdateAccountPk");
+            verify::<UpdateAccountPkData>(input, invoke_context)
+        }
+        ProofInstruction::VerifyCloseAccount => {
+            ic_msg!(invoke_context, "VerifyCloseAccount");
+            verify::<CloseAccountData>(input, invoke_context)
         }
     }
 }