api: upgrade ring management

- Return pointer to ring buffer with packet data allowing clients to
  read/write directly. This eliminates one memcpy().

- Make sending/receiving packets thread-safe.

Signed-off-by: Simon Rozman <simon@rozman.si>

3 files changed, 187 insertions, 108 deletions

diff --git a/api/exports.def b/api/exports.def
index 85e4453..6dbd89a 100644
--- a/api/exports.def
+++ b/api/exports.def
@@ -1,4 +1,5 @@
 EXPORTS
+	WintunAllocateSendPacket
 	WintunCreateAdapter
 	WintunDeleteAdapter
 	WintunEndSession
@@ -11,8 +12,9 @@ EXPORTS
 	WintunGetAdapterName
 	WintunGetVersion
 	WintunIsPacketAvailable
-	WintunReceivePackets
-	WintunSendPackets
+	WintunReceivePacket
+	WintunReceiveRelease
+	WintunSendPacket
 	WintunSetAdapterName
 	WintunSetLogger
 	WintunStartSession
diff --git a/api/session.c b/api/session.c
index f58b69c..02aa5be 100644
--- a/api/session.c
+++ b/api/session.c
@@ -14,12 +14,13 @@
 #define TUN_RING_CAPACITY(Size) ((Size) - sizeof(TUN_RING) - (TUN_MAX_PACKET_SIZE - TUN_ALIGNMENT))
 #define TUN_RING_SIZE(Capacity) (sizeof(TUN_RING) + (Capacity) + (TUN_MAX_PACKET_SIZE - TUN_ALIGNMENT))
 #define TUN_RING_WRAP(Value, Capacity) ((Value) & (Capacity - 1))
+#define LOCK_SPIN_COUNT 0x10000
+#define TUN_PACKET_RELEASE ((DWORD)0x80000000)
 
 typedef struct _TUN_PACKET
 {
     ULONG Size;
-    UCHAR _Field_size_bytes_(Size)
-    Data[];
+    UCHAR Data[];
 } TUN_PACKET;
 
 typedef struct _TUN_RING
@@ -45,6 +46,20 @@ typedef struct _TUN_REGISTER_RINGS
 typedef struct _TUN_SESSION
 {
     ULONG Capacity;
+    struct
+    {
+        ULONG Tail;
+        ULONG TailRelease;
+        ULONG PacketsToRelease;
+        CRITICAL_SECTION Lock;
+    } Receive;
+    struct
+    {
+        ULONG Head;
+        ULONG HeadRelease;
+        ULONG PacketsToRelease;
+        CRITICAL_SECTION Lock;
+    } Send;
     TUN_REGISTER_RINGS Descriptor;
     HANDLE Handle;
 } TUN_SESSION;
