Add missing files, update style
This commit is contained in:
@@ -0,0 +1,216 @@
|
||||
package pcap
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"errors"
|
||||
"io"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// A reorderEntry is 16 bytes, which bounds the memory needed per packet.
|
||||
type reorderEntry struct {
|
||||
ts uint64 // seconds << 32 | fraction (µs or ns, consistent within a file)
|
||||
pos uint64 // record offset << lenBits | record length (including the 16 byte record header)
|
||||
}
|
||||
|
||||
const (
|
||||
lenBits = 28 // fits MaxRecordLen + 16
|
||||
maxOffset = 1<<(64-lenBits) - 1
|
||||
)
|
||||
|
||||
func (e reorderEntry) off() int64 { return int64(e.pos >> lenBits) }
|
||||
func (e reorderEntry) len() int { return int(e.pos & (1<<lenBits - 1)) }
|
||||
|
||||
// Tuning for WriteTo: packets are copied in batches of up to batchBytes,
|
||||
// using parallel reads since random reads from SSDs are much faster with
|
||||
// several requests in flight. The next batch is read while the previous one
|
||||
// is written.
|
||||
const (
|
||||
batchPackets = 16384
|
||||
readers = 16
|
||||
)
|
||||
|
||||
var batchBytes = 8 << 20 // variable for tests
|
||||
|
||||
// Reorder sorts the packets of a pcap file by timestamp. NewReorder reads
|
||||
// and indexes the file, so all parse errors are reported before WriteTo
|
||||
// produces any output.
|
||||
type Reorder struct {
|
||||
r io.ReaderAt
|
||||
header []byte
|
||||
entries []reorderEntry
|
||||
end int64 // offset just after the last complete record
|
||||
outOfOrder int
|
||||
truncated bool
|
||||
}
|
||||
|
||||
// NewReorder indexes the pcap file in r (of the given size) and sorts the
|
||||
// index by timestamp. The sort is stable: packets with equal timestamps keep
|
||||
// their original order. It needs 16 bytes of memory per packet.
|
||||
func NewReorder(r io.ReaderAt, size int64) (*Reorder, error) {
|
||||
pr, err := NewReader(io.NewSectionReader(r, 0, size))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ro := &Reorder{r: r, header: slices.Clone(pr.Header())}
|
||||
off := int64(len(ro.header))
|
||||
var latest uint64
|
||||
for {
|
||||
rec, err := pr.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if errors.Is(err, ErrTruncated) {
|
||||
ro.truncated = true
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if off > maxOffset {
|
||||
return nil, errors.New("file too large to reorder")
|
||||
}
|
||||
ts := uint64(pr.order.Uint32(rec[0:4]))<<32 | uint64(pr.order.Uint32(rec[4:8]))
|
||||
if ts < latest {
|
||||
ro.outOfOrder++
|
||||
}
|
||||
latest = max(latest, ts)
|
||||
ro.entries = append(ro.entries, reorderEntry{ts: ts, pos: uint64(off)<<lenBits | uint64(len(rec))})
|
||||
off += int64(len(rec))
|
||||
|
||||
if len(ro.entries) == 1024 {
|
||||
// Size the index from the average packet size so far, to avoid
|
||||
// repeatedly growing (and copying) a large slice.
|
||||
avg := (off - int64(len(ro.header))) / 1024
|
||||
est := size / avg
|
||||
ro.entries = slices.Grow(ro.entries, int(est+est/16))
|
||||
}
|
||||
}
|
||||
if len(ro.entries) == 0 {
|
||||
return nil, errors.New("the pcap file contains no packets")
|
||||
}
|
||||
ro.end = off
|
||||
ro.entries = slices.Clip(ro.entries)
|
||||
if ro.outOfOrder > 0 {
|
||||
slices.SortStableFunc(ro.entries, func(a, b reorderEntry) int { return cmp.Compare(a.ts, b.ts) })
|
||||
}
|
||||
return ro, nil
|
||||
}
|
||||
|
||||
// Packets returns the number of packets.
|
||||
func (ro *Reorder) Packets() int { return len(ro.entries) }
|
||||
|
||||
// OutOfOrder returns the number of packets with a timestamp earlier than a
|
||||
// packet before them in the input.
|
||||
func (ro *Reorder) OutOfOrder() int { return ro.outOfOrder }
|
||||
|
||||
// Truncated reports whether the input ended with an incomplete packet
|
||||
// record, which is left out of the output.
|
||||
func (ro *Reorder) Truncated() bool { return ro.truncated }
|
||||
|
||||
type batch struct {
|
||||
buf []byte
|
||||
err error
|
||||
}
|
||||
|
||||
// WriteTo writes the file header followed by the packets in timestamp order.
|
||||
func (ro *Reorder) WriteTo(w io.Writer) (int64, error) {
|
||||
if ro.outOfOrder == 0 {
|
||||
return io.Copy(w, io.NewSectionReader(ro.r, 0, ro.end))
|
||||
}
|
||||
|
||||
n, err := w.Write(ro.header)
|
||||
total := int64(n)
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
|
||||
// A producer goroutine reads batches into two alternating buffers while
|
||||
// this goroutine writes them out.
|
||||
full := make(chan batch)
|
||||
free := make(chan []byte, 2)
|
||||
free <- nil
|
||||
free <- nil
|
||||
done := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
defer func() {
|
||||
close(done)
|
||||
wg.Wait()
|
||||
}()
|
||||
wg.Go(func() {
|
||||
defer close(full)
|
||||
for i := 0; i < len(ro.entries); {
|
||||
var buf []byte
|
||||
select {
|
||||
case buf = <-free:
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
var err error
|
||||
i, buf, err = ro.readBatch(i, buf)
|
||||
select {
|
||||
case full <- batch{buf, err}:
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
for b := range full {
|
||||
if b.err != nil {
|
||||
return total, b.err
|
||||
}
|
||||
n, err := w.Write(b.buf)
|
||||
total += int64(n)
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
free <- b.buf
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
// readBatch reads the records of entries[i:j] into buf, choosing j so that
|
||||
// the batch stays within batchBytes and batchPackets (but holds at least one
|
||||
// record). It returns j and the filled buffer.
|
||||
func (ro *Reorder) readBatch(i int, buf []byte) (int, []byte, error) {
|
||||
j, size := i, 0
|
||||
for j < len(ro.entries) && j-i < batchPackets && (j == i || size+ro.entries[j].len() <= batchBytes) {
|
||||
size += ro.entries[j].len()
|
||||
j++
|
||||
}
|
||||
entries := ro.entries[i:j]
|
||||
dst := make([]int, len(entries)+1)
|
||||
for k, e := range entries {
|
||||
dst[k+1] = dst[k] + e.len()
|
||||
}
|
||||
buf = slices.Grow(buf[:0], size)[:size]
|
||||
|
||||
var (
|
||||
next atomic.Int64
|
||||
wg sync.WaitGroup
|
||||
errs [readers]error
|
||||
)
|
||||
for r := range min(readers, len(entries)) {
|
||||
wg.Go(func() {
|
||||
for {
|
||||
k := int(next.Add(1) - 1)
|
||||
if k >= len(entries) {
|
||||
return
|
||||
}
|
||||
e := entries[k]
|
||||
if _, err := ro.r.ReadAt(buf[dst[k]:dst[k+1]], e.off()); err != nil {
|
||||
errs[r] = err
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
return j, buf, errors.Join(errs[:]...)
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package pcap
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"math/rand/v2"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// reorderTestFile builds a pcap file with one packet per timestamp; each
|
||||
// packet's data is its index in the input.
|
||||
func reorderTestFile(t *testing.T, ts ...time.Time) []byte {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
w, err := NewWriter(&buf, LinkTypeEthernet)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i, tm := range ts {
|
||||
if err := w.WritePacket(tm, bytes.Repeat([]byte{byte(i)}, i+1)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// packetIDs returns the first data byte of each packet in a pcap file.
|
||||
func packetIDs(t *testing.T, file []byte) []int {
|
||||
t.Helper()
|
||||
r, err := NewReader(bytes.NewReader(file))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var ids []int
|
||||
for {
|
||||
rec, err := r.Next()
|
||||
if err == io.EOF {
|
||||
return ids
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ids = append(ids, int(rec[16]))
|
||||
}
|
||||
}
|
||||
|
||||
func TestReorder(t *testing.T) {
|
||||
s := func(sec int64, usec int) time.Time { return time.Unix(sec, int64(usec)*1000) }
|
||||
tests := []struct {
|
||||
name string
|
||||
ts []time.Time
|
||||
want []int
|
||||
outOfOrder int
|
||||
}{
|
||||
{"sorted", []time.Time{s(1, 0), s(1, 1), s(2, 0)}, []int{0, 1, 2}, 0},
|
||||
{"reversed", []time.Time{s(3, 0), s(2, 0), s(1, 0)}, []int{2, 1, 0}, 2},
|
||||
{"microseconds", []time.Time{s(5, 999999), s(5, 2), s(5, 10)}, []int{1, 2, 0}, 2},
|
||||
{"stable for equal timestamps", []time.Time{s(2, 0), s(1, 0), s(2, 0), s(1, 0)}, []int{1, 3, 0, 2}, 2},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
in := reorderTestFile(t, tt.ts...)
|
||||
ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out bytes.Buffer
|
||||
n, err := ro.WriteTo(&out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != int64(len(in)) || out.Len() != len(in) {
|
||||
t.Errorf("wrote %d (%d) bytes, want %d", n, out.Len(), len(in))
|
||||
}
|
||||
if !bytes.Equal(out.Bytes()[:24], in[:24]) {
|
||||
t.Error("file header changed")
|
||||
}
|
||||
if got := packetIDs(t, out.Bytes()); !slices.Equal(got, tt.want) {
|
||||
t.Errorf("order %v, want %v", got, tt.want)
|
||||
}
|
||||
if ro.Packets() != len(tt.ts) || ro.OutOfOrder() != tt.outOfOrder || ro.Truncated() {
|
||||
t.Errorf("packets %d outOfOrder %d truncated %v", ro.Packets(), ro.OutOfOrder(), ro.Truncated())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReorderTruncated(t *testing.T) {
|
||||
in := reorderTestFile(t, time.Unix(2, 0), time.Unix(1, 0), time.Unix(0, 0))
|
||||
in = in[:len(in)-1]
|
||||
ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out bytes.Buffer
|
||||
if _, err := ro.WriteTo(&out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !ro.Truncated() || !slices.Equal(packetIDs(t, out.Bytes()), []int{1, 0}) {
|
||||
t.Errorf("truncated %v, order %v", ro.Truncated(), packetIDs(t, out.Bytes()))
|
||||
}
|
||||
}
|
||||
|
||||
func TestReorderErrors(t *testing.T) {
|
||||
for name, tc := range map[string]struct {
|
||||
in []byte
|
||||
want string
|
||||
}{
|
||||
"no packets": {reorderTestFile(t), "no packets"},
|
||||
"garbage": {bytes.Repeat([]byte("y"), 30), "unknown magic"},
|
||||
} {
|
||||
_, err := NewReorder(bytes.NewReader(tc.in), int64(len(tc.in)))
|
||||
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
||||
t.Errorf("%s: error %v, want %q", name, err, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestReorderManyPackets covers multiple parallel batches, including packets
|
||||
// that are bigger than a whole batch, against a simple reference sort.
|
||||
func TestReorderManyPackets(t *testing.T) {
|
||||
defer func(n int) { batchBytes = n }(batchBytes)
|
||||
batchBytes = 64 << 10
|
||||
|
||||
r := rand.New(rand.NewPCG(1, 2))
|
||||
var buf bytes.Buffer
|
||||
w, err := NewWriter(&buf, LinkTypeEthernet)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
type pkt struct {
|
||||
ts time.Time
|
||||
data []byte
|
||||
}
|
||||
var pkts []pkt
|
||||
for i := range 50000 {
|
||||
size := 8 + r.IntN(200)
|
||||
if i%10000 == 5 {
|
||||
size = batchBytes + 1000 // gets a batch by itself
|
||||
}
|
||||
data := make([]byte, size)
|
||||
binary.LittleEndian.PutUint64(data, uint64(i))
|
||||
p := pkt{time.Unix(int64(r.IntN(1000)), int64(r.IntN(1000))*1000), data}
|
||||
pkts = append(pkts, p)
|
||||
if err := w.WritePacket(p.ts, p.data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
slices.SortStableFunc(pkts, func(a, b pkt) int { return a.ts.Compare(b.ts) })
|
||||
|
||||
in := buf.Bytes()
|
||||
ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out bytes.Buffer
|
||||
if _, err := ro.WriteTo(&out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var want bytes.Buffer
|
||||
ww, _ := NewWriter(&want, LinkTypeEthernet)
|
||||
for _, p := range pkts {
|
||||
ww.WritePacket(p.ts, p.data)
|
||||
}
|
||||
if !bytes.Equal(out.Bytes(), want.Bytes()) {
|
||||
t.Error("reordered output differs from reference sort")
|
||||
}
|
||||
}
|
||||
|
||||
type failingWriter struct{ n int }
|
||||
|
||||
func (f *failingWriter) Write(p []byte) (int, error) {
|
||||
if f.n--; f.n < 0 {
|
||||
return 0, errors.New("write failed")
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func TestReorderWriteError(t *testing.T) {
|
||||
in := reorderTestFile(t, time.Unix(2, 0), time.Unix(1, 0))
|
||||
ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for n := range 2 {
|
||||
if _, err := ro.WriteTo(&failingWriter{n: n}); err == nil {
|
||||
t.Errorf("failing after %d writes: no error", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user