Repository: sancus-tee/sancus-compiler
Issue category: Security Vulnerability

Issue:
Check for integer overflow in sancus_is_outside_sm macro



Hints:
thanks for the PR! Not sure about your first clause "( __OUTSIDE_SM(p, sm) && ((len <= 0) ||". If the length is zero, the condition should always be true and the outside_sm is not necessary I'd say. A negative length doesn't make any sense and the disadvantage of a macro is that we can't easily enforce an `unsigned` type as in an inlined function..

The logical OR is flawed I think. You never check the first condition `__OUTSIDE_SM(p, sm)` in the OR clause? So only the end pointer is checked, which renders the check vulnerable AFAIS?

I feel the logic is getting to complex for a macro and I think this will be better implemented as an always_inline function with typed arguments, making use of the OUTSIDE_SM helper macros?

Maybe a good place to start is the logic from the SGX-SDK:

https://github.com/intel/linux-sgx/blob/4589daddd58bec7367a6a9de3fe301e6de17671a/sdk/trts/trts.cpp#L101

```C
// sgx_is_outside_enclave()
// Parameters:
//      addr - the start address of the buffer
//      size - the size of the buffer
// Return Value:
//      1 - the buffer is strictly outside the enclave
//      0 - the whole buffer or part of the buffer is not outside the enclave,
//          or the buffer is wrap around
//
int sgx_is_outside_enclave(const void *addr, size_t size)
{
    size_t start = reinterpret_cast<size_t>(addr);
    size_t end = 0;
    size_t enclave_start = (size_t)&__ImageBase;
    size_t enclave_end = enclave_start + g_global_data.enclave_size - 1;
    // g_global_data.enclave_end = enclave_base + enclave_size - 1;
    // so the enclave range is [enclave_start, enclave_end] inclusively

    if(size > 0)
    {
        end = start + size - 1;
    }
    else
    {
        end = start;
    }
    if( (start <= end) && ((end < enclave_start) || (start > enclave_end)) )
    {
        return 1;
    }
    return 0;
}
```
Thanks for the comments! I think you might have mentally skipped a bracket, as the start address is always checked, and then there is (size is below or equal to zero) OR (buffer does not wrap around and end address is outside of enclave), which is similar to the SGX logic. Nevertheless, this might be the proof this should be in a function, I'll take care of it! 
thanks, yes you're absolutely right I mis-parsed the brackets my bad! :sweat_smile: It's indeed an indication that we'll be better off with an inlined function :)
hm important update--re-reading the Tale of 2 worlds analysis, I remember the problem w this code goes beyond merely integer overflows:

> Sancus. We found both logical errors and integer overflow vulnerabilities
in the sancus_is_outside_sm() function provided by the trusted runtime.
Particularly, the current implementation does not properly detect an untrusted
buffer that spans the entire enclave address range, or a rogue length specifier
that triggers an integer overflow to wrap around the 16-bit address space.

The code needs to be rewritten as in the Intel SGX code stub above. Currently, the code is fine if your buffer starts below the enclave and then (without overflowing) ends after the enclave, overwriting the entire enclave(!)

Patch:
diff --git a/src/drivers/linker.py b/src/drivers/linker.py
index 4f3ec85..c369e43 100644
--- a/src/drivers/linker.py
+++ b/src/drivers/linker.py
@@ -17,6 +17,7 @@
 
 MAC_SIZE = int(sancus.config.SECURITY / 8)
 KEY_SIZE = MAC_SIZE
+CONNECTION_STRUCT_SIZE = 6 + KEY_SIZE
 
 
 class SmEntry:
@@ -63,7 +64,7 @@ def add_sym(file, sym_map):
 
     args += [file, file]
     call_prog('msp430-elf-objcopy', args)
-    return file 
+    return file
 
 
 def parse_size(val):