@@ -52,7 +67,7 @@ typedef struct _TUN_SESSION
 WINTUN_STATUS WINAPI
 WintunStartSession(_In_ const WINTUN_ADAPTER *Adapter, _In_ DWORD Capacity, _Out_ TUN_SESSION **Session)
 {
-    *Session = HeapAlloc(ModuleHeap, 0, sizeof(TUN_SESSION));
+    *Session = HeapAlloc(ModuleHeap, HEAP_ZERO_MEMORY, sizeof(TUN_SESSION));
     if (!*Session)
         return LOG(WINTUN_LOG_ERR, L"Out of memory"), ERROR_OUTOFMEMORY;
     const ULONG RingSize = TUN_RING_SIZE(Capacity);
@@ -102,6 +117,8 @@ WintunStartSession(_In_ const WINTUN_ADAPTER *Adapter, _In_ DWORD Capacity, _Out
         goto cleanupHandle;
     }
     (*Session)->Capacity = Capacity;
+    (void)InitializeCriticalSectionAndSpinCount(&(*Session)->Receive.Lock, LOCK_SPIN_COUNT);
+    (void)InitializeCriticalSectionAndSpinCount(&(*Session)->Send.Lock, LOCK_SPIN_COUNT);
     return ERROR_SUCCESS;
 cleanupHandle:
     CloseHandle((*Session)->Handle);
@@ -121,6 +138,8 @@ void WINAPI
 WintunEndSession(_In_ TUN_SESSION *Session)
 {
     SetEvent(Session->Descriptor.Send.TailMoved); // wake the reader if it's sleeping
+    DeleteCriticalSection(&Session->Send.Lock);
+    DeleteCriticalSection(&Session->Receive.Lock);
     CloseHandle(Session->Handle);
     CloseHandle(Session->Descriptor.Send.TailMoved);
     CloseHandle(Session->Descriptor.Receive.TailMoved);
@@ -131,8 +150,7 @@ WintunEndSession(_In_ TUN_SESSION *Session)
 BOOL WINAPI
 WintunIsPacketAvailable(_In_ TUN_SESSION *Session)
 {
-    return InterlockedGetU(&Session->Descriptor.Send.Ring->Head) !=
-           InterlockedGetU(&Session->Descriptor.Send.Ring->Tail);
+    return Session->Send.Head != InterlockedGetU(&Session->Descriptor.Send.Ring->Tail);
 }
 
 WINTUN_STATUS WINAPI
@@ -142,71 +160,126 @@ WintunWaitForPacket(_In_ TUN_SESSION *Session, _In_ DWORD Milliseconds)
 }
 
 WINTUN_STATUS WINAPI
-WintunReceivePackets(_In_ TUN_SESSION *Session, _Inout_ WINTUN_PACKET *Queue)
+WintunReceivePacket(_In_ TUN_SESSION *Session, _Out_bytecapcount_(*PacketSize) BYTE **Packet, _Out_ DWORD *PacketSize)
 {
-    ULONG BuffHead = InterlockedGetU(&Session->Descriptor.Send.Ring->Head);
-    if (BuffHead >= Session->Capacity)
-        return ERROR_HANDLE_EOF;
-
-    for (; Queue; Queue = Queue->Next)
+    DWORD Result;
+    EnterCriticalSection(&Session->Send.Lock);
+    if (Session->Send.Head >= Session->Capacity)
     {
-        const ULONG BuffTail = InterlockedGetU(&Session->Descriptor.Send.Ring->Tail);
-        if (BuffTail >= Session->Capacity)
-            return ERROR_HANDLE_EOF;
-
-        if (BuffHead == BuffTail)
-            return ERROR_NO_MORE_ITEMS;
-
-        const ULONG BuffContent = TUN_RING_WRAP(BuffTail - BuffHead, Session->Capacity);
-        if (BuffContent < sizeof(TUN_PACKET))
-            return ERROR_INVALID_DATA;
-
-        const TUN_PACKET *Packet = (TUN_PACKET *)&Session->Descriptor.Send.Ring->Data[BuffHead];
-        if (Packet->Size > WINTUN_MAX_IP_PACKET_SIZE)
-            return ERROR_INVALID_DATA;
-
-        const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + Packet->Size);
-        if (AlignedPacketSize > BuffContent)
-            return ERROR_INVALID_DATA;
-
-        Queue->Size = Packet->Size;
-        memcpy(Queue->Data, Packet->Data, Packet->Size);
-        BuffHead = TUN_RING_WRAP(BuffHead + AlignedPacketSize, Session->Capacity);
-        InterlockedSetU(&Session->Descriptor.Send.Ring->Head, BuffHead);
+        Result = ERROR_HANDLE_EOF;
+        goto cleanup;
     }
+    const ULONG BuffTail = InterlockedGetU(&Session->Descriptor.Send.Ring->Tail);
+    if (BuffTail >= Session->Capacity)
+    {
+        Result = ERROR_HANDLE_EOF;
+        goto cleanup;
+    }
+    if (Session->Send.Head == BuffTail)
+    {
+        Result = ERROR_NO_MORE_ITEMS;
+        goto cleanup;
+    }
+    const ULONG BuffContent = TUN_RING_WRAP(BuffTail - Session->Send.Head, Session->Capacity);
+    if (BuffContent < sizeof(TUN_PACKET))
+    {
+        Result = ERROR_INVALID_DATA;
+        goto cleanup;
+    }
+    TUN_PACKET *BuffPacket = (TUN_PACKET *)&Session->Descriptor.Send.Ring->Data[Session->Send.Head];
+    if (BuffPacket->Size > WINTUN_MAX_IP_PACKET_SIZE)
+    {
+        Result = ERROR_INVALID_DATA;
+        goto cleanup;
+    }
+    const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + BuffPacket->Size);
+    if (AlignedPacketSize > BuffContent)
+    {
+        Result = ERROR_INVALID_DATA;
+        goto cleanup;
+    }
+    *PacketSize = BuffPacket->Size;
+    *Packet = BuffPacket->Data;
+    Session->Send.Head = TUN_RING_WRAP(Session->Send.Head + AlignedPacketSize, Session->Capacity);
+    Session->Send.PacketsToRelease++;
+    Result = ERROR_SUCCESS;
+cleanup:
+    LeaveCriticalSection(&Session->Send.Lock);
+    return Result;
+}
 
