From c564dae8dbab43d9ef7140a2933f8d197069e5f7 Mon Sep 17 00:00:00 2001
From: Moody Salem <moody.salem@gmail.com>
Date: Wed, 19 Jun 2024 09:50:09 -0400
Subject: [PATCH] fix the critical issue in governor account contract
 implementation

---
 src/governor.cairo      |  3 ++-
 src/governor_test.cairo | 27 ++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/governor.cairo b/src/governor.cairo
index 90f8c59..bece142 100644
--- a/src/governor.cairo
+++ b/src/governor.cairo
@@ -90,7 +90,7 @@ pub trait IGovernor<TContractState> {
     fn upgrade(ref self: TContractState, class_hash: ClassHash);
 }
 
-#[starknet::contract]
+#[starknet::contract(account)]
 pub mod Governor {
     use core::hash::{HashStateTrait, HashStateExTrait};
     use core::num::traits::zero::{Zero};
@@ -461,6 +461,7 @@ pub mod Governor {
             0
         }
         fn __execute__(ref self: ContractState, mut calls: Array<Call>) -> Array<Span<felt252>> {
+            assert(get_caller_address().is_zero(), 'Invalid caller');
             let mut results: Array<Span<felt252>> = array![];
             while let Option::Some(call) = calls.pop_front() {
                 results.append(call.execute());
diff --git a/src/governor_test.cairo b/src/governor_test.cairo
index 6063e81..a063e49 100644
--- a/src/governor_test.cairo
+++ b/src/governor_test.cairo
@@ -11,7 +11,8 @@ use starknet::account::{Call};
 use starknet::{
     get_contract_address, syscalls::deploy_syscall, ClassHash, contract_address_const,
     ContractAddress, get_block_timestamp,
-    testing::{set_block_timestamp, set_contract_address, pop_log}
+    testing::{set_block_timestamp, set_contract_address, pop_log},
+    account::{AccountContractDispatcher, AccountContractDispatcherTrait}
 };
 
 fn recipient() -> ContractAddress {
@@ -1014,6 +1015,30 @@ fn test_reconfigure_fails_if_not_self_call() {
         );
 }
 
+#[test]
+#[should_panic(expected: ("Not allowed", 'ENTRYPOINT_FAILED'))]
+fn test_governor_validate_fails() {
+    let (_staker, _token, governor, _config) = setup();
+    AccountContractDispatcher { contract_address: governor.contract_address }
+        .__validate__(array![]);
+}
+
+#[test]
+#[should_panic(expected: ("Not allowed", 'ENTRYPOINT_FAILED'))]
+fn test_governor_validate_declare_fails() {
+    let (_staker, _token, governor, _config) = setup();
+    AccountContractDispatcher { contract_address: governor.contract_address }
+        .__validate_declare__(123);
+}
+
+#[test]
+#[should_panic(expected: ('Invalid caller', 'ENTRYPOINT_FAILED'))]
+fn test_governor_execute_fails_from_non_zero() {
+    let (_staker, _token, governor, _config) = setup();
+    set_contract_address(contract_address_const::<1>());
+    AccountContractDispatcher { contract_address: governor.contract_address }.__execute__(array![]);
+}
+
 #[test]
 fn test_reconfigure_succeeds_self_call() {
     let (staker, token, governor, config) = setup();