Add Wireshark packet dissector plugin and associated tools

2026-08-20 04:22:53 +00:00 · 2024-08-23 11:02:40 -07:00
parent fadd2c4259
commit e857bbb09f
26 changed files with 1475 additions and 0 deletions
@@ -0,0 +1,102 @@
+strip_nothing_(::Type{Union{Nothing, T}}) where T = T
+strip_nothing_(::Type{T}) where T = T
+
+progress_noop_(n) = nothing
+
+mutable struct SplitCapOutput{S}
+    work_buffer::Vector{UInt8}
+    complete_buffers::Channel{Vector{UInt8}}
+    stream::S
+end
+
+# Since Julia (as of 1.5) doesn't support task migration,
+# continue with new task after each buffer write, to rebalance across threads
+function write_one_and_continue_(output::SplitCapOutput, free_buffers::Channel, pending::Threads.Atomic{Int})
+    i = iterate(output.complete_buffers)
+    if i === nothing
+        Threads.atomic_sub!(pending, 1)
+        return nothing
+    end
+    b, _ = i
+    write(output.stream, b)
+    empty!(b)
+    put!(free_buffers, b)
+    Threads.@spawn write_one_and_continue_($output, $free_buffers, $pending)
+end
+
+function splitcap(
+    ::Type{KeyType},
+    ::Type{StreamType},
+    reader::PcapReader,
+    record2key,
+    key2stream,
+    progress_callback = progress_noop_;
+    own_streams::Bool = true
+) where {KeyType, StreamType}
+    buffer_size = 1024 * 1024 * 2
+    max_pending_buffers = 4
+    outputs = Dict{KeyType, SplitCapOutput{StreamType}}()
+    free_buffers = Channel{Vector{UInt8}}(Inf)
+    n = 0
+    pending = Threads.Atomic{Int}(0)
+    try
+        while !eof(reader)
+            record = read(reader)
+            dst = record2key(record)
+            if dst isa KeyType
+                output = get!(outputs, dst) do
+                    stream = key2stream(dst)
+                    buffer = sizehint!(UInt8[], buffer_size + 1500)
+                    own_streams && append!(buffer, reader.raw_header)
+                    output = SplitCapOutput{StreamType}(
+                        buffer,
+                        Channel{Vector{UInt8}}(max_pending_buffers),
+                        stream)
+                    Threads.atomic_add!(pending, 1)
+                    Threads.@spawn write_one_and_continue_($output, $free_buffers, $pending)
+                    output
+                end
+                append!(output.work_buffer, record.raw)
+                if length(output.work_buffer) >= buffer_size
+                    put!(output.complete_buffers, output.work_buffer)
+                    if isready(free_buffers)
+                        output.work_buffer = take!(free_buffers)
+                    else
+                        output.work_buffer = sizehint!(UInt8[], buffer_size + 1500)
+                    end
+                end
+            end
+            n += 1
+            progress_callback(n)
+            GC.safepoint()
+        end
+        for output in values(outputs)
+            if !isempty(output.work_buffer)
+                put!(output.complete_buffers, output.work_buffer)
+            end
+            close(output.complete_buffers)
+        end
+        while pending[] != 0
+            sleep(0.1)
+        end
+    finally
+        if own_streams
+            for output in values(outputs)
+                close(output.stream)
+            end
+        end
+    end
+    nothing
+end
+
+function splitcap(
+    reader::PcapReader,
+    record2key,
+    key2stream,
+    progress_callback = progress_noop_;
+    kwargs...
+)
+    KeyType = strip_nothing_(Core.Compiler.return_type(record2key, Tuple{PcapRecord}))
+    StreamType = Core.Compiler.return_type(key2stream, Tuple{KeyType})
+    splitcap(KeyType, StreamType, reader, record2key, key2stream, progress_callback; kwargs...)
+end