-    return ERROR_SUCCESS;
+void WINAPI
+WintunReceiveRelease(_In_ TUN_SESSION *Session, _In_ const BYTE *Packet)
+{
+    EnterCriticalSection(&Session->Send.Lock);
+    TUN_PACKET *ReleasedBuffPacket = (TUN_PACKET *)(Packet - offsetof(TUN_PACKET, Data));
+    ReleasedBuffPacket->Size |= TUN_PACKET_RELEASE;
+    while (Session->Send.PacketsToRelease)
+    {
+        const TUN_PACKET *BuffPacket = (TUN_PACKET *)&Session->Descriptor.Send.Ring->Data[Session->Send.HeadRelease];
+        if ((BuffPacket->Size & TUN_PACKET_RELEASE) == 0)
+            break;
+        const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + (BuffPacket->Size & ~TUN_PACKET_RELEASE));
+        Session->Send.HeadRelease = TUN_RING_WRAP(Session->Send.HeadRelease + AlignedPacketSize, Session->Capacity);
+        Session->Send.PacketsToRelease--;
+    }
+    InterlockedSetU(&Session->Descriptor.Send.Ring->Head, Session->Send.HeadRelease);
+    LeaveCriticalSection(&Session->Send.Lock);
 }
 
 WINTUN_STATUS WINAPI
-WintunSendPackets(_In_ TUN_SESSION *Session, _In_ const WINTUN_PACKET *Queue)
+WintunAllocateSendPacket(_In_ TUN_SESSION *Session, _In_ DWORD PacketSize, _Out_bytecapcount_(PacketSize) BYTE **Packet)
 {
-    ULONG BuffTail = InterlockedGetU(&Session->Descriptor.Receive.Ring->Tail);
-    if (BuffTail >= Session->Capacity)
-        return ERROR_HANDLE_EOF;
-
-    for (; Queue; Queue = Queue->Next)
+    DWORD Result;
+    EnterCriticalSection(&Session->Receive.Lock);
+    if (Session->Receive.Tail >= Session->Capacity)
     {
-        const ULONG PacketSize = Queue->Size;
-        const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + PacketSize);
-
-        const ULONG BuffHead = InterlockedGetU(&Session->Descriptor.Receive.Ring->Head);
-        if (BuffHead >= Session->Capacity)
-            return ERROR_HANDLE_EOF;
-
-        const ULONG BuffSpace = TUN_RING_WRAP(BuffHead - BuffTail - TUN_ALIGNMENT, Session->Capacity);
-        if (AlignedPacketSize > BuffSpace)
-            return ERROR_BUFFER_OVERFLOW; /* Dropping when ring is full. */
-
-        TUN_PACKET *Packet = (TUN_PACKET *)&Session->Descriptor.Receive.Ring->Data[BuffTail];
-        Packet->Size = PacketSize;
-        memcpy(Packet->Data, Queue->Data, PacketSize);
-        BuffTail = TUN_RING_WRAP(BuffTail + AlignedPacketSize, Session->Capacity);
-        InterlockedSetU(&Session->Descriptor.Receive.Ring->Tail, BuffTail);
-
-        if (InterlockedGet(&Session->Descriptor.Receive.Ring->Alertable))
-            SetEvent(Session->Descriptor.Receive.TailMoved);
+        Result = ERROR_HANDLE_EOF;
+        goto cleanup;
     }
+    const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + PacketSize);
+    const ULONG BuffHead = InterlockedGetU(&Session->Descriptor.Receive.Ring->Head);
+    if (BuffHead >= Session->Capacity)
+    {
+        Result = ERROR_HANDLE_EOF;
+        goto cleanup;
+    }
+    const ULONG BuffSpace = TUN_RING_WRAP(BuffHead - Session->Receive.Tail - TUN_ALIGNMENT, Session->Capacity);
+    if (AlignedPacketSize > BuffSpace)
+    {
+        Result = ERROR_BUFFER_OVERFLOW;
+        goto cleanup;
+    }
+    TUN_PACKET *BuffPacket = (TUN_PACKET *)&Session->Descriptor.Receive.Ring->Data[Session->Receive.Tail];
+    BuffPacket->Size = PacketSize | TUN_PACKET_RELEASE;
+    *Packet = BuffPacket->Data;
+    Session->Receive.Tail = TUN_RING_WRAP(Session->Receive.Tail + AlignedPacketSize, Session->Capacity);
+    Session->Receive.PacketsToRelease++;
+    Result = ERROR_SUCCESS;
+cleanup:
+    LeaveCriticalSection(&Session->Receive.Lock);
+    return Result;
+}
 
