#include <linux/ptrace.h>
 #include <kern_util.h>
 #include <sysdep/ptrace.h>
+#include <sysdep/ptrace_user.h>
 #include <sysdep/syscalls.h>
-#include <os.h>
 
 void handle_syscall(struct uml_pt_regs *r)
 {
        struct pt_regs *regs = container_of(r, struct pt_regs, regs);
-       long result;
        int syscall;
 
-       if (syscall_trace_enter(regs)) {
-               result = -ENOSYS;
+       /* Initialize the syscall number and default return value. */
+       UPT_SYSCALL_NR(r) = PT_SYSCALL_NR(r->gp);
+       PT_REGS_SET_SYSCALL_RETURN(regs, -ENOSYS);
+
+       if (syscall_trace_enter(regs))
                goto out;
-       }
 
-       syscall = get_syscall(r);
+       /* Update the syscall number after orig_ax has potentially been updated
+        * with ptrace.
+        */
+       UPT_SYSCALL_NR(r) = PT_SYSCALL_NR(r->gp);
+       syscall = UPT_SYSCALL_NR(r);
 
-       if ((syscall > __NR_syscall_max) || syscall < 0)
-               result = -ENOSYS;
-       else
-               result = EXECUTE_SYSCALL(syscall, regs);
+       if (syscall >= 0 && syscall <= __NR_syscall_max)
+               PT_REGS_SET_SYSCALL_RETURN(regs,
+                               EXECUTE_SYSCALL(syscall, regs));
 
 out:
-       PT_REGS_SET_SYSCALL_RETURN(regs, result);
-
        syscall_trace_leave(regs);
 }
 
        handle_syscall(regs);
 }
 
-int get_syscall(struct uml_pt_regs *regs)
-{
-       UPT_SYSCALL_NR(regs) = PT_SYSCALL_NR(regs->gp);
-
-       return UPT_SYSCALL_NR(regs);
-}
-
 extern char __syscall_stub_start[];
 
 static int userspace_tramp(void *stack)
 
        [EFL] = HOST_EFLAGS,
        [UESP] = HOST_SP,
        [SS] = HOST_SS,
+       [ORIG_EAX] = HOST_ORIG_AX,
 };
 
 int putreg(struct task_struct *child, int regno, unsigned long value)
        case EAX:
        case EIP:
        case UESP:
+       case ORIG_EAX:
                break;
        case FS:
                if (value && (value & 3) != 3)
                value &= FLAG_MASK;
                child->thread.regs.regs.gp[HOST_EFLAGS] |= value;
                return 0;
-       case ORIG_EAX:
-               child->thread.regs.regs.syscall = value;
-               return 0;
        default :
                panic("Bad register in putreg() : %d\n", regno);
        }
 
        regno >>= 2;
        switch (regno) {
-       case ORIG_EAX:
-               return child->thread.regs.regs.syscall;
        case FS:
        case GS:
        case DS:
        case EDI:
        case EBP:
        case EFL:
+       case ORIG_EAX:
                break;
        default:
                panic("Bad register in getreg() : %d\n", regno);