From 6833f41424bf8eb0d50235ea055a90a04f6fedf5 Mon Sep 17 00:00:00 2001
From: Jayant Krishnamurthy <jayantkrishnamurthy@gmail.com>
Date: Thu, 29 Dec 2022 12:04:50 -0800
Subject: [PATCH] [cosmos] Require governance instruction bytes to be
 completely consumed (#448)

* add test

* fix  build

Co-authored-by: Jayant Krishnamurthy <jkrishnamurthy@jumptrading.com>
---
 cosmwasm/contracts/pyth/src/governance.rs | 44 +++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/cosmwasm/contracts/pyth/src/governance.rs b/cosmwasm/contracts/pyth/src/governance.rs
index d9e436af..41e39311 100644
--- a/cosmwasm/contracts/pyth/src/governance.rs
+++ b/cosmwasm/contracts/pyth/src/governance.rs
@@ -137,6 +137,16 @@ impl GovernanceInstruction {
             _ => Err(format!("Unknown governance action type: {action_type}",)),
         };
 
+        // Check that we're at the end of the buffer (to ensure that this contract knows how to
+        // interpret every field in the governance message). The logic is a little janky
+        // but seems to be the simplest way to check that the reader is at EOF.
+        let mut next_byte = [0_u8; 1];
+        let read_result = bytes.read(&mut next_byte);
+        match read_result {
+            Ok(0) => (),
+            _ => Err("Governance action had an unexpectedly long payload.".to_string())?,
+        }
+
         Ok(GovernanceInstruction {
             module,
             action: action?,
@@ -205,3 +215,37 @@ impl GovernanceInstruction {
         Ok(buf)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::governance::{
+        GovernanceAction,
+        GovernanceInstruction,
+        GovernanceModule,
+    };
+
+    #[test]
+    fn test_payload_wrong_size() {
+        let instruction = GovernanceInstruction {
+            module:          GovernanceModule::Target,
+            action:          GovernanceAction::SetFee {
+                val:  100,
+                expo: 200,
+            },
+            target_chain_id: 7,
+        };
+
+        let mut buf: Vec<u8> = instruction.serialize().unwrap();
+
+        let result = GovernanceInstruction::deserialize(buf.as_slice());
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), instruction);
+
+        buf.push(0);
+        let result = GovernanceInstruction::deserialize(buf.as_slice());
+        assert!(result.is_err());
+
+        let result = GovernanceInstruction::deserialize(&buf[0..buf.len() - 2]);
+        assert!(result.is_err());
+    }
+}