司徒的教學網站

驅動程式 - Windows NT Driver (Legacy) - 使用範例 - C/C++ (DDK) - Winsock Kernel(TCP) - Register Event

參考資訊：
https://github.com/wbenny/KSOCKET
https://learn.microsoft.com/en-us/windows-hardware/drivers/network/establishing-a-connection-with-a-destination
https://learn.microsoft.com/en-us/windows-hardware/drivers/network/roadmap-for-developing-network-drivers-with-winsock-kernel
https://community.osr.com/t/wsksocketconnect-how-to-safely-use-the-associated-irp-to-get-the-resulting-wsk-socket-object/53185

main.c

#include <ntddk.h>
#include <wsk.h>

WSK_REGISTRATION wreg;

static PIRP irp = NULL;
static KEVENT event = { 0 };
static PVOID pthread = NULL;
static WSK_CLIENT_NPI w = { 0 };
static WSK_PROVIDER_NPI npi = { 0 };
static PWSK_SOCKET socket = NULL;
static PWSK_PROVIDER_CONNECTION_DISPATCH sd = NULL;
static const WSK_CLIENT_DISPATCH wdispatch = {
    MAKE_WSK_VERSION(1, 0),
    0,
    NULL
};

#define SERVER_IP "192.168.229.128"
#define SERVER_PORT 9999

static NTSTATUS complete_routine(PDEVICE_OBJECT	DeviceObject, PIRP Irp, PKEVENT evt)
{
	KeSetEvent(evt, IO_NO_INCREMENT, FALSE);

	return STATUS_MORE_PROCESSING_REQUIRED;
}

static void reuse_irp(void)
{
	KeResetEvent(&event);
	IoReuseIrp(irp, STATUS_UNSUCCESSFUL);
	IoSetCompletionRoutine(irp, &complete_routine, &event, TRUE, TRUE, TRUE);
}

static NTSTATUS wsk_receive_event(PVOID SocketContext,
	ULONG Flags,
	PWSK_DATA_INDICATION DataIndication,
	SIZE_T BytesIndicated,
	SIZE_T* BytesAccepted)
{
	if (DataIndication == NULL) {
		return STATUS_SUCCESS;
	}

	PWSK_DATA_INDICATION currentIndication = DataIndication;
	SIZE_T totalBytesProcessed = 0;

	while (currentIndication != NULL) {
		SIZE_T length = currentIndication->Buffer.Length;
		PMDL mdl = currentIndication->Buffer.Mdl;

		if (length > 0 && mdl != NULL) {
			PVOID buffer = MmGetSystemAddressForMdlSafe(mdl, NormalPagePriority | MdlMappingNoExecute);

			if (buffer != NULL) {
				char* p = (char*)buffer + currentIndication->Buffer.Offset;

				DbgPrint("Received %.*s (%lu bytes)\n", (ULONG)length, p, (ULONG)length);
				totalBytesProcessed += length;
			}
			else {
				DbgPrint("Failed to map MDL system address safely.\n");
				break;
			}
		}
		currentIndication = currentIndication->Next;
	}

	*BytesAccepted = totalBytesProcessed;
	return STATUS_SUCCESS;
}

static NTSTATUS wsk_disconnect_event(PVOID SocketContext, ULONG Flags)
{
	DbgPrint("Disconnect event\n");
	return STATUS_SUCCESS;
}

const WSK_CLIENT_CONNECTION_DISPATCH cb = {
	wsk_receive_event,
	wsk_disconnect_event,
	NULL
};

