diff --git a/ksm.c b/ksm.c
index d5f2ce7..0f746e2 100644
--- a/ksm.c
+++ b/ksm.c
@@ -142,6 +142,11 @@ int __ksm_init_cpu(struct ksm *k)
 	}
 
 	vcpu = ksm_cpu(k);
+	if (vcpu->subverted) {
+		VCPU_DEBUG_RAW("CPU already subverted\n");
+		return 0;
+	}
+
 	ret = vcpu_create(vcpu);
 	if (ret < 0) {
 		VCPU_DEBUG_RAW("failed to create vcpu, oom?\n");
diff --git a/ksm/ksm/ksm.vcxproj b/ksm/ksm/ksm.vcxproj
index 0076a33..13743c3 100644
--- a/ksm/ksm/ksm.vcxproj
+++ b/ksm/ksm/ksm.vcxproj
@@ -21,8 +21,7 @@
   </PropertyGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.Default.props" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'" Label="Configuration">
-    <TargetVersion>
-    </TargetVersion>
+    <TargetVersion>Windows7</TargetVersion>
     <UseDebugLibraries>true</UseDebugLibraries>
     <PlatformToolset>WindowsKernelModeDriver10.0</PlatformToolset>
     <ConfigurationType>Driver</ConfigurationType>
diff --git a/main_linux.c b/main_linux.c
index c878919..aefebb9 100644
--- a/main_linux.c
+++ b/main_linux.c
@@ -25,6 +25,9 @@
 #include "um/um.h"
 
 static struct mm_struct *mm = NULL;
+static int major_no = 0;
+static struct class *class;
+
 static long ksm_ioctl(struct file *filp, unsigned int cmd, unsigned long args)
 {
 	int ret = -EINVAL;
@@ -103,8 +106,17 @@ static struct file_operations ksm_fops = {
 	.release = ksm_release,
 	.unlocked_ioctl = ksm_ioctl,
 };
-static int major_no = 0;
-static struct class *class;
+
+static int ksm_reboot(struct notifier_block *nb, unsigned long action,
+		      void *data)
+{
+	kms_exit(ksm);
+	return 0;
+}
+
+static struct notifier_block reboot_notify = {
+	.notifier_call = ksm_reboot,
+};
 
 static int __init ksm_start(void)
 {
@@ -126,6 +138,7 @@ static int __init ksm_start(void)
 
 	dev = device_create(class, NULL, MKDEV(major_no, 0), NULL, UM_DEVICE_NAME);
 	if (dev) {
+		register_reboot_notifier(&reboot_notify);
 		VCPU_DEBUG_RAW("ready\n");
 		return ret;
 	}
@@ -150,6 +163,7 @@ static void __exit ksm_cleanup(void)
 	class_unregister(class);
 	class_destroy(class);
 	unregister_chrdev(major_no, UM_DEVICE_NAME);
+	unregister_reboot_notifier(&reboot_notify);
 
 	active = ksm->active_vcpus;
 	ret = ksm_free(ksm);
diff --git a/main_nt.c b/main_nt.c
index ac0d01d..2e6bc79 100644
--- a/main_nt.c
+++ b/main_nt.c
@@ -109,6 +109,7 @@ static void DriverUnload(PDRIVER_OBJECT driverObject)
 	print_exit();
 #endif
 	IoDeleteSymbolicLink(&deviceLink);
+	IoUnregisterShutdownNotification(driverObject->DeviceObject);
 	IoDeleteDevice(driverObject->DeviceObject);
 }
 
@@ -120,7 +121,8 @@ static NTSTATUS DriverDispatch(PDEVICE_OBJECT deviceObject, PIRP irp)
 	u32 inlen = stack->Parameters.DeviceIoControl.InputBufferLength;
 	u32 ioctl;
 
-	if (stack->MajorFunction == IRP_MJ_DEVICE_CONTROL) {
+	switch (stack->MajorFunction) {
+	case IRP_MJ_DEVICE_CONTROL:	
 		ioctl = stack->Parameters.DeviceIoControl.IoControlCode;
 		VCPU_DEBUG("%s: IOCTL: 0x%08X\n", proc_name(), ioctl);
 
@@ -143,6 +145,11 @@ static NTSTATUS DriverDispatch(PDEVICE_OBJECT deviceObject, PIRP irp)
 			status = STATUS_NOT_SUPPORTED;
 			break;
 		}
+		break;
+	case IRP_MJ_SHUTDOWN:
+		/* Ignore return value  */
+		ksm_free(ksm);
+		break;
 	}
 
 	irp->IoStatus.Status = status;
@@ -187,6 +194,9 @@ NTSTATUS DriverEntry(PDRIVER_OBJECT driverObject, PUNICODE_STRING registryPath)
 	if (!NT_SUCCESS(status))
 		goto exit;
 
+	if (!NT_SUCCESS(IoRegisterShutdownNotification(deviceObject)))
+		goto err2;
+
 	driverObject->DriverUnload = DriverUnload;
 	driverObject->MajorFunction[IRP_MJ_CREATE] =
 		driverObject->MajorFunction[IRP_MJ_CLOSE] =
@@ -198,6 +208,8 @@ NTSTATUS DriverEntry(PDRIVER_OBJECT driverObject, PUNICODE_STRING registryPath)
 		goto out;
 	}
 
+	IoUnregisterShutdownNotification(deviceObject);
+err2:
 	IoDeleteDevice(deviceObject);
 exit:
 	ksm_free(ksm);
diff --git a/mm.h b/mm.h
index fb14b6d..9e3e012 100644
--- a/mm.h
+++ b/mm.h
@@ -245,6 +245,11 @@ static inline void __mm_free_pool(void *v)
 	kfree(v);
 }
 