@@ -92,18 +93,23 @@ def get_symbol(elf_file, name):
 
 def get_io_sym_map(sm_name):
     sym_map = {
-        '__sm_handle_input':    '__sm_{}_handle_input'.format(sm_name),
-        '__sm_num_inputs':      '__sm_{}_num_inputs'.format(sm_name),
-        '__sm_num_connections': '__sm_{}_num_connections'.format(sm_name),
-        '__sm_io_keys':         '__sm_{}_io_keys'.format(sm_name),
-        '__sm_input_callbacks': '__sm_{}_input_callbacks'.format(sm_name),
-        '__sm_output_nonce':    '__sm_{}_output_nonce'.format(sm_name),
-        '__sm_send_output':     '__sm_{}_send_output'.format(sm_name),
-        '__sm_set_key':         '__sm_{}_set_key'.format(sm_name),
-        '__sm_X_exit':          '__sm_{}_exit'.format(sm_name),
-        '__sm_X_stub_malloc':   '__sm_{}_stub_malloc'.format(sm_name),
+        '__sm_handle_input':        '__sm_{}_handle_input'.format(sm_name),
+        '__sm_num_inputs':          '__sm_{}_num_inputs'.format(sm_name),
+        '__sm_num_connections':     '__sm_{}_num_connections'.format(sm_name),
+        '__sm_max_connections':     '__sm_{}_max_connections'.format(sm_name),
+        '__sm_io_connections':      '__sm_{}_io_connections'.format(sm_name),
+        '__sm_input_callbacks':     '__sm_{}_input_callbacks'.format(sm_name),
+        '__sm_send_output':         '__sm_{}_send_output'.format(sm_name),
+        '__sm_set_key':             '__sm_{}_set_key'.format(sm_name),
+        '__sm_attest':              '__sm_{}_attest'.format(sm_name),
+        '__sm_X_exit':              '__sm_{}_exit'.format(sm_name),
+        '__sm_X_stub_malloc':       '__sm_{}_stub_malloc'.format(sm_name),
         '__sm_X_stub_reactive_handle_output':
-            '__sm_{}_stub_reactive_handle_output'.format(sm_name)
+            '__sm_{}_stub_reactive_handle_output'.format(sm_name),
+        '__sm_X_public_start':      '__sm_{}_public_start'.format(sm_name),
+        '__sm_X_public_end':        '__sm_{}_public_end'.format(sm_name),
+        '__sm_X_secret_start':      '__sm_{}_secret_start'.format(sm_name),
+        '__sm_X_secret_end':        '__sm_{}_secret_end'.format(sm_name)
     }
 
     return sym_map
@@ -115,7 +121,7 @@ def get_io_sect_map(sm_name):
         '.rela.sm.X.text':  '.rela.sm.{}.text'.format(sm_name),
     }
 
-    for entry in ('__sm{}_set_key', '__sm{}_handle_input'):
+    for entry in ('__sm{}_set_key', '__sm{}_attest', '__sm{}_handle_input'):
         map['.sm.X.{}.table'.format(entry.format(''))] = \
             '.sm.{}.{}.table'.format(sm_name, entry.format('_' + sm_name))
         map['.rela.sm.X.{}.table'.format(entry.format(''))] = \
@@ -136,15 +142,19 @@ def create_io_stub(sm, stub):
 
 
 def sort_entries(entries):
-    # If the set_key entry exists, it should have index 0 and if the
-    # handle_input entry exists, it should have index 1. This is accomplished by
-    # mapping those entries to __ and ___ respectively since those come
+    # If the set_key entry exists, it should have index 0, if the
+    # attest entry exists, it should have index 1 and if
+    # handle_input entry exists, it should have index 2. This is accomplished by
+    # mapping those entries to __, ___ and ___ respectively since those come
     # alphabetically before any valid entry name.
     def sort_key(entry):
         if re.match(r'__sm_\w+_set_key', entry.name):
             return '__'
-        if re.match(r'__sm_\w+_handle_input', entry.name):
+        if re.match(r'__sm_\w+_attest', entry.name):
             return '___'
+        if re.match(r'__sm_\w+_handle_input', entry.name):
+            return '____'
+
         return entry.name
 
     entries.sort(key=sort_key)
@@ -212,7 +222,7 @@ def sort_key(entry):
     for a in archive_files:
         debug("Unpacking archive for Sancus SM inspection: " + a)
         file_name = a
-        if ':' in a: 
+        if ':' in a:
             # support calls such as -lib:/full/path
             file_name = file_name.split(':')[1]
 
