From 4527ec08dc05a2f7e9387d5b5253a2532db39285 Mon Sep 17 00:00:00 2001
From: rasmus-kirk <mail@rasmuskirk.com>
Date: Thu, 17 Oct 2024 17:30:48 +0200
Subject: [PATCH] Added init_origin functions

---
 concordium-std/CHANGELOG.md    |  2 ++
 concordium-std/src/prims.rs    |  6 ++++-
 concordium-std/src/test_env.rs | 46 +++++++++++++++++++++++-----------
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/concordium-std/CHANGELOG.md b/concordium-std/CHANGELOG.md
index 5e6df9dd..be0e7e8a 100644
--- a/concordium-std/CHANGELOG.md
+++ b/concordium-std/CHANGELOG.md
@@ -7,6 +7,7 @@
   - `set_receive_self_address`
   - `set_receive_self_balance`
   - `set_parameter`
+  - `set_init_origin`
   - `get_event`
   - `get_event_size`
 
@@ -17,6 +18,7 @@
   - `set_receive_balance`
   - `set_receive_self_address`
   - `set_parameter`
+  - `set_init_origin`
   - `get_event`
 
 ## concordium-std 10.1.0 (2024-04-04)
diff --git a/concordium-std/src/prims.rs b/concordium-std/src/prims.rs
index 776976bd..efb88c9f 100644
--- a/concordium-std/src/prims.rs
+++ b/concordium-std/src/prims.rs
@@ -305,10 +305,14 @@ extern "C" {
     #[cfg(all(feature = "wasm-test", target_arch = "wasm32"))]
     pub(crate) fn set_parameter(i: u32, start: *const u8, length: u32);
 
+    /// Sets the address of the sender, 32 bytes
+    #[cfg(all(feature = "wasm-test", target_arch = "wasm32"))]
+    pub(crate) fn set_init_origin(start: *const u8);
+
     /// Gets event number `i` in the smart contract state. Returns `-1` if `i`
     /// is an invalid index. Otherwise returns bytes written.
     #[cfg(all(feature = "wasm-test", target_arch = "wasm32"))]
-    pub(crate) fn get_event(i: u32, start: *const u8) -> i32;
+    pub(crate) fn get_event(i: u32, start: *mut u8) -> i32;
 
     /// Gets the size of event number `i` in the smart contract state. Returns `-1` if `i`
     /// is an invalid index.
diff --git a/concordium-std/src/test_env.rs b/concordium-std/src/test_env.rs
index 76a305d4..8954ea4b 100644
--- a/concordium-std/src/test_env.rs
+++ b/concordium-std/src/test_env.rs
@@ -1,4 +1,6 @@
-use concordium_contracts_common::{Amount, ContractAddress, SlotTime};
+use concordium_contracts_common::{
+    AccountAddress, Amount, ContractAddress, Serial, SlotTime,
+};
 
 use crate::{prims, to_bytes};
 
@@ -49,6 +51,11 @@ impl TestEnv {
             Some(buf)
         }
     }
+
+    pub fn set_init_origin(self, address: AccountAddress) -> Self {
+        unsafe { prims::set_init_origin(address.0.as_ptr()) };
+        self
+    }
 }
 
 #[cfg(feature = "internal-wasm-test")]
@@ -70,19 +77,12 @@ mod wasm_test {
         }
     }
 
-    fn extern_chain_meta() -> ExternChainMeta { ExternChainMeta {} }
-
-    fn external_call_response() -> ExternCallResponse {
-        ExternCallResponse::new(unsafe { NonZeroU32::new_unchecked(1) })
-    }
-
-    fn receive_context() -> ReceiveContext { ExternContext::default() }
-
     #[concordium_test]
     fn set_get_slot_time() {
+        let extern_chain_meta = ExternChainMeta {};
         let original = Timestamp::from_timestamp_millis(10);
         TestEnv.set_slot_time(original);
-        let stored = extern_chain_meta().block_time();
+        let stored = extern_chain_meta.block_time();
         claim_eq!(original, stored)
     }
 
@@ -96,14 +96,17 @@ mod wasm_test {
 
     #[concordium_test]
     fn set_get_receive_self_address() {
+        let receive_context: ReceiveContext = ExternContext::default();
         let original = ContractAddress::new(5040, 12);
         TestEnv.set_receive_self_address(original);
-        let stored = receive_context().self_address();
+        let stored = receive_context.self_address();
         claim_eq!(original, stored);
     }
 
     #[concordium_test]
     fn set_get_parameter() {
+        let mut external_call_response =
+            ExternCallResponse::new(unsafe { NonZeroU32::new_unchecked(1) });
         let i = 1;
         let param = vec![3u8; 7];
         let mut buf = vec![0u8; 10];
@@ -122,8 +125,7 @@ mod wasm_test {
         // Use the external call to test the actual function that would be called in
         // most real scenarios
         buf = vec![0u8; 7];
-        let bytes_written =
-            external_call_response().read(&mut buf).expect("Unable to read bytes to buffer");
+        let bytes_written = external_call_response.read(&mut buf).unwrap();
         let expected = param;
 
         claim_eq!(buf, expected);
@@ -136,7 +138,7 @@ mod wasm_test {
         let mut buf = vec![0u8; 10];
 
         let param_size = unsafe { prims::get_parameter_size(i) };
-        let bytes_written = unsafe { prims::get_parameter_section(i, buf.as_mut_ptr(), 5, 2) };
+        let bytes_written = unsafe { prims::get_parameter_section(i, buf.as_mut_ptr(), 0, 0) };
 
         claim_eq!(param_size, -1);
         claim_eq!(param_size, bytes_written);
@@ -185,4 +187,20 @@ mod wasm_test {
         let stored = TestEnv.get_event(0);
         claim_eq!(stored, None);
     }
+
+    #[concordium_test]
+    fn set_get_init_origin() {
+        let init_context: InitContext = ExternContext::default();
+        // 3kBx2h5Y2veb4hZgAJWPrr8RyQESKm5TjzF3ti1QQ4VSYLwK1G
+        let address = AccountAddress([
+            105, 117, 36, 6, 204, 147, 159, 201, 12, 166, 167, 59, 87, 206, 225, 9, 150, 53, 71,
+            249, 66, 0, 109, 33, 145, 68, 146, 79, 132, 133, 251, 13,
+        ]);
+        TestEnv.set_init_origin(address);
+
+        let stored = init_context.init_origin();
+        let expected = address;
+
+        claim_eq!(stored, expected);
+    }
 }