[LibOS] Fixing signal handling scheme

LibOS/shim/include/shim_internal.h

@@ -195,9 +195,16 @@ long convert_pal_errno (long err);
 
 void check_stack_hook (void);
 
+static inline uint64_t get_cur_preempt (void) {
+    shim_tcb_t* tcb = SHIM_GET_TLS();
+    assert(tcb);
+    return tcb->context.preempt;
+}
+
 #define BEGIN_SHIM(name, args ...)                          \
     SHIM_ARG_TYPE __shim_##name (args) {                    \
         SHIM_ARG_TYPE ret = 0;                              \
+        uint64_t preempt = get_cur_preempt();               \
         /* handle_signal(true); */                          \
         /* check_stack_hook(); */                           \
         BEGIN_SYSCALL_PROFILE();
@@ -205,6 +212,7 @@ void check_stack_hook (void);
 #define END_SHIM(name)                                      \
         END_SYSCALL_PROFILE(name);                          \
         handle_signal(false);                               \
+        assert(preempt == get_cur_preempt());               \
         return ret;                                         \
     }
 
@@ -390,14 +398,13 @@ void parse_syscall_after (int sysno, const char * name, int nr, ...);
  * container_of - cast a member of a structure out to the containing structure
  * @ptr:	the pointer to the member.
  * @type:	the type of the container struct this is embedded in.
- * @member:	the name of the member within the struct.
+ * @field:	the name of the field within the struct.
  *
  */
-#define container_of(ptr, type, member) ({			\
-	const typeof( ((type *)0)->member ) *__mptr = (ptr);	\
-	(type *)( (char *)__mptr - offsetof(type,member) );})
+#define container_of(ptr, type, field) ((type *)((char *)(ptr) - offsetof(type, field)))
 #endif
 
+
 #define CONCAT2(t1, t2) __CONCAT2(t1, t2)
 #define __CONCAT2(t1, t2) t1##_##t2
 
@@ -431,6 +438,8 @@ static inline PAL_HANDLE thread_create (void * func, void * arg, int option)
 static inline void __disable_preempt (shim_tcb_t * tcb)
 {
     //tcb->context.syscall_nr += SYSCALL_NR_PREEMPT_INC;
+    /* Assert if this counter overflows */
+    assert((tcb->context.preempt & ~SIGNAL_DELAYED) != ~SIGNAL_DELAYED);
     tcb->context.preempt++;
     //debug("disable preempt: %d\n", tcb->context.preempt & ~SIGNAL_DELAYED);
 }
@@ -446,6 +455,8 @@ static inline void disable_preempt (shim_tcb_t * tcb)
 static inline void __enable_preempt (shim_tcb_t * tcb)
 {
     //tcb->context.syscall_nr -= SYSCALL_NR_PREEMPT_INC;
+    /* Assert if this counter underflows */
+    assert(tcb->context.preempt > 0);
     tcb->context.preempt--;
     //debug("enable preempt: %d\n", tcb->context.preempt & ~SIGNAL_DELAYED);
 }

LibOS/shim/include/shim_thread.h

@@ -162,6 +162,14 @@ struct shim_thread * get_cur_thread (void)
     return SHIM_THREAD_SELF();
 }
 
+static inline
+__attribute__((always_inline))
+bool cur_thread_is_alive (void)
+{
+    struct shim_thread * thread = get_cur_thread();
+    return thread ? thread->is_alive : false;
+}
+
 static inline
 __attribute__((always_inline))
 void set_cur_thread (struct shim_thread * thread)