@@ -269,6 +279,7 @@ def sort_key(entry):
 elf_relocations = defaultdict(list)
 
 added_set_key_stub = False
+added_attest_stub = False
 added_input_stub = False
 added_output_stub = False
 
@@ -409,6 +420,14 @@ def sort_key(entry):
                         input_files_to_scan.append(generated_file)
                         added_set_key_stub = True
 
+                    if not added_attest_stub:
+                        # Generate the attest stub file
+                        generated_file = create_io_stub(sm, 'sm_attest.o')
+                        generated_object_files.append(generated_file)
+                        # And register it to also be scanned by this loop later
+                        input_files_to_scan.append(generated_file)
+                        added_attest_stub = True
+
                     if which == 'input':
                         dest = sms_inputs
 
@@ -463,7 +482,7 @@ def sort_key(entry):
 The YAML file allows to set the following things:
   1) warn whenever an ocall is performed and abort linking process
   2) Swap out assembly stubs for custom project dependent values
-  3) Set a peripheral offset for the first SM to 
+  3) Set a peripheral offset for the first SM to
 """
 try:
     file_path = ''
@@ -829,24 +848,36 @@ def sort_key(entry):
 
         call_prog('msp430-gcc', ['-c', '-o', o_file, c_file])
 
-        input_callbacks += '    {}(.sm.{}.callbacks)\n'.format(o_file, sm)
+        input_callbacks += '    KEEP({}(.sm.{}.callbacks))\n'.format(o_file, sm)
         input_callbacks += '    . = ALIGN(2);'
 
-    # Table of connection keys
-    io_keys = ''
+    """
+    Table of connections: in a reactive application, a connection links the
+    output of a SM (defined using the macro `SM_OUTPUT`) to the input of another
+    (defined using the macro `SM_INPUT`).
+    These connections are stored in a `Connection` array on each SM (see
+    `reactive_stubs_support.h`). The array is allocated here with a fixed size,
+    according to the `num_connections` parameter in the SM config (default 0).
+    """
 
-    if len(ios) > 0:
-        io_keys += '__sm_{}_io_keys = .;\n'.format(sm)
-        io_keys += '    . += {};\n'.format(len(ios) * KEY_SIZE)
-        io_keys += '    . = ALIGN(2);'
+    num_connections = ''
+    io_connections = ''
+
+    if hasattr(sm_config[sm], "num_connections"):
+        sm_num_connections = sm_config[sm].num_connections
+    else:
+        sm_num_connections = 0
 
-    # Nonce used by outputs
-    outputs_nonce = ''
+    if len(ios) > 0:
+        # make sure we allocate space even if num_connections is zero
+        io_connections_size = max(sm_num_connections * CONNECTION_STRUCT_SIZE, 2)
 
-    if len(outputs) > 0:
-        outputs_nonce += '__sm_{}_output_nonce = .;\n'.format(sm)
-        outputs_nonce += '    . += 2;\n'
-        outputs_nonce += '    . = ALIGN(2);'
+        num_connections += '__sm_{}_num_connections = .;\n'.format(sm)
+        num_connections += '    . += 2;\n'
+        num_connections += '    . = ALIGN(2);'
+        io_connections += '__sm_{}_io_connections = .;\n'.format(sm)
+        io_connections += '    . += {};\n'.format(io_connections_size)
+        io_connections += '    . = ALIGN(2);'
 
     # Set data section offset if peripheral access is set
     # in that case, optionally provide first SM with exclusive access to last peripheral
@@ -859,10 +890,10 @@ def sort_key(entry):
                                              exit_file, '\n    '.join(tables),
                                              input_callbacks,
                                              '\n    '.join(extra_labels)))
-    
+
     data_sections.append(data_section.format(sm, '\n    '.join(id_syms),
-                                             args.sm_stack_size, io_keys,
-                                             outputs_nonce,
+                                             args.sm_stack_size, num_connections,
+                                             io_connections,
                                              data_section_start))
 
     if sm in sms_entries:
@@ -880,7 +911,7 @@ def sort_key(entry):
         symbols.append('__sm_{}_io_{}_idx = {};'.format(sm, io, index))
 
     # Add symbols for the number of connections/inputs
-    symbols.append('__sm_{}_num_connections = {};'.format(sm, len(ios)))
+    symbols.append('__sm_{}_max_connections = {};'.format(sm, sm_num_connections))
     symbols.append('__sm_{}_num_inputs = {};'.format(sm, len(inputs)))
 
     if args.prepare_for_sm_text_section_wrapping:
diff --git a/src/drivers/sancus/sancus_config.py b/src/drivers/sancus/sancus_config.py
index c9826e1..b29bf0c 100644
--- a/src/drivers/sancus/sancus_config.py
+++ b/src/drivers/sancus/sancus_config.py
@@ -12,6 +12,7 @@
 SM_CONFIG_DEFAULT_SM_EXIT=""
 SM_CONFIG_DEFAULT_SM_ISR=""
 SM_CONFIG_DEFAULT_PERIPHERAL_OFFSET=0
+SM_CONFIG_DEFAULT_NUM_CONNECTIONS=None
 
 class SmParserError(Exception):
     """ 
