initial commit

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

7 files changed, 162 insertions, 153 deletions

diff --git a/README.md b/README.md
index 26f2bbd..da8c953 100644
--- a/README.md
+++ b/README.md
@@ -1,14 +1,37 @@
-# Mingw64 C++ Kernel Sockets
+# Mingw64 Kernel Sockets
 
-Requires [mingw-w64-dpp](https://github.com/utoni/mingw-w64-dpp).
+Requires [mingw-w64-dpp](https://github.com/utoni/mingw-w64-dpp) and Windows 10. Windows 11 may or may not work, not tested yet.
 
 Did you ever ask yourself the question: how-the-hell can I open a socket in kernel mode with a Mingw-w64 compiled driver?
 Well, this is the solution. `mingw-w64-ksocket` simplifies the use of stream/datagram client/server sockets.
 The API is similiar to the BSD socket API.
 Greetings to [KSOCKET](https://github.com/wbenny/KSOCKET) for this sweet approach.
 
+Compatible with C and C++ drivers.
+
 # HowTo
 
+Build it:
+
+```shell
+make DPP_ROOT="[path-to-mingw-w64-dpp-template-dir]" all
+
+```
+
+To install & (self-)sign the driver:
+
 ```shell
-make -C . DPP_ROOT="[path-to-mingw-w64-ddk-template-dir]" clean all
+make DPP_ROOT="[path-to-mingw-w64-dpp-template-dir]" DESTDIR="[path-to-install-dir]" install
 ```
+
+The directory `[path-to-install-dir]` should now contain three new files:
+
+ * `driver.bat`: setup the driver service, start it, stop it when it's done and delete it
+ * `driver.sys`: example driver that uses kernel sockets
+ * `userspace_client.exe`: example userspace application which communicates with the driver via TCP socket
+
+Start `driver.bat` as `Administrator` and then `userspace_client.exe`.
+
+If everything works fine, there should be a text displayed in `userspace_client.exe` console window, received from the driver.
+
+For more debug output, it is recommended to use a debugger or log viewer like `dbgview`.
diff --git a/berkeley.h b/berkeley.h
index 75cdb0f..0d09c77 100644
--- a/berkeley.h
+++ b/berkeley.h
@@ -1,38 +1,41 @@
 #pragma once
 #include <ntddk.h>
-#include <wsk.h>
 #include <stdint.h>
+#include <wsk.h>
+
+#define socket socket_connection
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-typedef int       socklen_t;
-typedef intptr_t  ssize_t;
+typedef int socklen_t;
+typedef intptr_t ssize_t;
 
 uint32_t htonl(uint32_t hostlong);
 uint16_t htons(uint16_t hostshort);
 uint32_t ntohl(uint32_t netlong);
 uint16_t ntohs(uint16_t netshort);
 
-int getaddrinfo(const char* node, const char* service, const struct addrinfo* hints, struct addrinfo** res);
+int getaddrinfo(const char *node, const char *service,
+                const struct addrinfo *hints, struct addrinfo **res);
 void freeaddrinfo(struct addrinfo *res);
 
 int socket_connection(int domain, int type, int protocol);
 int socket_listen(int domain, int type, int protocol);
 int socket_datagram(int domain, int type, int protocol);
-int connect(int sockfd, const struct sockaddr* addr, socklen_t addrlen);
+int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen);
 int listen(int sockfd, int backlog);
 int bind(int sockfd, const struct sockaddr *addr, socklen_t addrlen);
 int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen);
-int send(int sockfd, const void* buf, size_t len, int flags);
-int sendto(int sockfd, const void *buf, size_t len, int flags, const struct sockaddr *dest_addr, socklen_t addrlen);
-int recv(int sockfd, void* buf, size_t len, int flags);
-int recvfrom(int sockfd, void *buf, size_t len, int flags, struct sockaddr *src_addr, socklen_t *addrlen);
+int send(int sockfd, const void *buf, size_t len, int flags);
+int sendto(int sockfd, const void *buf, size_t len, int flags,
+           const struct sockaddr *dest_addr, socklen_t addrlen);
+int recv(int sockfd, void *buf, size_t len, int flags);
+int recvfrom(int sockfd, void *buf, size_t len, int flags,
+             struct sockaddr *src_addr, socklen_t *addrlen);
 int closesocket(int sockfd);
 
-#define socket  socket_connection
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/driver.cpp b/driver.cpp
index 005468d..b7ecbc1 100644
--- a/driver.cpp
+++ b/driver.cpp
@@ -55,12 +55,30 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject,
 
     int sockfd;
     sockfd = socket_connection(AF_INET, SOCK_STREAM, IPPROTO_TCP);
+    if (sockfd < 0) {
+      DebuggerPrint("TCP client socket_connection failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    }
+
     result = connect(sockfd, res->ai_addr, (int)res->ai_addrlen);
+    if (result != 0) {
+      DebuggerPrint("TCP client connect failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    }
+
     result = send(sockfd, send_buffer, sizeof(send_buffer), 0);
-    result = recv(sockfd, recv_buffer, sizeof(recv_buffer), 0);
-    recv_buffer[sizeof(recv_buffer) - 1] = '\0';
+    if (result <= 0) {
+      DebuggerPrint("TCP client send failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    }
 
-    DebuggerPrint("TCP client:\n%s\n", recv_buffer);
+    result = recv(sockfd, recv_buffer, sizeof(recv_buffer), 0);
+    if (result <= 0) {
+      DebuggerPrint("TCP client recv failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    } else {
+      DebuggerPrint("TCP client:\n%.*s\n", result, recv_buffer);
+    }
 
     closesocket(sockfd);
   }
@@ -73,7 +91,6 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject,
 
   {
     int result;
-    UNREFERENCED_PARAMETER(result);
 
     char send_buffer[] = "Hello from WSK!";
     char recv_buffer[1024] = {0};
@@ -86,19 +103,34 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject,
     addr.sin_port = htons(9095);
 
     result = bind(server_sockfd, (struct sockaddr *)&addr, sizeof(addr));
+    if (result != 0) {
+      DebuggerPrint("TCP server bind failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    }
+
     result = listen(server_sockfd, 1);
+    if (result != 0) {
+      DebuggerPrint("TCP server listen failed\n");
+      return STATUS_FAILED_DRIVER_ENTRY;
+    }
 
     socklen_t addrlen = sizeof(addr);
     int client_sockfd =
         accept(server_sockfd, (struct sockaddr *)&addr, &addrlen);
 
     result = recv(client_sockfd, recv_buffer, sizeof(recv_buffer) - 1, 0);
-    recv_buffer[sizeof(recv_buffer) - 1] = '\0';
-
-    DebuggerPrint("TCP server:\n%s\n", recv_buffer);
+    if (result > 0) {
+      DebuggerPrint("TCP server:\n%.*s\n", result, recv_buffer);
+    } else {
+      DebuggerPrint("TCP server recv failed\n");
+    }
 
     result = send(client_sockfd, send_buffer, sizeof(send_buffer), 0);
 
+    // Wait for the client to terminate the connection.
+    do {
+    } while (recv(client_sockfd, recv_buffer, sizeof(recv_buffer) - 1, 0) > 0);
+
     closesocket(client_sockfd);
     closesocket(server_sockfd);
   }
diff --git a/ksocket.c b/ksocket.c
index 5262333..3b6ba4d 100644
--- a/ksocket.c
+++ b/ksocket.c
@@ -172,8 +172,21 @@ KspAsyncContextWaitForCompletion(_In_ PKSOCKET_ASYNC_CONTEXT AsyncContext,
 NTSTATUS
 NTAPI
 KsInitialize(VOID) {
+  OSVERSIONINFOEXW OsVersionInfo;
   NTSTATUS Status;
 
+  OsVersionInfo.dwOSVersionInfoSize = sizeof(OsVersionInfo);
+  Status = RtlGetVersion((POSVERSIONINFOW)&OsVersionInfo);
+
+  if (!NT_SUCCESS(Status)) {
+    return Status;
+  }
+
+  if (OsVersionInfo.wProductType != VER_NT_WORKSTATION ||
+      OsVersionInfo.dwMajorVersion != 10 || OsVersionInfo.dwMinorVersion != 0) {
+    return STATUS_UNSUPPORTED_WINDOWS_VERSION;
+  }
+
   //
   // Register as a WSK client.
   //
diff --git a/ksocket.h b/ksocket.h
index a4bd0f7..0a6717e 100644
--- a/ksocket.h
+++ b/ksocket.h
@@ -1,6 +1,9 @@
 #pragma once
-#include <ntddk.h>
 #include "wsk.h"
+#include <ntddk.h>
+
+#define STATUS_UNSUPPORTED_WINDOWS_VERSION                                     \
+  (STATUS_SEVERITY_ERROR | (1 << 29) | (1 << 16) | 1)
 
 #ifdef __cplusplus
 extern "C" {
@@ -10,155 +13,86 @@ typedef struct _KSOCKET KSOCKET, *PKSOCKET;
 
 NTSTATUS
 NTAPI
-KsInitialize(
-  VOID
-  );
+KsInitialize(VOID);
 
-VOID
-NTAPI
-KsDestroy(
-  VOID
-  );
+VOID NTAPI KsDestroy(VOID);
 
 NTSTATUS
 NTAPI
-KsGetAddrInfo(
-  _In_ PUNICODE_STRING NodeName,
-  _In_ PUNICODE_STRING ServiceName,
-  _In_ PADDRINFOEXW Hints,
-  _Out_ PADDRINFOEXW* Result
-  );
+KsGetAddrInfo(_In_ PUNICODE_STRING NodeName, _In_ PUNICODE_STRING ServiceName,
+              _In_ PADDRINFOEXW Hints, _Out_ PADDRINFOEXW *Result);
 
-VOID
-NTAPI
-KsFreeAddrInfo(
-  _In_ PADDRINFOEXW AddrInfo
-  );
+VOID NTAPI KsFreeAddrInfo(_In_ PADDRINFOEXW AddrInfo);
 
 NTSTATUS
 NTAPI
-KsCreateSocket(
-  _Out_ PKSOCKET* Socket,
-  _In_ ADDRESS_FAMILY AddressFamily,
-  _In_ USHORT SocketType,
-  _In_ ULONG Protocol,
-  _In_ ULONG Flags
-  );
+KsCreateSocket(_Out_ PKSOCKET *Socket, _In_ ADDRESS_FAMILY AddressFamily,
+               _In_ USHORT SocketType, _In_ ULONG Protocol, _In_ ULONG Flags);
 
 NTSTATUS
 NTAPI
-KsCreateConnectionSocket(
-  _Out_ PKSOCKET* Socket,
-  _In_ ADDRESS_FAMILY AddressFamily,
-  _In_ USHORT SocketType,
-  _In_ ULONG Protocol
-  );
+KsCreateConnectionSocket(_Out_ PKSOCKET *Socket,
+                         _In_ ADDRESS_FAMILY AddressFamily,
+                         _In_ USHORT SocketType, _In_ ULONG Protocol);
 
 NTSTATUS
 NTAPI
-KsCreateListenSocket(
-  _Out_ PKSOCKET* Socket,
-  _In_ ADDRESS_FAMILY AddressFamily,
-  _In_ USHORT SocketType,
-  _In_ ULONG Protocol
-  );
+KsCreateListenSocket(_Out_ PKSOCKET *Socket, _In_ ADDRESS_FAMILY AddressFamily,
+                     _In_ USHORT SocketType, _In_ ULONG Protocol);
 
 NTSTATUS
 NTAPI
-KsCreateDatagramSocket(
-  _Out_ PKSOCKET* Socket,
-  _In_ ADDRESS_FAMILY AddressFamily,
-  _In_ USHORT SocketType,
-  _In_ ULONG Protocol
-  );
+KsCreateDatagramSocket(_Out_ PKSOCKET *Socket,
+                       _In_ ADDRESS_FAMILY AddressFamily,
+                       _In_ USHORT SocketType, _In_ ULONG Protocol);
 
 NTSTATUS
 NTAPI
-KsCloseSocket(
-  _In_ PKSOCKET Socket
-  );
+KsCloseSocket(_In_ PKSOCKET Socket);
 
 NTSTATUS
 NTAPI
-KsBind(
-  _In_ PKSOCKET Socket,
-  _In_ PSOCKADDR LocalAddress
-  );
+KsBind(_In_ PKSOCKET Socket, _In_ PSOCKADDR LocalAddress);
 
 NTSTATUS
 NTAPI
-KsAccept(
-  _In_ PKSOCKET Socket,
-  _Out_ PKSOCKET* NewSocket,
-  _Out_opt_ PSOCKADDR LocalAddress,
-  _Out_opt_ PSOCKADDR RemoteAddress
-  );
+KsAccept(_In_ PKSOCKET Socket, _Out_ PKSOCKET *NewSocket,
+         _Out_opt_ PSOCKADDR LocalAddress, _Out_opt_ PSOCKADDR RemoteAddress);
 
 NTSTATUS
 NTAPI
-KsConnect(
-  _In_ PKSOCKET Socket,
-  _In_ PSOCKADDR RemoteAddress
-  );
+KsConnect(_In_ PKSOCKET Socket, _In_ PSOCKADDR RemoteAddress);
 
 NTSTATUS
 NTAPI
-KsSendRecv(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags,
-  _In_ BOOLEAN Send
-  );
+KsSendRecv(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+           _In_ ULONG Flags, _In_ BOOLEAN Send);
 
 NTSTATUS
 NTAPI
-KsSendRecvUdp(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags,
-  _In_ PSOCKADDR RemoteAddress,
-  _In_ BOOLEAN Send
-  );
+KsSendRecvUdp(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+              _In_ ULONG Flags, _In_ PSOCKADDR RemoteAddress,
+              _In_ BOOLEAN Send);
 
 NTSTATUS
 NTAPI
-KsSend(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags
-  );
+KsSend(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+       _In_ ULONG Flags);
 
 NTSTATUS
 NTAPI
-KsRecv(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags
-  );
+KsRecv(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+       _In_ ULONG Flags);
 
 NTSTATUS
 NTAPI
-KsSendTo(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags,
-  _In_ PSOCKADDR RemoteAddress
-  );
+KsSendTo(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+         _In_ ULONG Flags, _In_ PSOCKADDR RemoteAddress);
 
 NTSTATUS
 NTAPI
-KsRecvFrom(
-  _In_ PKSOCKET Socket,
-  _In_ PVOID Buffer,
-  _Inout_ PULONG Length,
-  _In_ ULONG Flags,
-  _In_ PSOCKADDR RemoteAddress
-  );
+KsRecvFrom(_In_ PKSOCKET Socket, _In_ PVOID Buffer, _Inout_ PULONG Length,
+           _In_ ULONG Flags, _In_ PSOCKADDR RemoteAddress);
 
 #ifdef __cplusplus
 }
diff --git a/userspace_client.cpp b/userspace_client.cpp
index efcd782..cb2a3b4 100644
--- a/userspace_client.cpp
+++ b/userspace_client.cpp
@@ -47,7 +47,7 @@ int main(int argc, char **argv) {
     return 1;
   }
 
-  do {
+  {
     char recvbuf[1024];
 
     iResult = recv(sock, recvbuf, 1024, 0);
@@ -55,12 +55,12 @@ int main(int argc, char **argv) {
       wprintf(L"Bytes received: %d\n", iResult);
       wprintf(L"Data received: %.*s\n", iResult, recvbuf);
     } else if (iResult == 0) {
-      wprintf(L"Connection closed\n");
+      wprintf(L"Connection closed by remote\n");
     } else if (WSAGetLastError() != WSAECONNRESET) {
       wprintf(L"recv failed: %d\n", WSAGetLastError());
     }
-  } while (iResult > 0);
-  wprintf(L"connection closed by remote\n");
+  };
+  wprintf(L"Closing Connection ..\n");
 
   iResult = closesocket(sock);
   if (iResult == SOCKET_ERROR) {
diff --git a/wsk.h b/wsk.h
index f7e5fb6..ccea103 100644
--- a/wsk.h
+++ b/wsk.h
@@ -1,8 +1,11 @@
 #ifndef WSK_H
 #define WSK_H 1
 
+#include <ntddk.h>
+
 #if !defined(__MINGW64__)
-#error "This was designed to work **only** with Mingw-w64! For MSVC, please check https://github.com/wbenny/KSOCKET out"
+#error                                                                         \
+    "This was designed to work **only** with Mingw-w64! For MSVC, please check https://github.com/wbenny/KSOCKET out"
 #endif
 
 #if (NTDDI_VERSION < NTDDI_WIN10)
@@ -11,7 +14,8 @@
 #endif
 
 #if !defined(__BYTE_ORDER__) || __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
-#error "This project requires a little endian system. Does Windows support any other?"
+#error                                                                         \
+    "This project requires a little endian system. Does Windows support any other?"
 #endif
 
 // ---------------------------------------------------------------
@@ -110,38 +114,38 @@ typedef enum {
 // ---------------------------------------------------------------
 
 typedef struct addrinfo {
-  int             ai_flags;
-  int             ai_family;
-  int             ai_socktype;
-  int             ai_protocol;
-  size_t          ai_addrlen;
-  char            *ai_canonname;
+  int ai_flags;
+  int ai_family;
+  int ai_socktype;
+  int ai_protocol;
+  size_t ai_addrlen;
+  char *ai_canonname;
   struct sockaddr *ai_addr;
   struct addrinfo *ai_next;
 } ADDRINFOA, *PADDRINFOA;
 
 typedef struct addrinfoexA {
-  int                ai_flags;
-  int                ai_family;
-  int                ai_socktype;
-  int                ai_protocol;
-  size_t             ai_addrlen;
-  char               *ai_canonname;
-  struct sockaddr    *ai_addr;
-  void               *ai_blob;
-  size_t             ai_bloblen;
-  LPGUID             ai_provider;
+  int ai_flags;
+  int ai_family;
+  int ai_socktype;
+  int ai_protocol;
+  size_t ai_addrlen;
+  char *ai_canonname;
+  struct sockaddr *ai_addr;
+  void *ai_blob;
+  size_t ai_bloblen;
+  LPGUID ai_provider;
   struct addrinfoexA *ai_next;
 } ADDRINFOEXA, *PADDRINFOEXA, *LPADDRINFOEXA;
 
 typedef struct addrinfoW {
-  int              ai_flags;
-  int              ai_family;
-  int              ai_socktype;
-  int              ai_protocol;
-  size_t           ai_addrlen;
-  PWSTR            ai_canonname;
-  struct sockaddr  *ai_addr;
+  int ai_flags;
+  int ai_family;
+  int ai_socktype;
+  int ai_protocol;
+  size_t ai_addrlen;
+  PWSTR ai_canonname;
+  struct sockaddr *ai_addr;
   struct addrinfoW *ai_next;
 } ADDRINFOW, *PADDRINFOW;