From 5e051888914e8250a6a256135b7a90e227be6b06 Mon Sep 17 00:00:00 2001
From: Gui Andrade <gh@archshift.com>
Date: Thu, 8 Apr 2021 15:57:39 -0700
Subject: [PATCH] Update multimem SBI calls to get output in arg0

---
 paging.c | 14 ++++++++------
 sbi.c    | 12 ++++++------
 sbi.h    |  5 +++--
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/paging.c b/paging.c
index 0645bef..8dcc4c8 100644
--- a/paging.c
+++ b/paging.c
@@ -36,22 +36,24 @@ void paging_dec_user_page(void)
 
 void init_paging(uintptr_t user_pa_start, uintptr_t user_pa_end)
 {
-  uintptr_t addr;
-  uintptr_t size;
+  uintptr_t addr = 0;
+  uintptr_t size = 0;
   uintptr_t* trap_table = &rt_trap_table;
 
   /* query if there is backing store */
-  size = MEGAPAGE_DOWN(sbi_query_multimem());
+  int ret = sbi_query_multimem(&size);
+  size = MEGAPAGE_DOWN(size);
 
-  if (!size) {
+  if (ret || !size) {
 		warn("no backing store found\n");
     return;
   }
 
   /* query the backing store PA */
-  addr = MEGAPAGE_UP(sbi_query_multimem_addr());
+  ret = sbi_query_multimem_addr(&addr);
+  addr = MEGAPAGE_UP(addr);
 
-  if (!addr) {
+  if (ret || !addr) {
 		warn("address is zero\n");
     return;
   }
diff --git a/sbi.c b/sbi.c
index ead4206..97506f2 100644
--- a/sbi.c
+++ b/sbi.c
@@ -56,15 +56,15 @@ sbi_random() {
 }
 
 uintptr_t
-sbi_query_multimem() {
-  return SBI_CALL_2(SBI_EXT_EXPERIMENTAL_KEYSTONE_ENCLAVE,
-      SBI_SM_CALL_PLUGIN, SM_MULTIMEM_PLUGIN_ID, SM_MULTIMEM_CALL_GET_SIZE);
+sbi_query_multimem(size_t *size) {
+  return SBI_CALL_3(SBI_EXT_EXPERIMENTAL_KEYSTONE_ENCLAVE,
+      SBI_SM_CALL_PLUGIN, SM_MULTIMEM_PLUGIN_ID, SM_MULTIMEM_CALL_GET_SIZE, size);
 }
 
 uintptr_t
-sbi_query_multimem_addr() {
-  return SBI_CALL_2(SBI_EXT_EXPERIMENTAL_KEYSTONE_ENCLAVE,
-      SBI_SM_CALL_PLUGIN, SM_MULTIMEM_PLUGIN_ID, SM_MULTIMEM_CALL_GET_ADDR);
+sbi_query_multimem_addr(uintptr_t *addr) {
+  return SBI_CALL_3(SBI_EXT_EXPERIMENTAL_KEYSTONE_ENCLAVE,
+      SBI_SM_CALL_PLUGIN, SM_MULTIMEM_PLUGIN_ID, SM_MULTIMEM_CALL_GET_ADDR, addr);
 }
 
 uintptr_t
diff --git a/sbi.h b/sbi.h
index 30fc840..b3dffd0 100644
--- a/sbi.h
+++ b/sbi.h
@@ -6,6 +6,7 @@
 #define __SBI_H_
 
 #include <stdint.h>
+#include <stddef.h>
 
 #define SBI_SET_TIMER 0
 #define SBI_CONSOLE_PUTCHAR 1
@@ -38,9 +39,9 @@ sbi_exit_enclave(uint64_t retval);
 uintptr_t
 sbi_random();
 uintptr_t
-sbi_query_multimem();
+sbi_query_multimem(size_t *size);
 uintptr_t
-sbi_query_multimem_addr();
+sbi_query_multimem_addr(uintptr_t *addr);
 uintptr_t
 sbi_attest_enclave(void* report, void* buf, uintptr_t len);
 uintptr_t