-    return ERROR_SUCCESS;
+void WINAPI
+WintunSendPacket(_In_ TUN_SESSION *Session, _In_ const BYTE *Packet)
+{
+    EnterCriticalSection(&Session->Receive.Lock);
+    TUN_PACKET *ReleasedBuffPacket = (TUN_PACKET *)(Packet - offsetof(TUN_PACKET, Data));
+    ReleasedBuffPacket->Size &= ~TUN_PACKET_RELEASE;
+    while (Session->Receive.PacketsToRelease)
+    {
+        const TUN_PACKET *BuffPacket =
+            (TUN_PACKET *)&Session->Descriptor.Receive.Ring->Data[Session->Receive.TailRelease];
+        if (BuffPacket->Size & TUN_PACKET_RELEASE)
+            break;
+        const ULONG AlignedPacketSize = TUN_ALIGN(sizeof(TUN_PACKET) + BuffPacket->Size);
+        Session->Receive.TailRelease =
+            TUN_RING_WRAP(Session->Receive.TailRelease + AlignedPacketSize, Session->Capacity);
+        Session->Receive.PacketsToRelease--;
+    }
+    InterlockedSetU(&Session->Descriptor.Receive.Ring->Tail, Session->Receive.TailRelease);
+    if (InterlockedGet(&Session->Descriptor.Receive.Ring->Alertable))
+        SetEvent(Session->Descriptor.Receive.TailMoved);
+    LeaveCriticalSection(&Session->Receive.Lock);
 }