static void wsk_thread(PVOID Context)
{
	NTSTATUS status = 0;

	do {
		w.ClientContext = NULL;
		w.Dispatch = &wdispatch;
		status = WskRegister(&w, &wreg);
		if (!NT_SUCCESS(status)) {
			break;
		}
		
		status = WskCaptureProviderNPI(&wreg, WSK_INFINITE_WAIT, &npi);
		if (!NT_SUCCESS(status)) {
			break;
		}
		KeInitializeEvent(&event, SynchronizationEvent, FALSE);

		irp = IoAllocateIrp(1, FALSE);
		if (irp == NULL) {
			status = STATUS_INSUFFICIENT_RESOURCES;
			break;
		}
		IoSetCompletionRoutine(irp, &complete_routine, &event, TRUE, TRUE, TRUE);

		SOCKADDR_IN local = { 0 };
		local.sin_family = AF_INET;
		local.sin_addr.s_addr = INADDR_ANY;
		local.sin_port = 0;

		PCWSTR term = NULL;
		SOCKADDR_IN remote = { 0 };
		remote.sin_family = AF_INET;
		remote.sin_port = RtlUshortByteSwap(SERVER_PORT);
		status = RtlIpv4StringToAddress(SERVER_IP, TRUE, &term, &remote.sin_addr);
		if (!NT_SUCCESS(status)) {
			break;
		}

		status = npi.Dispatch->WskSocketConnect(
			npi.Client,
			SOCK_STREAM,
			IPPROTO_TCP,
			(PSOCKADDR)&local,
			(PSOCKADDR)&remote,
			WSK_FLAG_CONNECTION_SOCKET,
			NULL,
			&cb,
			NULL,
			NULL,
			NULL,
			irp
		);

		if (status == STATUS_PENDING){
			KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
			status = irp->IoStatus.Status;
			socket = (PWSK_SOCKET)irp->IoStatus.Information;
			sd = socket->Dispatch;
		}

		if (!NT_SUCCESS(status)) {
			break;
		}
		DbgPrint("Connected to server !");


		WSK_EVENT_CALLBACK_CONTROL e;
		e.NpiId = &NPI_WSK_INTERFACE_ID;
		e.EventMask = WSK_EVENT_RECEIVE | WSK_EVENT_DISCONNECT;

		reuse_irp();
		status = sd->WskControlSocket(
			socket,
			WskSetOption,
			SO_WSK_EVENT_CALLBACK,
			SOL_SOCKET,
			sizeof(e),
			&e,
			0,
			NULL,
			NULL,
			irp);

		if (status == STATUS_PENDING) {
			KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
			status = irp->IoStatus.Status;
		}

		if (!NT_SUCCESS(status)) {
			break;
		}

		char buffer[] = { "Hello, world!" };
		WSK_BUF WskBuffer;
		WskBuffer.Offset = 0;
		WskBuffer.Length = strlen(buffer);
		WskBuffer.Mdl = IoAllocateMdl(buffer, (ULONG)WskBuffer.Length, FALSE, FALSE, NULL);

		__try {
			MmProbeAndLockPages(WskBuffer.Mdl, KernelMode, IoWriteAccess);
		}
		__except (EXCEPTION_EXECUTE_HANDLER) {
			status = STATUS_ACCESS_VIOLATION;
			break;
		}

		reuse_irp();
		status = sd->WskSend(socket, &WskBuffer, 0, irp);
		if (status == STATUS_PENDING) {
			KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
			status = irp->IoStatus.Status;
		}

		if (!NT_SUCCESS(status)) {
			break;
		}
		MmUnlockPages(WskBuffer.Mdl);
		DbgPrint("Sent data !");

		LARGE_INTEGER stTime;
		stTime.HighPart |= -1;
		stTime.LowPart = -30000000;
		KeDelayExecutionThread(KernelMode, FALSE, &stTime);

		reuse_irp(); 
		status = sd->WskCloseSocket(socket, irp);
		if (status == STATUS_PENDING) {
			KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
			status = irp->IoStatus.Status;
		}

		if (!NT_SUCCESS(status)) {
			break;
		}
		DbgPrint("Closed socket !");

	} while(0);
	
	if (irp) {
		IoFreeIrp(irp);
	}

	WskReleaseProviderNPI(&wreg);
	WskDeregister(&wreg);
	PsTerminateSystemThread(STATUS_SUCCESS);
}

void Unload(PDRIVER_OBJECT pMyDriver)
{
	if (pthread != NULL) {
		KeWaitForSingleObject(pthread, Executive, KernelMode, FALSE, NULL);
		ObDereferenceObject(pthread);
	}
}
  
NTSTATUS DriverEntry(PDRIVER_OBJECT pMyDriver, PUNICODE_STRING pMyRegistry)
{
    HANDLE thread = 0;
	NTSTATUS status = 0;

	status = PsCreateSystemThread(&thread,
		THREAD_ALL_ACCESS,
		NULL,
		NULL,
		NULL,
		wsk_thread,
		NULL
	);

    if (!NT_SUCCESS(status)) {
		DbgPrint("failed to create thread\n");
        return status;
    }
	
	status = ObReferenceObjectByHandle(thread, THREAD_ALL_ACCESS, NULL, KernelMode, &pthread, NULL);
    ZwClose(thread);
    if (!NT_SUCCESS(status)) {
		DbgPrint("failed to reference thread object\n");
		return status;
    }

    pMyDriver->DriverUnload = Unload;
	
    return STATUS_SUCCESS;
}

server.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
  
int main(int argc, char const* argv[])
{
    ssize_t r = 0;
    int opt = 1;
    int fd = -1;
    int new_socket = 0;
    char buf[255] = { 0 };
    struct sockaddr_in addr = { 0 };
    const char *hello = "hello from server !";
  
    if (argc != 3) {
        printf("Usage: %s IP Port\n", argv[0]);
        return 0;
    }
  
    socklen_t addrlen = sizeof(addr);
    fd = socket(AF_INET, SOCK_STREAM, 0);
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt));
  
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = inet_addr(argv[1]);
    addr.sin_port = htons(atoi(argv[2]));
  
    bind(fd, (struct sockaddr *)&addr, sizeof(addr));
    listen(fd, 3);
    new_socket = accept(fd, (struct sockaddr *)&addr, &addrlen);
    printf("new_socket %d\n", new_socket);
  
    r = read(new_socket, buf, sizeof(buf));
    printf("%s (r=%d)\n", buf, r);
    send(new_socket, hello, strlen(hello), 0);
    close(new_socket);
    close(fd);
 
    return 0;
}

Steward

分享是一種喜悅、更是一種幸福

驅動程式 - Windows NT Driver (Legacy) - 使用範例 - C/C++ (DDK) - Winsock Kernel(TCP) - Register Event