+static inline bool mm_is_kernel_addr(void *va)
+{
+	return va >= PAGE_OFFSET;
+}
+
 extern void *mm_remap(u64 phys, size_t size);
 extern void mm_unmap(void *addr, size_t size);
 extern void *kmap_virt(void *addr, size_t len, pgprot_t prot);
@@ -377,6 +382,11 @@ static inline void __mm_free_pool(void *v)
 {
 	ExFreePool(v);
 }
+
+static inline bool mm_is_kernel_addr(void *va)
+{
+	return va >= MmSystemRangeStart;
+}
 #endif
 
 static inline void mm_free_pool(void *v, size_t size)
diff --git a/sandbox.c b/sandbox.c
index 74a4c5a..a256d04 100644
--- a/sandbox.c
+++ b/sandbox.c
@@ -284,8 +284,8 @@ bool ksm_sandbox_handle_ept(struct ept *ept, int dpl, u64 gpa,
 	pid = proc_id();
 	task = find_sa_task_pgd_pid(k, pid, cr3 & PAGE_PA_MASK);
 	if (!task) {
-		dbgbreak();
-		return false;
+		*eptp_switch = EPTP_DEFAULT;
+		return true;
 	}
 
 	eptp = task_eptp(task);
diff --git a/vcpu.c b/vcpu.c
index f2c72d3..9af4c13 100644
--- a/vcpu.c
+++ b/vcpu.c
@@ -151,9 +151,14 @@ static bool setup_pml4(struct ept *ept, int access, u16 eptp)
 
 	for (i = 0; i < ksm->range_count; ++i) {
 		range = &ksm->ranges[i];
-		for (addr = range->start; addr < range->end; addr += PAGE_SIZE)
-			if (!ept_alloc_page(EPT4(ept, eptp), access, addr, addr))
+		for (addr = range->start; addr < range->end; addr += PAGE_SIZE) {
+			int r = access;
+			if (mm_is_kernel_addr(__va(addr)))
+				r = EPT_ACCESS_ALL;
+
+			if (!ept_alloc_page(EPT4(ept, eptp), r, addr, addr))
 				return false;
+		}
 	}
 
 	/* Allocate APIC page  */