From e23100a93513112cdbf7bff6b142e51bec20431f Mon Sep 17 00:00:00 2001
From: Sebastian Solnica <ssolnica@gmail.com>
Date: Tue, 2 Mar 2021 19:50:38 +0100
Subject: [PATCH] Added support for UDP (#14)

---
 wtrace.cmd/Events/UdpIp.fs   | 65 ++++++++++++++++++++++++++++++++++++
 wtrace.cmd/Program.fs        |  4 ++-
 wtrace.cmd/wtrace.cmd.fsproj |  1 +
 wtrace/wtrace.csproj         |  4 +--
 4 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 wtrace.cmd/Events/UdpIp.fs

diff --git a/wtrace.cmd/Events/UdpIp.fs b/wtrace.cmd/Events/UdpIp.fs
new file mode 100644
index 0000000..a7a06ae
--- /dev/null
+++ b/wtrace.cmd/Events/UdpIp.fs
@@ -0,0 +1,65 @@
+module LowLevelDesign.WTrace.Events.UdpIp
+
+open System
+open Microsoft.Diagnostics.Tracing
+open Microsoft.Diagnostics.Tracing.Parsers.Kernel
+open LowLevelDesign.WTrace
+open LowLevelDesign.WTrace.Events
+open LowLevelDesign.WTrace.Events.HandlerCommons
+
+type private UdpIpHandlerState = {
+    Broadcast : EventBroadcast
+}
+
+[<AutoOpen>]
+module private H =
+
+    let noFields = Array.empty<TraceEventField>
+    
+    let handleUdpIpFail id state (ev : UdpIpFailTraceData) =
+        let ev = toEvent ev id "" "" "" ev.FailureCode
+        state.Broadcast.publishTraceEvent (TraceEventWithFields (ev, noFields))
+
+    let handleUdpIpData id state (ev : UdpIpTraceData) =
+        let fields = [|
+            struct (nameof ev.size, FI32 ev.size) |]
+
+        let details = sprintf "size: %d" ev.size
+        let path = sprintf "%s:%d -> %s:%d" (ev.saddr.ToString()) ev.sport (ev.daddr.ToString()) ev.dport
+        let ev = toEvent ev id "" path details WinApi.eventStatusUndefined
+        state.Broadcast.publishTraceEvent (TraceEventWithFields (ev, fields |> Array.map (toEventField id)))
+
+    let handleUdpIp6Data id state (ev : UpdIpV6TraceData) =
+        let fields = [|
+            struct (nameof ev.size, FI32 ev.size) |]
+
+        let details = sprintf "size: %d" ev.size
+        let path = sprintf "%s:%d -> %s:%d" (ev.saddr.ToString()) ev.sport (ev.daddr.ToString()) ev.dport
+        let activityId = sprintf "conn#%d" ev.connid
+        let ev = toEvent ev id activityId path details WinApi.eventStatusUndefined
+        state.Broadcast.publishTraceEvent (TraceEventWithFields (ev, fields |> Array.map (toEventField id)))
+
+    let subscribe (source : TraceEventSource, isRundown, idgen, state : obj) =
+        let state = state :?> UdpIpHandlerState
+        let handleEvent h = Action<_>(handleEvent idgen state h)
+        if not isRundown then
+            source.Kernel.add_UdpIpFail(handleEvent handleUdpIpFail)
+            source.Kernel.add_UdpIpRecv(handleEvent handleUdpIpData)
+            source.Kernel.add_UdpIpRecvIPV6(handleEvent handleUdpIp6Data)
+            source.Kernel.add_UdpIpSend(handleEvent handleUdpIpData)
+            source.Kernel.add_UdpIpSendIPV6(handleEvent handleUdpIp6Data)
+
+
+let createEtwHandler () =
+    {
+        KernelFlags = NtKeywords.NetworkTCPIP
+        KernelStackFlags = NtKeywords.NetworkTCPIP
+        KernelRundownFlags = NtKeywords.None
+        Providers = Array.empty<EtwEventProvider>
+        Initialize = 
+            fun (broadcast) -> ({
+                Broadcast = broadcast
+            } :> obj)
+        Subscribe = subscribe
+    }
+
diff --git a/wtrace.cmd/Program.fs b/wtrace.cmd/Program.fs
index 25f8dc6..d04da7b 100644
--- a/wtrace.cmd/Program.fs
+++ b/wtrace.cmd/Program.fs
@@ -47,6 +47,7 @@ Options:
     registry  - to receive Registry events (voluminous, disabled by default)
     rpc       - to receive RPC events
     tcp       - to receive TCP/IP events
+    udp       - to receive UDP events
 
   Example: --handlers 'tcp,file,registry'
 
@@ -79,6 +80,7 @@ let parseHandlers args =
             elif name >=< "registry" then Registry.createEtwHandler ()
             elif name >=< "rpc" then Rpc.createEtwHandler ()
             elif name >=< "tcp" then TcpIp.createEtwHandler ()
+            elif name >=< "udp" then UdpIp.createEtwHandler ()
             else failwith (sprintf "Invalid handler name: '%s'" name)
 
         try
@@ -96,7 +98,7 @@ let parseHandlers args =
         | Failure msg -> Error msg
 
     match args |> Map.tryFind "handlers" with
-    | None -> createHandlers "process,file,rpc,tcp"
+    | None -> createHandlers "process,file,rpc,tcp,udp"
     | Some [ handler ] ->
         if isSystemTrace args then
             Error ("Handlers are not allowed in the system trace.")
diff --git a/wtrace.cmd/wtrace.cmd.fsproj b/wtrace.cmd/wtrace.cmd.fsproj
index 427857c..43c3fad 100644
--- a/wtrace.cmd/wtrace.cmd.fsproj
+++ b/wtrace.cmd/wtrace.cmd.fsproj
@@ -21,6 +21,7 @@
     <Compile Include="Events\IsrDpc.fs" />
     <Compile Include="Events\Rpc.fs" />
     <Compile Include="Events\TcpIp.fs" />
+    <Compile Include="Events\UdpIp.fs" />
     <Compile Include="Tracing\Commons.fs" />
     <Compile Include="Tracing\EtwTraceSession.fs" />
     <Compile Include="Tracing\EventFilter.fs" />
diff --git a/wtrace/wtrace.csproj b/wtrace/wtrace.csproj
index e69c585..fee1d55 100644
--- a/wtrace/wtrace.csproj
+++ b/wtrace/wtrace.csproj
@@ -6,8 +6,8 @@
     <RootNamespace>LowLevelDesign.WTrace</RootNamespace>
     <Authors>Sebastian Solnica</Authors>
     <Company>Sebastian Solnica (lowleveldesign.org)</Company>
-    <AssemblyVersion>3.0.0.0</AssemblyVersion>
-    <FileVersion>3.0.0.0</FileVersion>
+    <AssemblyVersion>3.1.0.0</AssemblyVersion>
+    <FileVersion>3.1.0.0</FileVersion>
     <SatelliteResourceLanguages>en</SatelliteResourceLanguages>
     <OutputPath>..\bin\wtrace</OutputPath>
     <TreatWarningsAsErrors>true</TreatWarningsAsErrors>