Nucleus

diff --git a/include/nucleus/cpu.h b/include/nucleus/cpu.h
index 88c7506..b8293b8 100644
--- a/include/nucleus/cpu.h
+++ b/include/nucleus/cpu.h
@@ -23,20 +23,6 @@ extern cpu_t lapicNums[];
 #define CPUID (apic ? lapicNums[(cpu_t) (LAPIC(0x20) >> 24)] : 0)
 #define MAX_CPUS 2
 
-/* Push/pop interrupts */
-static inline uintptr_t
-irq_disable(void)
-{
-	uintptr_t flags;
-	asm volatile("pushf; cli; pop %0" : "=r" (flags) :: "memory");
-	return flags;
-}
-static inline void
-irq_restore(uintptr_t flags)
-{
-	asm volatile("push %0; popf" :: "rm" (flags) : "memory","cc");
-}
-
 void register_exception(int num, exc_handler_t addr);
 void register_interrupt(int num, int_handler_t addr);
 
diff --git a/include/nucleus/task.h b/include/nucleus/task.h
index cfc0c99..11ad4ae 100644
--- a/include/nucleus/task.h
+++ b/include/nucleus/task.h
@@ -35,7 +35,9 @@ struct Task {
 	gid_t gid, egid, sgid;
 	enum Priority priority;
 	enum State state;
+	uint32_t inCriticalSection;
 
+	uintptr_t esi, edi, ebx;
 	uintptr_t esp, ebp, eip;
 	page_dir_t pageDir;
 
@@ -56,6 +58,19 @@ super_user(void)
 	return (current->euid == 0);
 }
 
+/* Enter a critical section */
+static inline void
+enter_critical_section(void)
+{
+	__atomic_add_fetch(&current->inCriticalSection, 1, __ATOMIC_RELAXED);
+}
+/* Exit a critical section */
+static inline void
+exit_critical_section(void)
+{
+	__atomic_sub_fetch(&current->inCriticalSection, 1, __ATOMIC_RELAXED);
+}
+
 void init_tasking(void);
 void schedule(void);
 pid_t clone(int flags);
diff --git a/kernel/start.S b/kernel/start.S
index 725d6e6..96048e2 100644
--- a/kernel/start.S
+++ b/kernel/start.S
@@ -2,10 +2,11 @@
 .global header
 header:
 	.long 0x1BADB002
-	.long 1 | 2
-	.long -(0x1BADB002 + (1 | 2))
+	.long 1 | 2 | 4
+	.long -(0x1BADB002 + (1 | 2 | 4))
 
 	.long 0, 0, 0, 0, 0
+	.long 0, 1280, 1024, 32
 
 .section .bss, "aw", @nobits
 .global stackTop
@@ -20,7 +21,6 @@ _start:
 	mov $stackTop, %ebp
 	mov %ebp, %esp
 	push %ebx
-	push %esp
 	call kmain
 	cli
 _end:
diff --git a/memory/fault.c b/memory/fault.c
index dd7953f..0ed452d 100644
--- a/memory/fault.c
+++ b/memory/fault.c
@@ -22,6 +22,7 @@ early_page_fault_handler(struct InterruptFrame *frame, uint32_t err)
 	if (!PAGE_ADDR(addr))
 		panic("Null dereference @ %#.8x", frame->eip);
 	ASSERT(!present);
+	kprintf("Allocating frame for %#.8x [%#.8x]", addr, frame->eip);
 	/* Allocate a page */
 	set_page(addr, alloc_frame() | PTE_PRESENT | PTE_WRITE | PTE_GLOBAL);
 }
diff --git a/task/clone.c b/task/clone.c
index 4dd15e7..22d737b 100644
--- a/task/clone.c
+++ b/task/clone.c
@@ -7,6 +7,7 @@
 #include <sys/sched.h>
 #include <nucleus/task.h>
 #include <nucleus/vfs.h>
+#include <io.h>
 
 extern ObjectList *readyQueue[];
 
@@ -14,7 +15,7 @@ extern ObjectList *readyQueue[];
 pid_t
 clone(int flags)
 {
-	uintptr_t ints = irq_disable();
+	enter_critical_section();
 
 	Task *parent = current, *child = new(&taskType), *tmp;
 	pid_t tid = 0;
@@ -41,7 +42,12 @@ clone(int flags)
 	add(readyQueue[child->priority], child);
 	tid = child->tid;
 	put(child);
+	exit_critical_section();
 end:
-	irq_restore(ints);
+	if (!tid) {
+		outb(0x20, 0x20);
+		LAPIC(0xB0) = 0;
+	}
+
 	return tid;
 }
diff --git a/task/scheduler.c b/task/scheduler.c
index 7b08fc2..1bc0f55 100644
--- a/task/scheduler.c
+++ b/task/scheduler.c
@@ -38,6 +38,8 @@ switch_to_task(Task *task)
 		   "g" (current->ebp), "g" (current->pageDir)
 	);
 end:
+	/* This prevents GCC from optimising the jump to be after the return */
+	__atomic_thread_fence(__ATOMIC_ACQ_REL);
 }
 
 /* Find the next schedulable ready queue */
@@ -56,6 +58,9 @@ highest_priority_queue(void)
 void
 schedule(void)
 {
+	if (current->inCriticalSection)
+		return;
+
 	Task *task = current;
 	ObjectList *queue = highest_priority_queue();