Cleanup hook irp

src/driver/hypervisor.cpp

@@ -61,6 +61,7 @@ namespace
 	}
 
         // This absolutely needs to be inlined. Otherwise the stack might be broken upon restoration
+        // See: https://github.com/ionescu007/SimpleVisor/issues/48
         #define capture_cpu_context(launch_context) \
 	      cpature_special_registers((launch_context).special_registers);\
 	      RtlCaptureContext(&(launch_context).context_frame);

src/driver/irp.cpp

@@ -36,11 +36,10 @@ namespace
 
 		return STATUS_SUCCESS;
 	}
-
+	
-	// TODO: This is vulnerable as fuck. Optimize!
+	void apply_hook(const hook_request& request)
-	void apply_hook(const hook_request* request)
 	{
-		auto* buffer = new uint8_t[request->source_data_size];
+		auto* buffer = new uint8_t[request.source_data_size];
 		if (!buffer)
 		{
 			throw std::runtime_error("Failed to copy buffer");
@@ -53,9 +52,15 @@ namespace
 			vmx::ept::free_translation_hints(translation_hints);
 		});
 
-		memcpy(buffer, request->source_data, request->source_data_size);
+		memcpy(buffer, request.source_data, request.source_data_size);
 
-		thread::kernel_thread t([&translation_hints, r = *request]()
+		auto* hypervisor = hypervisor::get_instance();
+	        if(!hypervisor)
+		{
+		  throw std::runtime_error("Hypervisor not installed");
+	        }
+
+		thread::kernel_thread t([&translation_hints, r = request]
 		{
 			debug_log("Pid: %d | Address: %p\n", r.process_id, r.target_address);
 
@@ -69,11 +74,9 @@ namespace
 			const auto name = process_handle.get_image_filename();
 			if (name)
 			{
-				debug_log("Attaching to %s\n", name);
+			  debug_log("Attaching to %s\n", name);
 			}
 
-			debug_log("Level: %d\n", static_cast<int>(KeGetCurrentIrql()));
-
 			process::scoped_process_attacher attacher{process_handle};
 			translation_hints = vmx::ept::generate_translation_hints(r.target_address, r.source_data_size);
 		});
@@ -82,14 +85,12 @@ namespace
 
 		if (!translation_hints)
 		{
-			debug_log("Failed to generate tranlsation hints");
+			debug_log("Failed to generate tranlsation hints\n");
 			return;
 		}
 
-		hypervisor::get_instance()->install_ept_hook(request->target_address, buffer, request->source_data_size,
+		hypervisor->install_ept_hook(request.target_address, buffer, request.source_data_size,
 		                                             translation_hints);
-
-		debug_log("Done1\n");
 	}
 
 	void unhook()
@@ -101,37 +102,69 @@ namespace
 		}
 	}
 
