Added an interface to query the state of the connection buffers

catchpoint · Nov 11, 2016 · e09d172 · e09d172
1 parent 3b0c54e
commit e09d172
Show file tree

Hide file tree

Showing 10 changed files with 50 additions and 2 deletions.
diff --git a/bin/x64/shaper.cat b/bin/x64/shaper.cat
diff --git a/bin/x64/shaper.inf b/bin/x64/shaper.inf
@@ -8,7 +8,7 @@ Class       = WFPCALLOUTS
 ClassGuid   = {57465043-616C-6C6F-7574-5F636C617373}
 Provider    = %ProviderString%
 CatalogFile = shaper.cat
-DriverVer=11/10/2016,9.42.52.816
+DriverVer=11/11/2016,15.27.52.723
 
 [SourceDisksNames]
    1 = %ShaperDisk%,,,""

diff --git a/bin/x64/shaper.sys b/bin/x64/shaper.sys
diff --git a/bin/x86/shaper.cat b/bin/x86/shaper.cat
diff --git a/bin/x86/shaper.inf b/bin/x86/shaper.inf
@@ -8,7 +8,7 @@ Class       = WFPCALLOUTS
 ClassGuid   = {57465043-616C-6C6F-7574-5F636C617373}
 Provider    = %ProviderString%
 CatalogFile = shaper.cat
-DriverVer=11/10/2016,9.44.0.366
+DriverVer=11/11/2016,15.27.51.31
 
 [SourceDisksNames]
    1 = %ShaperDisk%,,,""

diff --git a/bin/x86/shaper.sys b/bin/x86/shaper.sys
diff --git a/driver/interface.h b/driver/interface.h
@@ -16,6 +16,13 @@ typedef struct {
   unsigned __int64 inBufferBytes;  // Size of inbound packet buffer in bytes (drop packets that overflow). 150,000 Matches the dummynet default
   unsigned __int64 outBufferBytes; // Size of outbound packet buffer in bytes (drop packets that overflow). 150,000 Matches the dummynet default
 } SHAPER_PARAMS;
+
+typedef struct {
+  BOOLEAN       enabled;           // If traffic shaping is enabled
+  SHAPER_PARAMS params;            // connection settings
+  unsigned __int64 inQueuedBytes;  // Size of the pending data in the inbound queue
+  unsigned __int64 outQueuedBytes; // Size of the pending data in the inbound queue
+} SHAPER_STATUS;
 #pragma pack(pop)
 
 // from ntifs.h
@@ -27,6 +34,10 @@ typedef struct {
 #define METHOD_BUFFERED                 0
 #endif
 
+#ifndef FILE_READ_ACCESS
+#define FILE_READ_ACCESS          ( 0x0001 )    // file & pipe
+#endif
+
 #ifndef FILE_WRITE_ACCESS
 #define FILE_WRITE_ACCESS         ( 0x0002 )    // file & pipe
 #endif
@@ -39,3 +50,4 @@ typedef struct {
 
 #define	SHAPER_IOCTL_DISABLE  CTL_CODE(FILE_DEVICE_NETWORK, 0x801, METHOD_BUFFERED, FILE_WRITE_ACCESS)
 #define	SHAPER_IOCTL_ENABLE CTL_CODE(FILE_DEVICE_NETWORK, 0x802, METHOD_BUFFERED, FILE_WRITE_ACCESS)
+#define	SHAPER_IOCTL_GET_STATUS CTL_CODE(FILE_DEVICE_NETWORK, 0x803, METHOD_BUFFERED, FILE_READ_ACCESS)
diff --git a/driver/packet_queue.c b/driver/packet_queue.c
@@ -443,3 +443,18 @@ BOOLEAN ShaperQueuePacket(_In_ const FWPS_INCOMING_VALUES* inFixedValues,
   return queued;
 }
 
+/*-----------------------------------------------------------------------------
+-----------------------------------------------------------------------------*/
+void ShaperGetStatus(SHAPER_STATUS *status) {
+  // Don't bother locking to read the values out
+  status->enabled = traffic_shaping_enabled;
+  status->params.inBps = inbound_queue.bps;
+  status->params.outBps = outbound_queue.bps;
+  status->params.inLatency = inbound_queue.latency;
+  status->params.outLatency = outbound_queue.latency;
+  status->params.plr = inbound_queue.plr;
+  status->params.inBufferBytes = inbound_queue.bufferBytes;
+  status->params.outBufferBytes = outbound_queue.bufferBytes;
+  status->inQueuedBytes = inbound_queue.queued_bytes;
+  status->outQueuedBytes = outbound_queue.queued_bytes;
+}
diff --git a/driver/packet_queue.h b/driver/packet_queue.h
@@ -1,4 +1,5 @@
 #pragma once
+#include "interface.h"
 
 NTSTATUS InitializePacketQueues(WDFDEVICE timer_parent);
 void DestroyPacketQueues();
@@ -18,3 +19,4 @@ BOOLEAN ShaperEnable(_In_ unsigned short plr,
                      _In_ unsigned __int64 outBufferBytes);
 
 BOOLEAN ShaperDisable();
+void ShaperGetStatus(SHAPER_STATUS *status);
diff --git a/driver/wdm_driver.c b/driver/wdm_driver.c
@@ -210,6 +210,25 @@ VOID EvtDeviceIOCtl(_In_ WDFQUEUE Queue,
       }
       break;
     }
+    case SHAPER_IOCTL_GET_STATUS: {
+      if (OutputBufferLength >= sizeof(SHAPER_STATUS)) {
+        WDFMEMORY pMemory;
+        SHAPER_STATUS* shaper_status;
+        status = WdfRequestRetrieveOutputMemory(Request, &pMemory);
+        if (NT_SUCCESS(status)) {
+          shaper_status = (SHAPER_STATUS *)WdfMemoryGetBuffer(pMemory, NULL);
+          if (shaper_status) {
+            ShaperGetStatus(shaper_status);
+            WdfRequestSetInformation(Request, sizeof(SHAPER_STATUS));
+          } else {
+            status = STATUS_INVALID_PARAMETER;
+          }
+        }
+      } else {
+        status = STATUS_INVALID_PARAMETER;
+      }
+      break;
+    }
     default: {
       DbgPrint("[shaper] IOCTL unknown command - %lu\n", IoControlCode);
       status = STATUS_INVALID_PARAMETER;