diff --git a/example/smart-contract/src/lib.rs b/example/smart-contract/src/lib.rs
index 3b7a6adb4..53c5ba425 100644
--- a/example/smart-contract/src/lib.rs
+++ b/example/smart-contract/src/lib.rs
@@ -1,10 +1,11 @@
 use darkfi_sdk::{
     crypto::Nullifier,
-    initialize, entrypoint, update_state,
+    entrypoint,
     error::{ContractError, ContractResult},
-    msg,
+    initialize, msg,
     pasta::pallas,
-    state::{set_update, nullifier_exists},
+    state::{nullifier_exists, set_update},
+    update_state,
 };
 use darkfi_serial::{deserialize, serialize, SerialDecodable, SerialEncodable};
 
@@ -41,7 +42,7 @@ pub struct BarArgs {
 #[derive(SerialEncodable, SerialDecodable)]
 pub struct FooUpdate {
     pub name: String,
-    pub y: u32
+    pub y: u32,
 }
 
 initialize!(init_contract);
@@ -68,10 +69,7 @@ fn process_instruction(ix: &[u8]) -> ContractResult {
             // ...
             let args: FooArgs = deserialize(tx_data)?;
             // ...
-            let update = FooUpdate {
-                name: "john_doe".to_string(),
-                y: 110
-            };
+            let update = FooUpdate { name: "john_doe".to_string(), y: 110 };
 
             let mut update_data = vec![Function::Foo as u8];
             update_data.extend_from_slice(&serialize(&update));
@@ -115,20 +113,16 @@ fn process_instruction(ix: &[u8]) -> ContractResult {
 }
 
 update_state!(process_update);
-fn process_update() -> ContractResult {
+fn process_update(update_data: &[u8]) -> ContractResult {
     msg!("Make update!");
 
-    /*
-    let (func_id, update_data) = get_update()?;
-
-    match Function::from(func_id) {
+    match Function::from(update_data[0]) {
         Function::Foo => {
-            let update: FooUpdate = deserialize(update_data)?;
+            let update: FooUpdate = deserialize(&update_data[1..])?;
             // update.apply()
         }
-        _ => unreachable!()
-    };
-    */
+        _ => unreachable!(),
+    }
 
     Ok(())
 }
@@ -140,4 +134,3 @@ fn process_update() -> ContractResult {
 //fn apply(update) {
 //    // writes happen here
 //}
-
diff --git a/example/smart-contract/tests/runtime.rs b/example/smart-contract/tests/runtime.rs
index 7b7b32d58..0e15c9cdf 100644
--- a/example/smart-contract/tests/runtime.rs
+++ b/example/smart-contract/tests/runtime.rs
@@ -65,9 +65,9 @@ fn run_contract() -> Result<()> {
     // ============================================================
     // Serialize the payload into the runtime format and execute it
     // ============================================================
-    //let update = runtime.exec(&serialize_payload(&payload))?;
+    runtime.exec(&serialize_payload(&payload))?;
 
-    //runtime.apply(update);
+    runtime.apply()?;
     //Ok(())
 
     //runtime.exec(&serialize_payload(&payload))?;
@@ -76,4 +76,3 @@ fn run_contract() -> Result<()> {
 
     Ok(())
 }
-
diff --git a/src/runtime/chain_state.rs b/src/runtime/chain_state.rs
index 0baa30027..3c1c7c561 100644
--- a/src/runtime/chain_state.rs
+++ b/src/runtime/chain_state.rs
@@ -26,7 +26,7 @@ use super::{
 };
 use crate::node::state::ProgramState;
 
-pub(crate) fn set_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
+pub(super) fn set_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
     let env = ctx.data();
     match env.contract_section {
         ContractSection::Exec => {
@@ -35,7 +35,7 @@ pub(crate) fn set_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u3
             // FIXME: make me preettty!
             let slice = ptr.slice(&memory_view, len);
             if slice.is_err() {
-                return -2;
+                return -2
             }
             let slice = slice.unwrap();
 
@@ -53,36 +53,21 @@ pub(crate) fn set_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u3
             };
             //
 
+            // FIXME: Shouldn't assert here, but rather return an error.
+            // An assert would make the host panic.
             assert!(env.contract_update.take().is_none());
             let func_id = update_data[0];
             let update_data = &update_data[1..];
             env.contract_update.set(Some((func_id, update_data.to_vec())));
             0
         }
-        _ => {
-            -1
-        }
-    }
-}
-
-pub(crate) fn get_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
-    let env = ctx.data();
-    match env.contract_section {
-        ContractSection::Update => {
-           let memory_view = env.memory_view(&ctx);
-
-
-
-           0
-
-        }
-        _ => { -1 }
+        _ => -1,
     }
 }
 
 /// Try to read a `Nullifier` from the given pointer and check if it's
 /// an existing nullifier in the blockchain state machine.
-pub fn nullifier_exists(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32 {
+pub(super) fn nullifier_exists(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32 {
     let env = ctx.data();
     match env.contract_section {
         ContractSection::Null => {
@@ -127,7 +112,7 @@ pub fn nullifier_exists(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32
 
 /// Try to read a `MerkleNode` from the given pointer and check if it's
 /// a valid Merkle root in the chain's Merkle tree.
-pub fn is_valid_merkle(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32 {
+pub(super) fn is_valid_merkle(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32 {
     /*
     if let Some(bytes) = env.memory.get_ref().unwrap().read(ptr, len as usize) {
         debug!(target: "wasm_runtime::is_valid_merkle", "Read bytes: {:?}", bytes);
diff --git a/src/runtime/vm_runtime.rs b/src/runtime/vm_runtime.rs
index f66547180..68f22205f 100644
--- a/src/runtime/vm_runtime.rs
+++ b/src/runtime/vm_runtime.rs
@@ -18,14 +18,14 @@
 
 use std::{
     cell::{Cell, RefCell},
-    sync::{Arc, Mutex},
+    sync::Arc,
 };
 
 use darkfi_sdk::entrypoint;
 use log::{debug, info};
 use wasmer::{
     imports, wasmparser::Operator, AsStoreRef, CompilerConfig, Function, FunctionEnv, Instance,
-    Memory, MemoryView, Module, Pages, Store, Value, WasmPtr, WASM_PAGE_SIZE,
+    Memory, MemoryView, Module, Pages, Store, Value, WASM_PAGE_SIZE,
 };
 use wasmer_compiler_singlepass::Singlepass;
 use wasmer_middlewares::{
@@ -34,9 +34,9 @@ use wasmer_middlewares::{
 };
 
 use super::{
-    chain_state::{set_update, is_valid_merkle, nullifier_exists},
+    chain_state::{is_valid_merkle, nullifier_exists, set_update},
     memory::MemoryManipulation,
-    util::drk_log,
+    util::{drk_log, serialize_payload},
 };
 use crate::{crypto::contract_id::ContractId, Error, Result};
 
@@ -260,10 +260,23 @@ impl Runtime {
         env_mut.contract_section = ContractSection::Update;
 
         debug!(target: "wasm_runtime::run", "Getting initialize function");
-        let entrypoint = self.instance.exports.get_function(INITIALIZE)?;
+        let entrypoint = self.instance.exports.get_function(UPDATE)?;
+
+        let update_data = env_mut.contract_update.take().unwrap();
+        // FIXME: Less realloc
+        let mut payload = vec![update_data.0];
+        payload.extend_from_slice(&update_data.1);
+        let payload = serialize_payload(&payload);
+
+        // TODO: Test if this works when state update is larger than the initial payload
+        // The question is if we need to allocate more memory or if it's ok to just
+        // overwrite from zero (and even if overwrite - is there enough space?)
+        let pages_required = payload.len() / WASM_PAGE_SIZE + 1;
+        //self.set_memory_page_size(pages_required as u32)?;
+        self.copy_to_memory(&payload)?;
 
         debug!(target: "wasm_runtime::run", "Executing wasm");
-        let ret = match entrypoint.call(&mut self.store, &[]) {
+        let ret = match entrypoint.call(&mut self.store, &[Value::I32(0 as i32)]) {
             Ok(retvals) => {
                 self.print_logs();
                 debug!(target: "wasm_runtime::run", "{}", self.gas_info());
@@ -331,6 +344,7 @@ impl Runtime {
 
     /// Copy payload to the start of the memory
     fn copy_to_memory(&self, payload: &[u8]) -> Result<()> {
+        // TODO: Maybe should write to first zero memory and return the pointer/offset?
         // Get the memory view
         let env = self.ctx.as_ref(&self.store);
         let memory_view = env.memory_view(&self.store);
diff --git a/src/sdk/src/entrypoint.rs b/src/sdk/src/entrypoint.rs
index c7ac85159..e0c56ba38 100644
--- a/src/sdk/src/entrypoint.rs
+++ b/src/sdk/src/entrypoint.rs
@@ -61,8 +61,10 @@ macro_rules! update_state {
     ($process_update:ident) => {
         /// # Safety
         #[no_mangle]
-        pub unsafe extern "C" fn __update() -> u64 {
-            match $process_update() {
+        pub unsafe extern "C" fn __update(input: *mut u8) -> u64 {
+            let update_data = $crate::entrypoint::deserialize(input);
+
+            match $process_update(&update_data) {
                 Ok(()) => $crate::entrypoint::SUCCESS,
                 Err(e) => e.into(),
             }
diff --git a/src/sdk/src/state.rs b/src/sdk/src/state.rs
index 098a86e78..e389cf838 100644
--- a/src/sdk/src/state.rs
+++ b/src/sdk/src/state.rs
@@ -35,18 +35,6 @@ pub fn set_update(update_data: &[u8]) -> Result<(), ContractError> {
     unimplemented!();
 }
 
-pub fn get_update() -> Result<(u8, Vec<u8>), ContractError> {
-    #[cfg(target_arch = "wasm32")]
-    // get_update_ needs to take a buffer?
-    // get pointer for contract_update
-    // piece back into (u8, Vec<u8>)
-    // return
-    return Ok((0, vec![]));
-
-    #[cfg(not(target_arch = "wasm32"))]
-    unimplemented!();
-}
-
 pub fn nullifier_exists(nullifier: &Nullifier) -> Result<bool, ContractError> {
     #[cfg(target_arch = "wasm32")]
     unsafe {