@@ -125,6 +126,7 @@ def _safe_config_parse(name, default_val):
         self._sm_exit  = _safe_config_parse('sm_exit',  SM_CONFIG_DEFAULT_SM_EXIT)
         self._sm_isr   = _safe_config_parse('sm_isr' ,  SM_CONFIG_DEFAULT_SM_ISR)
         self._peripheral_offset = _safe_config_parse('peripheral_offset', 0)
+        self._num_connections = _safe_config_parse('num_connections', SM_CONFIG_DEFAULT_NUM_CONNECTIONS)
 
     """
     Getters as properties for all class variables.
@@ -178,6 +180,13 @@ def peripheral_offset(self):
         else:
             raise AttributeError
     
+    @property
+    def num_connections(self):
+        if self._num_connections != SM_CONFIG_DEFAULT_NUM_CONNECTIONS:
+            return abs(self._num_connections)
+        else:
+            raise AttributeError
+
     def __str__(self):
         if self._contains_config:
             s = '' + self._name + ':\n'
@@ -187,6 +196,7 @@ def __str__(self):
             if hasattr(self, 'sm_exit'): s += '\t\t SM_EXIT: ' + self._sm_exit  + '\n'
             if hasattr(self, 'sm_isr'): s += '\t\t SM_ISR: ' + self._sm_isr  + '\n'
             if hasattr(self, 'peripheral_offset'): s += '\t\t Peripheral Offset: ' + str(self._peripheral_offset)  + '\n'
+            if hasattr(self, 'num_connections'): s += '\t\t Number of connections: ' + str(self._num_connections)  + '\n'
         else:
             s = 'No YAML config provided.'
         return s
\ No newline at end of file
diff --git a/src/drivers/sm-config-example.yaml b/src/drivers/sm-config-example.yaml
index 4c82eed..b2a8d13 100644
--- a/src/drivers/sm-config-example.yaml
+++ b/src/drivers/sm-config-example.yaml
@@ -23,6 +23,13 @@ my_sm_1:
     # amount of byte. This allows to map the data section over the
     # peripherals to grant exclusive and faster access to these peripherals
     - peripheral_offset: 159
+    # Set the maximum number of connections allowed for this SM in a reactive
+    # application. Connections link the output of a SM (defined using the macro
+    # `SM_OUTPUT`) with the input of another (defined using the macro
+    # `SM_INPUT`). If this SM does not have any inputs nor outputs, there is no
+    # need to specify this parameter (default value is zero). In the Authentic
+    # Execution framework, this parameter is automatically configured.
+    - num_connections: 3
 my_sm_2:
  	- disallow_outcalls: True
     
\ No newline at end of file
diff --git a/src/sancus_support/reactive.h b/src/sancus_support/reactive.h
index db60b61..d3bebc5 100644
--- a/src/sancus_support/reactive.h
+++ b/src/sancus_support/reactive.h
@@ -6,19 +6,20 @@
 #include <stdint.h>
 
 typedef uint16_t io_index;
+typedef uint16_t conn_index;
 typedef uint8_t  io_data __attribute__((aligned(2)));
 
 // The ASM symbols are used for the linker to be able to detect inputs/outputs
 
-#define SM_OUTPUT_AUX(sm, name)                                         \
-    asm("__sm_" #sm "_output_tag_" #name " = 0\n");                     \
-    SM_FUNC(sm) void name(const io_data* data, size_t len)              \
-    {                                                                   \
-        extern char __sm_##sm##_io_##name##_idx;                        \
-        SM_FUNC(sm) void __sm_##sm##_send_output(unsigned int,          \
-                                                 const void*, size_t);  \
-        __sm_##sm##_send_output((io_index)&__sm_##sm##_io_##name##_idx, \
-                                data, len);                             \
+#define SM_OUTPUT_AUX(sm, name)                                                \
+    asm("__sm_" #sm "_output_tag_" #name " = 0\n");                            \
+    SM_FUNC(sm) uint16_t name(const io_data* data, size_t len)                 \
+    {                                                                          \
+        extern char __sm_##sm##_io_##name##_idx;                               \
+        SM_FUNC(sm) uint16_t __sm_##sm##_send_output(unsigned int,             \
+                                                 const void*, size_t);         \
+        return __sm_##sm##_send_output((io_index)&__sm_##sm##_io_##name##_idx, \
+                                data, len);                                    \
     }
 
 #define SM_OUTPUT(sm, name) SM_OUTPUT_AUX(sm, name)
diff --git a/src/sancus_support/sm_support.h b/src/sancus_support/sm_support.h
index 9881973..12264ef 100644
--- a/src/sancus_support/sm_support.h
+++ b/src/sancus_support/sm_support.h
@@ -4,6 +4,7 @@
 #include "config.h"
 
 #include <stddef.h>
+#include <stdint.h>
 
 #if __GNUC__ >= 5 || __clang_major__ >= 5
 
@@ -199,15 +200,12 @@ extern char __unprotected_sp;
 
 #endif
 
-#define __OUTSIDE_SM( p, sm )                                                  \
-    ( ((void*) p < (void*) &__PS(sm)) || ((void*) p >= (void*) &__PE(sm)) ) && \
-    ( ((void*) p < (void*) &__SS(sm)) || ((void*) p >= (void*) &__SE(sm)) )
-
 /*
  * Returns true iff whole buffer [p,p+len-1] is outside of the sm SancusModule
  */
-#define sancus_is_outside_sm( sm, p, len) \
-    ( __OUTSIDE_SM(p, sm) && __OUTSIDE_SM((p+len-1), sm) )
+#define sancus_is_outside_sm(sm, p, len) \
+    ( is_buffer_outside_region(&__PS(sm), &__PE(sm), p, len) && \
+      is_buffer_outside_region(&__SS(sm), &__SE(sm), p, len) )
 
 /**
  * Interrupt vector for the Sancus violation ISR.
@@ -311,6 +309,37 @@ sm_id sancus_enable_wrapped(struct SancusModule* sm, unsigned nonce, void* tag);
 #undef always_inline
 #define always_inline static inline __attribute__((always_inline))
 
+/*
+ * Returns true if buf is outside the memory region [start, end)
+ * if start >= end, immediately return false
+ */
+always_inline int is_buffer_outside_region(void *start_p, void *end_p,
+    void *buf_p, size_t len) {
+  uintptr_t start = (uintptr_t) start_p;
+  uintptr_t end   = (uintptr_t) end_p;
+  uintptr_t buf   = (uintptr_t) buf_p;
+  uintptr_t buf_end;
+  
+  // make sure start < end, otherwise return false
+  if (start >= end) {
+    return 0;
+  }
+
+  if(len > 0) {
+    buf_end = buf + len - 1;
+  }
+  else {
+    buf_end = buf;
+  }
+
+  /* check for int overflow and finally validate `buf` falls outside */ 
+  if( (buf <= buf_end) && ((end <= buf) || (start > buf_end))) {
+    return 1;
+  }
+
+  return 0;
+}
+
 /**
  * Performs constant time comparison between to buffers
  * Returns 0 if the first `n` bytes of the two buffers are equal, -1 otherwise
diff --git a/src/stubs/CMakeLists.txt b/src/stubs/CMakeLists.txt
index 9427a17..9172a8e 100644
--- a/src/stubs/CMakeLists.txt
+++ b/src/stubs/CMakeLists.txt
@@ -23,6 +23,7 @@ set(EXTRA_FLAGS -I${CMAKE_SOURCE_DIR}/src/sancus_support)
 add_object(sm_output.o sm_output.c ${EXTRA_FLAGS})
 add_object(sm_input.o sm_input.c ${EXTRA_FLAGS})
 add_object(sm_set_key.o sm_set_key.c ${EXTRA_FLAGS})
+add_object(sm_attest.o sm_attest.c ${EXTRA_FLAGS})
 
 set(STUBS
     ${CMAKE_CURRENT_BINARY_DIR}/sm_entry.o
@@ -34,6 +35,7 @@ set(STUBS
     ${CMAKE_CURRENT_BINARY_DIR}/sm_output.o
     ${CMAKE_CURRENT_BINARY_DIR}/sm_input.o
     ${CMAKE_CURRENT_BINARY_DIR}/sm_set_key.o
+    ${CMAKE_CURRENT_BINARY_DIR}/sm_attest.o
     ${CMAKE_CURRENT_BINARY_DIR}/sm_mmio_entry.o
     ${CMAKE_CURRENT_BINARY_DIR}/sm_mmio_exclusive.o
 
diff --git a/src/stubs/reactive_stubs_support.h b/src/stubs/reactive_stubs_support.h
index c5aca41..3c7a251 100644
--- a/src/stubs/reactive_stubs_support.h
+++ b/src/stubs/reactive_stubs_support.h
@@ -3,29 +3,54 @@
 
 #include "reactive.h"
 
-void reactive_handle_output(io_index output_id, void* data, size_t len);
+void reactive_handle_output(conn_index conn_id, void* data, size_t len);
 
 typedef uint8_t IoKey[SANCUS_KEY_SIZE];
 typedef void (*InputCallback)(const void*, size_t);
 
 typedef enum
 {
-    Ok                = 0x0,
-    IllegalConnection = 0x1,
-    MalformedPayload  = 0x2
+    Ok                    = 0x0,
+    IllegalConnection     = 0x1,
+    MalformedPayload      = 0x2,
+    IllegalParameters     = 0x3,
+    BufferInsideSM        = 0x4,
+    CryptoError           = 0x5,
+    InternalError         = 0x6
 } ResultCode;
 
+typedef struct Connection {
+    io_index    io_id;
+    conn_index  conn_id;
+    uint16_t    nonce;
+    IoKey       key;
+} Connection;
+
+// The size of the Connection struct is also hardcoded in linker.py. Hence,
+// we need to make sure that it does not change at compile time (e.g. due to
+// optimizations).
+// Besides, if the struct changes, we need to adjust this value here and in
+// linker.py (check the CONNECTION_STRUCT_SIZE global variable) as well.
+_Static_assert (sizeof(Connection) == 6 + SANCUS_KEY_SIZE,
+    "Size of Connection struct differs from the expected value");
+
 // These will be allocated by the linker
-extern IoKey __sm_io_keys[];
+extern Connection __sm_io_connections[];
 extern InputCallback __sm_input_callbacks[];
-extern uint16_t __sm_output_nonce;
 
-extern char __sm_num_connections;
-#define SM_NUM_CONNECTIONS (size_t)&__sm_num_connections
+extern char __sm_max_connections;
+#define SM_MAX_CONNECTIONS (size_t)&__sm_max_connections
+
+extern uint16_t __sm_num_connections;
 
 extern char __sm_num_inputs;
 #define SM_NUM_INPUTS (size_t)&__sm_num_inputs
 
 #define SM_NAME X
 
+// declare symbols for the public/secret regions
+#define __SECTION(sect, name) sect(name)
+extern char __SECTION(__PS, SM_NAME), __SECTION(__PE, SM_NAME),
+            __SECTION(__SS, SM_NAME), __SECTION(__SE, SM_NAME);
+
 #endif
diff --git a/src/stubs/sm_input.c b/src/stubs/sm_input.c
index 992d2e6..05f3f55 100644
--- a/src/stubs/sm_input.c
+++ b/src/stubs/sm_input.c
@@ -2,24 +2,42 @@
 
 #include <alloca.h>
 
-#define AD_SIZE 2
-
-void SM_ENTRY(SM_NAME) __sm_handle_input(uint16_t conn_id,
+uint16_t SM_ENTRY(SM_NAME) __sm_handle_input(uint16_t conn_idx,
                                          const void* payload, size_t len)
 {
-    if (conn_id >= SM_NUM_INPUTS)
-        return;
+    // sanitize input buffer
+    if(!sancus_is_outside_sm(SM_NAME, (void *) payload, len)) {
+      return BufferInsideSM;
+    }
 
-    const size_t data_len = len - AD_SIZE - SANCUS_TAG_SIZE;
-    const uint8_t* cipher = (uint8_t*)payload + AD_SIZE;
-    const uint8_t* tag = cipher + data_len;
+    // check correctness of other parameters
+    if(len < SANCUS_TAG_SIZE || conn_idx >= __sm_num_connections) {
+      return IllegalParameters;
+    }
+
+    Connection *conn = &__sm_io_connections[conn_idx];
 
+    // check if io_id is a valid input ID
+    if (conn->io_id >= SM_NUM_INPUTS) {
+      return IllegalConnection;
+    }
+
+    // associated data only contains the nonce, therefore we can use this
+    // this trick to build the array fastly (i.e. by swapping the bytes)
+    const uint16_t nonce_rev = conn->nonce << 8 | conn->nonce >> 8;
+    const size_t data_len = len - SANCUS_TAG_SIZE;
+    const uint8_t* cipher = payload;
+    const uint8_t* tag = cipher + data_len;
     // TODO check for stack overflow!
     uint8_t* input_buffer = alloca(data_len);
 
-    if (sancus_unwrap_with_key(__sm_io_keys[conn_id], payload, AD_SIZE,
-                               cipher, data_len, tag, input_buffer))
-    {
-        __sm_input_callbacks[conn_id](input_buffer, data_len);
+    if (sancus_unwrap_with_key(conn->key, &nonce_rev, sizeof(nonce_rev),
+                               cipher, data_len, tag, input_buffer)) {
+       conn->nonce++;
+       __sm_input_callbacks[conn->io_id](input_buffer, data_len);
+       return Ok;
     }
+
+    // here only if decryption fails
+    return CryptoError;
 }
diff --git a/src/stubs/sm_output.c b/src/stubs/sm_output.c
index 6febf29..cc71424 100644
--- a/src/stubs/sm_output.c
+++ b/src/stubs/sm_output.c
@@ -2,14 +2,38 @@
 
 #include <stdlib.h>
 
-SM_FUNC(SM_NAME) void __sm_send_output(io_index index,
+SM_FUNC(SM_NAME) uint16_t __sm_send_output(io_index index,
                                        const void* data, size_t len)
 {
-    const size_t nonce_len = sizeof(__sm_output_nonce);
-    const size_t payload_len = nonce_len + len + SANCUS_TAG_SIZE;
-    uint8_t* payload = malloc(payload_len);
-    *(uint16_t*)payload = __sm_output_nonce++;
-    sancus_wrap_with_key(__sm_io_keys[index], payload, nonce_len, data, len,
-                         payload + nonce_len, payload + nonce_len + len);
-    reactive_handle_output(index, payload, payload_len);
+    const size_t payload_len = len + SANCUS_TAG_SIZE;
+
+    // search for all the connections associated to the index.
+    // Unfortunately, this operation is O(n) with n = number of connections
+    int i;
+    for (i=0; i<__sm_num_connections; i++) {
+      Connection *conn = &__sm_io_connections[i];
+      if (conn->io_id != index) {
+          continue;
+      }
+
+      uint8_t* payload = malloc(payload_len);
+
+      if (payload == NULL) {
+          return InternalError;
+      }
+
+      if ( !sancus_is_outside_sm(SM_NAME, (void *) payload, payload_len) ) {
+          return BufferInsideSM;
+      }
+
+      // associated data only contains the nonce, therefore we can use this
+      // this trick to build the array fastly (i.e. by swapping the bytes)
+      uint16_t nonce_rev = conn->nonce << 8 | conn->nonce >> 8;
+      sancus_wrap_with_key(conn->key, &nonce_rev, sizeof(nonce_rev), data,
+                          len, payload, payload + len);
+      conn->nonce++;
+      reactive_handle_output(conn->conn_id, payload, payload_len);
+    }
+
+    return Ok;
 }
diff --git a/src/stubs/sm_set_key.c b/src/stubs/sm_set_key.c
index e982d4c..7694958 100644
--- a/src/stubs/sm_set_key.c
+++ b/src/stubs/sm_set_key.c
@@ -1,21 +1,43 @@
 #include "reactive_stubs_support.h"
 
-void SM_ENTRY(SM_NAME) __sm_set_key(const uint8_t* ad, const uint8_t* cipher,
-                                    const uint8_t* tag, uint8_t* result)
+#define AD_SIZE 6
+
+uint16_t SM_ENTRY(SM_NAME) __sm_set_key(const uint8_t* ad, const uint8_t* cipher,
+                                    const uint8_t* tag, uint16_t *conn_idx)
 {
-    uint16_t conn_id = (ad[2] << 8) | ad[3];
-    ResultCode code = Ok;
-
-    if (conn_id >= SM_NUM_CONNECTIONS)
-        code = IllegalConnection;
-    else if (!sancus_unwrap(ad, 4, cipher, SANCUS_KEY_SIZE, tag,
-             __sm_io_keys[conn_id]))
-    {
-        code = MalformedPayload;
+    if( !sancus_is_outside_sm(SM_NAME, (void *) ad, AD_SIZE) ||
+        !sancus_is_outside_sm(SM_NAME, (void *) cipher, SANCUS_KEY_SIZE) ||
+        !sancus_is_outside_sm(SM_NAME, (void *) tag, SANCUS_TAG_SIZE) ||
+        !sancus_is_outside_sm(SM_NAME, (void *) conn_idx, sizeof(uint16_t)) ) {
+      return BufferInsideSM;
+    }
+
+    // Note: make sure we only use AD_SIZE bytes of the buffer `ad`
+    conn_index conn_id = (ad[0] << 8) | ad[1];
+    io_index io_id = (ad[2] << 8) | ad[3];
+    uint16_t nonce = (ad[4] << 8) | ad[5];
+
+    // check if there is still space left in the array
+    if (__sm_num_connections == SM_MAX_CONNECTIONS) {
+      return InternalError;
+    }
+
+    // check nonce
+    if(nonce != __sm_num_connections) {
+      return MalformedPayload;
     }
 
-    result[0] = 0;
-    result[1] = code;
-    uint8_t result_ad[] = {ad[0], ad[1], result[0], result[1]};
-    sancus_tag(result_ad, sizeof(result_ad), result + 2);
+    Connection *conn = &__sm_io_connections[__sm_num_connections];
+    *conn_idx = __sm_num_connections;
+
+    if (!sancus_unwrap(ad, AD_SIZE, cipher, SANCUS_KEY_SIZE, tag, conn->key)) {
+      return CryptoError;
+    }
+
+    __sm_num_connections++;
+    conn->io_id = io_id;
+    conn->conn_id = conn_id;
+    conn->nonce = 0;
+
+    return Ok;
 }


Test patch:
diff --git a/src/stubs/sm_attest.c b/src/stubs/sm_attest.c
new file mode 100644
index 0000000..6bfec25
--- /dev/null
+++ b/src/stubs/sm_attest.c
@@ -0,0 +1,16 @@
+#include "reactive_stubs_support.h"
+
+uint16_t SM_ENTRY(SM_NAME) __sm_attest(const uint8_t* challenge, size_t len,
+  uint8_t *result)
+{
+    if( !sancus_is_outside_sm(SM_NAME, (void *) challenge, len) ||
+        !sancus_is_outside_sm(SM_NAME, (void *) result, SANCUS_TAG_SIZE) ) {
+      return BufferInsideSM;
+    }
+
+    if( !sancus_tag(challenge, len, result) ) {
+        return CryptoError;
+    }
+
+    return Ok;
+}