From ba92ba0e061edbca3fb37b057d69687875108d8b Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Tue, 5 Apr 2022 09:40:45 -0300
Subject: [PATCH] Zk instructions check length (#24103)

* zk-token-sdk: add a length check before decoding proof instruction

* zk-token-sdk: fix minor spelling

* zk-token-sdk: one-liner for length check

* zk-token-sdk: one-liner fix
---
 programs/zk-token-proof/src/lib.rs             | 2 +-
 zk-token-sdk/src/zk_token_proof_instruction.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/programs/zk-token-proof/src/lib.rs b/programs/zk-token-proof/src/lib.rs
index 7578b03266..d7c791c76c 100644
--- a/programs/zk-token-proof/src/lib.rs
+++ b/programs/zk-token-proof/src/lib.rs
@@ -56,7 +56,7 @@ pub fn process_instruction(
             verify::<WithdrawData>(invoke_context)
         }
         ProofInstruction::VerifyWithdrawWithheldTokens => {
-            ic_msg!(invoke_context, "VerifyWithdraw");
+            ic_msg!(invoke_context, "VerifyWithdrawWithheldTokens");
             verify::<WithdrawData>(invoke_context)
         }
         ProofInstruction::VerifyTransfer => {
diff --git a/zk-token-sdk/src/zk_token_proof_instruction.rs b/zk-token-sdk/src/zk_token_proof_instruction.rs
index 3dfdea18de..5bcd653967 100644
--- a/zk-token-sdk/src/zk_token_proof_instruction.rs
+++ b/zk-token-sdk/src/zk_token_proof_instruction.rs
@@ -73,7 +73,7 @@ impl ProofInstruction {
     }
 
     pub fn decode_type(input: &[u8]) -> Option<Self> {
-        FromPrimitive::from_u8(input[0])
+        input.get(0).and_then(|x| FromPrimitive::from_u8(*x))
     }
 
     pub fn decode_data<T: Pod>(input: &[u8]) -> Option<&T> {