diff --git a/api/wintun.h b/api/wintun.h
index d2fbb40..2fe5c2e 100644
--- a/api/wintun.h
+++ b/api/wintun.h
@@ -261,27 +261,6 @@ typedef void(WINAPI *WINTUN_END_SESSION_FUNC)(_In_ WINTUN_SESSION_HANDLE Session
 #define WINTUN_MAX_IP_PACKET_SIZE 0xFFFF
 
 /**
- * Packet with data
- */
-typedef struct _WINTUN_PACKET
-{
-    /**
-     * Pointer to next packet in queue
-     */
-    struct _WINTUN_PACKET *Next;
-
-    /**
-     * Size of packet (max WINTUN_MAX_IP_PACKET_SIZE).
-     */
-    DWORD Size;
-
-    /**
-     * Pointer to layer 3 IPv4 or IPv6 packet
-     */
-    BYTE *Data;
-} WINTUN_PACKET;
-
-/**
  * Peeks if there is a packet available for reading.
  *
  * @param Session       Wintun session handle obtained with WintunStartSession
@@ -302,42 +281,67 @@ BOOL(WINAPI *WINTUN_IS_PACKET_AVAILABLE_FUNC)(_In_ WINTUN_SESSION_HANDLE Session
  *
  * @return See WaitForSingleObject() for return values.
  */
-WINTUN_STATUS(WINAPI *WINTUN_WAIT_FOR_PACKET_FUNC)(_In_ WINTUN_SESSION_HANDLE Session, _In_ DWORD Milliseconds);
+typedef WINTUN_STATUS(WINAPI *WINTUN_WAIT_FOR_PACKET_FUNC)(_In_ WINTUN_SESSION_HANDLE Session, _In_ DWORD Milliseconds);
 
 /**
- * Reads one or more packets.
+ * Retrieves one or packet. After the packet content is consumed, call WintunReceiveRelease with Packet returned
+ * from this function to release internal buffer. This function is thread-safe.
  *
  * @param Session       Wintun session handle obtained with WintunStartSession
  *
- * @param Queue         A linked list of nodes to fill with packets read. The list must be NULL terminated. May
- *                      contain only one node to read one packet at a time. All nodes should be big enough to
- *                      accommodate WINTUN_MAX_IP_PACKET_SIZE packet. The Size field of every node should be
- *                      initialized to something bigger than WINTUN_MAX_IP_PACKET_SIZE allowing post-festum
- *                      detection which nodes were actually filled with packets.
- *
- * @return
- * ERROR_HANDLE_EOF     Wintun adapter is terminating.
- * ERROR_NO_MORE_ITEMS  Wintun buffer is exhausted.
- * ERROR_INVALID_DATA   Wintun buffer is corrupt.
- * ERROR_SUCCESS        Requested amount of packets was read successfully.
+ * @param Packet        Pointer to receive pointer to layer 3 IPv4 or IPv6 packet. Client may modify its content at
+ *                      will.
+ *
+ * @param PacketSize    Pointer to receive Packet size.
+ *
+ * @return Returns one of the following values:
+ * ERROR_HANDLE_EOF     Wintun adapter is terminating;
+ * ERROR_NO_MORE_ITEMS  Wintun buffer is exhausted;
+ * ERROR_INVALID_DATA   Wintun buffer is corrupt;
+ * ERROR_SUCCESS        on success.
  * Regardless, if the error was returned, some packets might have been read nevertheless.
  */
-WINTUN_STATUS(WINAPI *WINTUN_RECEIVE_PACKETS_FUNC)
-(_In_ WINTUN_SESSION_HANDLE Session, _Inout_ WINTUN_PACKET *Queue);
+typedef WINTUN_STATUS(WINAPI *WINTUN_RECEIVE_PACKETS_FUNC)
+(_In_ WINTUN_SESSION_HANDLE *Session, _Out_bytecapcount_(*PacketSize) BYTE **Packet, _Out_ DWORD *PacketSize);
 
 /**
- * Sends packets.
+ * Releases internal buffer after the received packet has been processed by the client. This function is thread-safe.
  *
  * @param Session       Wintun session handle obtained with WintunStartSession
- * 
- * @param Queue         Linked list of packets to send. The list must be NULL terminated.
- * 
- * @return
- * ERROR_HANDLE_EOF     Wintun adapter is terminating.
- * ERROR_BUFFER_OVERFLOW  Wintun buffer is full. One or more packets were dropped.
- * ERROR_SUCCESS        All packets were sent successfully.
+ *
+ * @param Packet        Packet obtained with WintunReceivePacket
+ */
+typedef void(WINAPI *WINTUN_RECEIVE_RELEASE_FUNC)(_In_ WINTUN_SESSION_HANDLE *Session, _In_ const BYTE *Packet);
+
+/**
+ * Allocates memory for a packet to send. After the memory is filled with packet data, call WintunSendPacket to send
+ * and release internal buffer. WintunAllocateSendPacket is thread-safe and the WintunAllocateSendPacket order of
+ * calls define the packet sending order.
+ *
+ * @param Session       Wintun session handle obtained with WintunStartSession
+ *
+ * @param PacketSize    Exact packet size. Must be less or equal to WINTUN_MAX_IP_PACKET_SIZE.
+ *
+ * @param Packet        Pointer to receive pointer to memory where to prepare layer 3 IPv4 or IPv6 packet for sending.
+ *
+ * @return Returns one of the following values:
+ * ERROR_HANDLE_EOF     Wintun adapter is terminating;
+ * ERROR_BUFFER_OVERFLOW  Wintun buffer is full;
+ * ERROR_SUCCESS        on success.
+ */
+typedef WINTUN_STATUS(WINAPI *WINTUN_ALLOCATE_SEND_PACKET_FUNC)
+(_In_ WINTUN_SESSION_HANDLE *Session, _In_ DWORD PacketSize, _Out_bytecapcount_(PacketSize) BYTE **Packet);
+
+/**
+ * Sends the packet and releases internal buffer. WintunSendPacket is thread-safe, but the WintunAllocateSendPacket
+ * order of calls define the packet sending order. This means the packet is not guaranteed to be sent in the
+ * WintunSendPacket yet.
+ *
+ * @param Session       Wintun session handle obtained with WintunStartSession
+ *
+ * @param Packet        Packet obtained with WintunAllocateSendPacket
  */
-WINTUN_STATUS(WINAPI *WINTUN_SEND_PACKETS_FUNC)(_In_ WINTUN_SESSION_HANDLE Session, _In_ const WINTUN_PACKET *Queue);
+typedef void(WINAPI *WINTUN_SEND_PACKET_FUNC)(_In_ WINTUN_SESSION_HANDLE *Session, _In_ const BYTE *Packet);
 
 #ifdef __cplusplus
 }