+        void try_apply_hook(const PIO_STACK_LOCATION irp_sp)
+	{
+	  if(irp_sp->Parameters.DeviceIoControl.InputBufferLength < sizeof(hook_request))
+	  {
+	    throw std::runtime_error("Invalid hook request");
+	  }
+
+          const auto& request = *static_cast<hook_request*>(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer);
+	  memory::assert_readability(request.source_data, request.source_data_size);
+	  memory::assert_readability(request.target_address, request.source_data_size);
+
+	  apply_hook(request);
+	}
+
+        void handle_irp(const PIRP irp)
+	{
+	  irp->IoStatus.Information = 0;
+	  irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
+
+	  const auto irp_sp = IoGetCurrentIrpStackLocation(irp);
+
+	  if (irp_sp)
+	  {
+	    const auto ioctr_code = irp_sp->Parameters.DeviceIoControl.IoControlCode;
+
+	    switch (ioctr_code)
+	    {
+	    case HELLO_DRV_IOCTL:
+	      debug_log("Hello from the Driver!\n");
+	      break;
+	    case HOOK_DRV_IOCTL:
+	      try_apply_hook(irp_sp);
+	      break;
+	    case UNHOOK_DRV_IOCTL:
+	      unhook();
+	      break;
+	    default:
+	      debug_log("Invalid IOCTL Code: 0x%X\n", ioctr_code);
+	      irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
+	      break;
+	    }
+	  }
+	}
+
 	_Function_class_(DRIVER_DISPATCH) NTSTATUS io_ctl_handler(
 		PDEVICE_OBJECT /*device_object*/, const PIRP irp)
 	{
 		PAGED_CODE()
 
-		irp->IoStatus.Information = 0;
+	        try
-		irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
-
-		const auto irp_sp = IoGetCurrentIrpStackLocation(irp);
-
-		if (irp_sp)
 		{
-			const auto ioctr_code = irp_sp->Parameters.DeviceIoControl.IoControlCode;
+		  handle_irp(irp);
-
+	        }
-			switch (ioctr_code)
+	        catch(std::exception& e)
-			{
+		{
-			case HELLO_DRV_IOCTL:
+		  debug_log("Handling IRP failed: %s\n", e.what());
-				debug_log("Hello from the Driver!\n");
+		  irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
-				break;
+	        }
-			case HOOK_DRV_IOCTL:
+	        catch(...)
-				apply_hook(static_cast<hook_request*>(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer));
+		{
-				break;
+		  debug_log("Handling IRP failed\n");
-			case UNHOOK_DRV_IOCTL:
+		  irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
-				unhook();
+	        }
-				break;
-			default:
-				debug_log("Invalid IOCTL Code: 0x%X\n", ioctr_code);
-				irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
-				break;
-			}
-		}
 
 		IoCompleteRequest(irp, IO_NO_INCREMENT);
 

src/driver/memory.cpp

@@ -104,4 +104,24 @@ namespace memory
 			ExFreePool(memory);
 		}
 	}
+
+	bool prope_for_read(const void* address, const size_t length, const uint64_t alignment)
+	{
+	  __try
+	  {
+	    ProbeForRead(const_cast<volatile void*>(address), length, static_cast<ULONG>(alignment));
+	    return true;
+	  }
+	  __except (EXCEPTION_EXECUTE_HANDLER)
+	  {
+	    return false;
+	  }
+	}
+
+	void assert_readability(const void* address, const size_t length, const uint64_t alignment)
+	{
+	  if(!prope_for_read(address, length, alignment)) {
+	    throw std::runtime_error("Access violation");
+	  }
+	}
 }

src/driver/memory.hpp

@@ -23,6 +23,9 @@ namespace memory
 	_IRQL_requires_max_(DISPATCH_LEVEL)
 	void free_non_paged_memory(void* memory);
 
+	bool prope_for_read(const void* address, size_t length, uint64_t alignment = 1);
+	void assert_readability(const void* address, size_t length, uint64_t alignment = 1);
+
 	template <typename T, typename... Args>
 	T* allocate_non_paged_object(Args ... args)
 	{

src/runner/main.cpp

@@ -10,30 +10,6 @@
 
 #pragma comment(lib, "Shlwapi.lib")
 
-BOOL send_ioctl(HANDLE device, DWORD ioctl_code)
-{
-	//prepare input buffer:
-	DWORD bufSize = 0x4;
-	BYTE* inBuffer = (BYTE*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, bufSize);
-
-	//fill the buffer with some content:
-	RtlFillMemory(inBuffer, bufSize, 'A');
-
-	DWORD size_returned = 0;
-	BOOL is_ok = DeviceIoControl(device,
-	                             ioctl_code,
-	                             inBuffer,
-	                             bufSize,
-	                             NULL, //outBuffer -> None
-	                             0, //outBuffer size -> 0
-	                             &size_returned,
-	                             NULL
-	);
-	//release the input bufffer:
-	HeapFree(GetProcessHeap(), 0, (LPVOID)inBuffer);
-	return is_ok;
-}
-
 std::filesystem::path get_current_path()
 {
 	const auto module = GetModuleHandleA(nullptr);