63 lines
2.4 KiB
Go
63 lines
2.4 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/netip"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/ti-mo/conntrack"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestConntrackRecordRequiresDeclaredPortAndReplyTarget(t *testing.T) {
|
||
|
|
flow := acceptedFlow()
|
||
|
|
record, ok := conntrackRecordForFlow(&flow, []conntrackService{{
|
||
|
|
Name: "web",
|
||
|
|
Ports: []uint16{80, 443},
|
||
|
|
TargetAddress: netip.MustParseAddr("10.104.40.127"),
|
||
|
|
}})
|
||
|
|
if !ok {
|
||
|
|
t.Fatal("accepted relay flow was not classified")
|
||
|
|
}
|
||
|
|
if record.Service != "web" || record.Original.SourceAddress != "198.18.0.1" || record.Reply.SourceAddress != "10.104.40.127" || record.DurationMilliseconds != 1500 || record.TCPState != "time_wait" {
|
||
|
|
t.Fatalf("unexpected record %#v", record)
|
||
|
|
}
|
||
|
|
flow.TupleReply.IP.SourceAddress = netip.MustParseAddr("203.0.113.10")
|
||
|
|
if _, ok := conntrackRecordForFlow(&flow, []conntrackService{{Name: "web", Ports: []uint16{443}, TargetAddress: netip.MustParseAddr("10.104.40.127")}}); ok {
|
||
|
|
t.Fatal("outbound connection was classified as accepted relay traffic")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestConntrackSamplingIsDeterministic(t *testing.T) {
|
||
|
|
flow := acceptedFlow()
|
||
|
|
first := sampleConntrackFlow(&flow, 10)
|
||
|
|
for attempt := 0; attempt < 100; attempt++ {
|
||
|
|
if sampleConntrackFlow(&flow, 10) != first {
|
||
|
|
t.Fatal("sampling decision changed for an unchanged tuple")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTokenBucketBoundsBurstAndRefills(t *testing.T) {
|
||
|
|
now := time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC)
|
||
|
|
bucket := newTokenBucket(1, 2, now)
|
||
|
|
if !bucket.allow(now) || !bucket.allow(now) || bucket.allow(now) {
|
||
|
|
t.Fatal("initial burst was not bounded")
|
||
|
|
}
|
||
|
|
if !bucket.allow(now.Add(time.Second)) || bucket.allow(now.Add(time.Second)) {
|
||
|
|
t.Fatal("sustained rate was not bounded")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func acceptedFlow() conntrack.Flow {
|
||
|
|
startedAt := time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC)
|
||
|
|
flow := conntrack.NewFlow(6, 0, netip.MustParseAddr("198.18.0.1"), netip.MustParseAddr("10.49.127.43"), 50000, 443, 60, 0)
|
||
|
|
flow.TupleReply.IP.SourceAddress = netip.MustParseAddr("10.104.40.127")
|
||
|
|
flow.TupleReply.IP.DestinationAddress = netip.MustParseAddr("198.18.0.1")
|
||
|
|
flow.TupleReply.Proto.SourcePort = 443
|
||
|
|
flow.TupleReply.Proto.DestinationPort = 50000
|
||
|
|
flow.Timestamp = conntrack.Timestamp{Start: startedAt, Stop: startedAt.Add(1500 * time.Millisecond)}
|
||
|
|
flow.CountersOrig = conntrack.Counter{Packets: 3, Bytes: 180}
|
||
|
|
flow.CountersReply = conntrack.Counter{Direction: true, Packets: 4, Bytes: 240}
|
||
|
|
flow.ProtoInfo.TCP = &conntrack.ProtoInfoTCP{State: 7}
|
||
|
|
return flow
|
||
|
|
}
|