test_reassembly_invariant.go

package main

import (
	"bufio"
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"github.com/google/gopacket/reassembly"
)

func NewAssembler() *reassembly.Assembler {
	streamFactory := &myFactory{}
	streamPool := reassembly.NewStreamPool(streamFactory)
	return reassembly.NewAssembler(streamPool)
}

type myFactory struct {
	counter int64
}

func (f *myFactory) New(netFlow, tcpFlow gopacket.Flow, tcp *layers.TCP, ac reassembly.AssemblerContext) reassembly.Stream {
	f.counter += 1

	fsmOptions := reassembly.TCPSimpleFSMOptions{
		SupportMissingEstablishment: false,
	}
	s := &myStream{
		i:        f.counter,
		tcpstate: reassembly.NewTCPSimpleFSM(fsmOptions),
		NetFlow:  netFlow,
		TCPFlow:  tcpFlow,
	}
	log.Printf("stream #%d: (created)", s.i)
	return s
}

type myStream struct {
	i int64
	b bool

	NetFlow gopacket.Flow
	TCPFlow gopacket.Flow

	tcpstate *reassembly.TCPSimpleFSM
}

func (s *myStream) Accept(tcp *layers.TCP, ci gopacket.CaptureInfo, dir reassembly.TCPFlowDirection, ackSeq reassembly.Sequence, start *bool, ac reassembly.AssemblerContext) bool {
	if !s.tcpstate.CheckState(tcp, dir) {
		log.Printf("stream #%d: Accept rejected by FSM", s.i)
		return false
	}

	log.Printf("stream #%d: Accept", s.i)
	return true
}

func (s *myStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.AssemblerContext) {
	log.Printf("stream #%d: ReassembledSG", s.i)
	if s.b {
		log.Fatalf("Got ReassembledSG after ReassemblyComplete %s %s!", s.NetFlow, s.TCPFlow)
	}
}

func (s *myStream) ReassemblyComplete(ac reassembly.AssemblerContext) bool {
	log.Printf("stream #%d: ReassemblyComplete", s.i)
	if s.b {
		log.Fatalf("Got ReassemblyComplete after ReassemblyComplete %s %s!", s.NetFlow, s.TCPFlow)
	}
	s.b = true
	return false // XXX?
}

type Context struct {
	CaptureInfo gopacket.CaptureInfo
}

func (c *Context) GetCaptureInfo() gopacket.CaptureInfo {
	return c.CaptureInfo
}

func main() {
	file, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}

	pcap, err := pcapgo.NewReader(bufio.NewReader(file))
	if err != nil {
		log.Fatal(err)
	}

	packetSource := gopacket.NewPacketSource(pcap, pcap.LinkType())
	packets := packetSource.Packets()

	assembler := NewAssembler()

	if false {
		log.SetOutput(ioutil.Discard)
	}

	var lastPacketTs time.Time

	var packetNo int64
	packetNo = 1
	for packet := range packets {
		if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
			continue
		}
		tcp := packet.TransportLayer().(*layers.TCP)
		log.Printf("packet %d: calling AssembleWithTimestamp()", packetNo)
		packetNo++

		assembler.AssembleWithContext(packet.NetworkLayer().NetworkFlow(), tcp, &Context{packet.Metadata().CaptureInfo})

		lastPacketTs = packet.Metadata().Timestamp
	}

	log.Print("Calling FlushCloseOlderThan")
	assembler.FlushCloseOlderThan(lastPacketTs.Add(time.Minute * 1))

	log.Print("calling FlushAll")

	assembler.FlushAll()
}