diff --git a/arch/sandbox/cpu/state.c b/arch/sandbox/cpu/state.c
index ef2e63f..a347cec 100644
--- a/arch/sandbox/cpu/state.c
+++ b/arch/sandbox/cpu/state.c
@@ -7,6 +7,7 @@
 #include <errno.h>
 #include <fdtdec.h>
 #include <os.h>
+#include <asm/malloc.h>
 #include <asm/state.h>
 
 /* Main state record for the sandbox */
@@ -27,17 +28,17 @@
 		return 0;
 
 	size = used + extra_size;
-	buf = os_malloc(size);
+	buf = malloc(size);
 	if (!buf)
 		return -ENOMEM;
 
 	ret = fdt_open_into(blob, buf, size);
 	if (ret) {
-		os_free(buf);
+		free(buf);
 		return -EIO;
 	}
 
-	os_free(blob);
+	free(blob);
 	state->state_fdt = buf;
 	return 0;
 }
@@ -53,7 +54,7 @@
 		printf("Cannot find sandbox state file '%s'\n", fname);
 		return -ENOENT;
 	}
-	state->state_fdt = os_malloc(size);
+	state->state_fdt = malloc(size);
 	if (!state->state_fdt) {
 		puts("No memory to read sandbox state\n");
 		return -ENOMEM;
@@ -75,7 +76,7 @@
 err_read:
 	os_close(fd);
 err_open:
-	os_free(state->state_fdt);
+	free(state->state_fdt);
 	state->state_fdt = NULL;
 
 	return ret;
@@ -242,7 +243,7 @@
 	/* Create a state FDT if we don't have one */
 	if (!state->state_fdt) {
 		size = 0x4000;
-		state->state_fdt = os_malloc(size);
+		state->state_fdt = malloc(size);
 		if (!state->state_fdt) {
 			puts("No memory to create FDT\n");
 			return -ENOMEM;
@@ -300,7 +301,7 @@
 err_write:
 	os_close(fd);
 err_create:
-	os_free(state->state_fdt);
+	free(state->state_fdt);
 
 	return ret;
 }
@@ -418,7 +419,7 @@
 		os_unlink(state->jumped_fname);
 
 	if (state->state_fdt)
-		os_free(state->state_fdt);
+		free(state->state_fdt);
 	memset(state, '\0', sizeof(*state));
 
 	return 0;