bcap/test_helpers_test.go
2026-03-24 23:39:55 +08:00

355 lines
8.0 KiB
Go

package bcap
import (
"net"
"testing"
"time"
"github.com/gopacket/gopacket"
"github.com/gopacket/gopacket/layers"
)
type observedPacket struct {
Observation
Key string
ReverseKey string
}
func wrapObservedPacket(obs Observation) observedPacket {
return observedPacket{
Observation: obs,
Key: obs.Flow.Forward.StableString(),
ReverseKey: obs.Flow.Reverse.StableString(),
}
}
func (p observedPacket) StateDescript() uint8 {
return legacyStateFromObservation(p.Observation)
}
func (p observedPacket) TcpPayloads() int {
if tcp := p.Packet.Transport.TCP; tcp != nil {
return tcp.Payload
}
return p.Packet.Transport.Payload
}
func (p observedPacket) TcpSeq() uint32 {
if tcp := p.Packet.Transport.TCP; tcp != nil {
return tcp.Seq
}
return 0
}
func (p observedPacket) TcpAck() uint32 {
if tcp := p.Packet.Transport.TCP; tcp != nil {
return tcp.Ack
}
return 0
}
func (p observedPacket) TcpWindow() uint16 {
if tcp := p.Packet.Transport.TCP; tcp != nil {
return tcp.Window
}
return 0
}
func legacyStateFromObservation(obs Observation) uint8 {
switch obs.Packet.Transport.Kind {
case ProtocolTCP:
if obs.Hints.TCP != nil {
return obs.Hints.TCP.LegacyState
}
case ProtocolUDP:
return StateUdp
case ProtocolICMPv4:
return StateIcmp
case ProtocolICMPv6:
return StateIcmpv6
}
return StateUnknown
}
type tcpPacketSpec struct {
srcIP string
dstIP string
srcPort uint16
dstPort uint16
seq uint32
ack uint32
syn bool
ackFlag bool
fin bool
rst bool
ece bool
cwr bool
window uint16
payload []byte
}
type udpPacketSpec struct {
srcIP string
dstIP string
srcPort uint16
dstPort uint16
payload []byte
}
type icmpPacketSpec struct {
srcIP string
dstIP string
typ uint8
code uint8
id uint16
seq uint16
payload []byte
}
type arpPacketSpec struct {
srcMAC string
dstMAC string
senderMAC string
targetMAC string
senderIP string
targetIP string
operation uint16
}
func newTestAnalyzer() *Analyzer {
cfg := DefaultConfig()
cfg.CleanupInterval = 0
return NewAnalyzerWithConfig(cfg)
}
func mustObservePacket(t *testing.T, analyzer *Analyzer, packet gopacket.Packet) observedPacket {
t.Helper()
obs, err := analyzer.ObservePacket(packet)
if err != nil {
t.Fatalf("observe packet: %v", err)
}
return wrapObservedPacket(obs)
}
func mustEstablishTCPConnection(
t *testing.T,
analyzer *Analyzer,
base time.Time,
clientIP string,
clientPort uint16,
serverIP string,
serverPort uint16,
clientSeq uint32,
serverSeq uint32,
) (string, string) {
t.Helper()
mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base, tcpPacketSpec{
srcIP: clientIP,
dstIP: serverIP,
srcPort: clientPort,
dstPort: serverPort,
seq: clientSeq,
syn: true,
window: 4096,
}))
serverInfo := mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base.Add(time.Millisecond), tcpPacketSpec{
srcIP: serverIP,
dstIP: clientIP,
srcPort: serverPort,
dstPort: clientPort,
seq: serverSeq,
ack: clientSeq + 1,
syn: true,
ackFlag: true,
window: 4096,
}))
clientInfo := mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base.Add(2*time.Millisecond), tcpPacketSpec{
srcIP: clientIP,
dstIP: serverIP,
srcPort: clientPort,
dstPort: serverPort,
seq: clientSeq + 1,
ack: serverSeq + 1,
ackFlag: true,
window: 4096,
}))
return clientInfo.Key, serverInfo.Key
}
func mustBuildTCPPacket(t *testing.T, ts time.Time, spec tcpPacketSpec) gopacket.Packet {
t.Helper()
ip := &layers.IPv4{
Version: 4,
IHL: 5,
TTL: 64,
SrcIP: net.ParseIP(spec.srcIP).To4(),
DstIP: net.ParseIP(spec.dstIP).To4(),
Protocol: layers.IPProtocolTCP,
}
tcp := &layers.TCP{
SrcPort: layers.TCPPort(spec.srcPort),
DstPort: layers.TCPPort(spec.dstPort),
Seq: spec.seq,
Ack: spec.ack,
SYN: spec.syn,
ACK: spec.ackFlag,
FIN: spec.fin,
RST: spec.rst,
ECE: spec.ece,
CWR: spec.cwr,
Window: spec.window,
}
tcp.SetNetworkLayerForChecksum(ip)
buf := gopacket.NewSerializeBuffer()
opts := gopacket.SerializeOptions{
FixLengths: true,
ComputeChecksums: true,
}
if err := gopacket.SerializeLayers(buf, opts, ip, tcp, gopacket.Payload(spec.payload)); err != nil {
t.Fatalf("serialize tcp packet: %v", err)
}
packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
packet.Metadata().Timestamp = ts
return packet
}
func mustBuildUDPPacket(t *testing.T, ts time.Time, spec udpPacketSpec) gopacket.Packet {
t.Helper()
ip := &layers.IPv4{
Version: 4,
IHL: 5,
TTL: 64,
SrcIP: net.ParseIP(spec.srcIP).To4(),
DstIP: net.ParseIP(spec.dstIP).To4(),
Protocol: layers.IPProtocolUDP,
}
udp := &layers.UDP{
SrcPort: layers.UDPPort(spec.srcPort),
DstPort: layers.UDPPort(spec.dstPort),
}
udp.SetNetworkLayerForChecksum(ip)
buf := gopacket.NewSerializeBuffer()
opts := gopacket.SerializeOptions{
FixLengths: true,
ComputeChecksums: true,
}
if err := gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(spec.payload)); err != nil {
t.Fatalf("serialize udp packet: %v", err)
}
packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
packet.Metadata().Timestamp = ts
return packet
}
func mustBuildICMPPacket(t *testing.T, ts time.Time, spec icmpPacketSpec) gopacket.Packet {
t.Helper()
ip := &layers.IPv4{
Version: 4,
IHL: 5,
TTL: 64,
SrcIP: net.ParseIP(spec.srcIP).To4(),
DstIP: net.ParseIP(spec.dstIP).To4(),
Protocol: layers.IPProtocolICMPv4,
}
icmp := &layers.ICMPv4{
TypeCode: layers.CreateICMPv4TypeCode(spec.typ, spec.code),
Id: spec.id,
Seq: spec.seq,
}
buf := gopacket.NewSerializeBuffer()
opts := gopacket.SerializeOptions{
FixLengths: true,
ComputeChecksums: true,
}
if err := gopacket.SerializeLayers(buf, opts, ip, icmp, gopacket.Payload(spec.payload)); err != nil {
t.Fatalf("serialize icmp packet: %v", err)
}
packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
packet.Metadata().Timestamp = ts
return packet
}
func mustBuildIPv4OnlyPacket(t *testing.T, ts time.Time, srcIP, dstIP string) gopacket.Packet {
t.Helper()
ip := &layers.IPv4{
Version: 4,
IHL: 5,
TTL: 64,
SrcIP: net.ParseIP(srcIP).To4(),
DstIP: net.ParseIP(dstIP).To4(),
Protocol: layers.IPProtocolNoNextHeader,
}
buf := gopacket.NewSerializeBuffer()
opts := gopacket.SerializeOptions{
FixLengths: true,
ComputeChecksums: true,
}
if err := gopacket.SerializeLayers(buf, opts, ip); err != nil {
t.Fatalf("serialize ipv4-only packet: %v", err)
}
packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
packet.Metadata().Timestamp = ts
return packet
}
func mustBuildARPPacket(t *testing.T, ts time.Time, spec arpPacketSpec) gopacket.Packet {
t.Helper()
eth := &layers.Ethernet{
SrcMAC: mustMAC(t, spec.srcMAC),
DstMAC: mustMAC(t, spec.dstMAC),
EthernetType: layers.EthernetTypeARP,
}
arp := &layers.ARP{
AddrType: layers.LinkTypeEthernet,
Protocol: layers.EthernetTypeIPv4,
HwAddressSize: 6,
ProtAddressSize: 4,
Operation: spec.operation,
SourceHwAddress: []byte(mustMAC(t, spec.senderMAC)),
SourceProtAddress: []byte(net.ParseIP(spec.senderIP).To4()),
DstHwAddress: []byte(mustMAC(t, spec.targetMAC)),
DstProtAddress: []byte(net.ParseIP(spec.targetIP).To4()),
}
buf := gopacket.NewSerializeBuffer()
opts := gopacket.SerializeOptions{
FixLengths: true,
ComputeChecksums: true,
}
if err := gopacket.SerializeLayers(buf, opts, eth, arp); err != nil {
t.Fatalf("serialize arp packet: %v", err)
}
packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
packet.Metadata().Timestamp = ts
return packet
}
func mustMAC(t *testing.T, raw string) net.HardwareAddr {
t.Helper()
hw, err := net.ParseMAC(raw)
if err != nil {
t.Fatalf("parse mac %q: %v", raw, err)
}
return hw
}