@@ -304,8 +312,6 @@ struct clone_args {
 
 int clone_implementation_wrapper(struct clone_args * arg);
 
-int clean_held_locks (struct shim_thread * self);
-
 void * allocate_stack (size_t size, size_t protect_size, bool user);
 int populate_user_stack (void * stack, size_t stack_size,
                          int nauxv, elf_auxv_t ** auxpp,

LibOS/shim/include/shim_tls.h

@@ -51,8 +51,8 @@ struct shim_context {
     void *                  ret_ip;
     struct shim_regs *      regs;
     struct shim_context *   next;
-    unsigned long           enter_time;
-    unsigned long           preempt;
+    uint64_t                enter_time;
+    uint64_t                preempt;
 };
 
 #ifdef IN_SHIM
@@ -71,8 +71,6 @@ typedef struct {
     unsigned int            tid;
     int                     pal_errno;
     void *                  debug_buf;
-    int                     last_lock;
-    struct lock_record      held_locks[NUM_LOCK_RECORD];
 
     /* This record is for testing the memory of user inputs.
      * If a segfault occurs with the range [start, end],

LibOS/shim/src/bookkeep/shim_signal.c

@@ -146,8 +146,11 @@ void __store_context (shim_tcb_t * tcb, PAL_CONTEXT * pal_context,
 void deliver_signal (siginfo_t * info, PAL_CONTEXT * context)
 {
     shim_tcb_t * tcb = SHIM_GET_TLS();
+    assert(tcb);
 
-    if (!tcb || !tcb->tp)
+    // Signals should not be delivered before the user process starts
+    // or after the user process dies.
+    if (!tcb->tp || !cur_thread_is_alive())
         return;
 
     struct shim_thread * cur_thread = (struct shim_thread *) tcb->tp;
@@ -248,6 +251,8 @@ ret_exception:
 static void memfault_upcall (PAL_PTR event, PAL_NUM arg, PAL_CONTEXT * context)
 {
     shim_tcb_t * tcb = SHIM_GET_TLS();
+    assert(tcb);
+
     if (tcb->test_range.cont_addr && arg
         && (void *) arg >= tcb->test_range.start
         && (void *) arg <= tcb->test_range.end) {
@@ -454,10 +459,7 @@ static void resume_upcall (PAL_PTR event, PAL_NUM arg, PAL_CONTEXT * context)
         goto ret_exception;
 
     shim_tcb_t * tcb = SHIM_GET_TLS();
-
-    if (!tcb || !tcb->tp)
-        return;
-
+    assert(tcb);
     __disable_preempt(tcb);
 
     if ((tcb->context.preempt & ~SIGNAL_DELAYED) > 1) {
@@ -614,16 +616,12 @@ void __handle_signal (shim_tcb_t * tcb, int sig, ucontext_t * uc)
 void handle_signal (bool delayed_only)
 {
     shim_tcb_t * tcb = SHIM_GET_TLS();
-
-    if (!tcb || !tcb->tp)
-        return;
+    assert(tcb);
 
     struct shim_thread * thread = (struct shim_thread *) tcb->tp;
 
-    debug("handle signal (counter = %d)\n", thread->has_signal.counter);
-
     /* Fast path */
-    if (!thread->has_signal.counter)
+    if (!thread || !thread->has_signal.counter)
         return;
 
     __disable_preempt(tcb);

LibOS/shim/src/bookkeep/shim_thread.c

@@ -148,6 +148,7 @@ static IDTYPE get_internal_pid (void)
     internal_tid_alloc_idx++;
     IDTYPE idx = internal_tid_alloc_idx;
     unlock(thread_list_lock);
+    assert(IS_INTERNAL_TID(idx));
     return idx;
 }
 
@@ -754,6 +755,8 @@ BEGIN_RS_FUNC(running_thread)
             assert(tcb->context.sp);
             tcb->debug_buf = SHIM_GET_TLS()->debug_buf;
             allocate_tls(libc_tcb, thread->user_tcb, thread);
+            /* Temporarily disable preemption until the thread resumes. */
+            __disable_preempt(tcb);
             debug_setprefix(tcb);
             debug("after resume, set tcb to %p\n", libc_tcb);
         } else {
@@ -763,7 +766,7 @@ BEGIN_RS_FUNC(running_thread)
         thread->in_vm = thread->is_alive = true;
         thread->pal_handle = PAL_CB(first_thread);
     }
-    
+
     DEBUG_RS("tid=%d", thread->tid);
 }
 END_RS_FUNC(running_thread)

LibOS/shim/src/elf/shim_rtld.c

@@ -1573,6 +1573,10 @@ int execute_elf_object (struct shim_handle * exec, int argc, const char ** argp,
 
     ElfW(Addr) entry = interp_map ? interp_map->l_entry : exec_map->l_entry;
 
+    /* Ready to start execution, re-enable preemption. */
+    shim_tcb_t * tcb = SHIM_GET_TLS();
+    __enable_preempt(tcb);
+
 #if defined(__x86_64__)
     asm volatile (
                     "movq %%rbx, %%rsp\r\n"

LibOS/shim/src/shim_checkpoint.c

@@ -1256,6 +1256,10 @@ void restore_context (struct shim_context * context)
     regs[nregs] = (void *) context->sp - 8;
     *(void **) (context->sp - 8) = context->ret_ip;
 
+    /* Ready to resume execution, re-enable preemption. */
+    shim_tcb_t * tcb = SHIM_GET_TLS();
+    __enable_preempt(tcb);
+
     memset(context, 0, sizeof(struct shim_context));
 
     asm volatile("movq %0, %%rsp\r\n"

LibOS/shim/src/shim_init.c

@@ -671,6 +671,8 @@ int shim_init (int argc, void * args, void ** return_stack)
     __libc_tcb_t tcb;
     memset(&tcb, 0, sizeof(__libc_tcb_t));
     allocate_tls(&tcb, false, NULL);
+    __disable_preempt(&tcb.shim_tcb); // Temporarily disable preemption for delaying any signal
+                                      // that arrives during initialization
     debug_setbuf(&tcb.shim_tcb, true);
     debug("set tcb to %p\n", &tcb);
 

LibOS/shim/src/sys/shim_clone.c

@@ -110,6 +110,8 @@ int clone_implementation_wrapper(struct clone_args * arg)
     }
     allocate_tls(my_thread->tcb, my_thread->user_tcb, my_thread);
     shim_tcb_t * tcb = &((__libc_tcb_t *) my_thread->tcb)->shim_tcb;
+    __disable_preempt(tcb); // Temporarily disable preemption, because the preemption
+                            // will be re-enabled when the thread starts.
     debug_setbuf(tcb, true);
     debug("set tcb to %p (stack allocated? %d)\n", my_thread->tcb, stack_allocated);
 

LibOS/shim/src/sys/shim_exec.c

@@ -91,6 +91,8 @@ int shim_do_execve_rtld (struct shim_handle * hdl, const char ** argv,
         return -ENOMEM;
 
     populate_tls(tcb, false);
+    __disable_preempt(&((__libc_tcb_t *) tcb)->shim_tcb); // Temporarily disable preemption
+                                                          // during execve().
     debug("set tcb to %p\n", tcb);
 
     put_handle(cur_thread->exec);