Fix loss of state, and a thorough test

This commit is contained in:
Neale Pickett 2020-09-24 20:23:43 -06:00
parent b9f024161f
commit 3a6c62985c
4 changed files with 32 additions and 24 deletions

View File

@ -2,13 +2,14 @@ package main
import (
"fmt"
"github.com/dirtbags/netshovel"
"github.com/google/gopacket"
"github.com/google/gopacket/tcpassembly"
"io"
"log"
"strings"
"sync"
"github.com/dirtbags/netshovel"
"github.com/google/gopacket"
"github.com/google/gopacket/tcpassembly"
)
var wg sync.WaitGroup
@ -17,7 +18,7 @@ type SimpleStreamFactory struct {
}
type SimpleStream struct {
netshovel.Stream
*netshovel.Stream
}
type SimplePacket struct {

View File

@ -31,7 +31,16 @@ func Shovel(factory tcpassembly.StreamFactory) {
assembler := tcpassembly.NewAssembler(streamPool)
for _, fn := range flag.Args() {
handle, err := pcap.OpenOffline(fn)
ShovelFile(fn, assembler)
}
assembler.FlushAll()
}
// ShovelFile shovels a single file.
// You must call assembler.FlushAll() at the end of this!
func ShovelFile(filename string, assembler *tcpassembly.Assembler) {
handle, err := pcap.OpenOffline(filename)
if err != nil {
log.Fatal(err)
}
@ -48,6 +57,4 @@ func Shovel(factory tcpassembly.StreamFactory) {
tcp := packet.TransportLayer().(*layers.TCP)
assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
}
}
assembler.FlushAll()
}

View File

@ -36,7 +36,7 @@ func TestHeaders(t *testing.T) {
t.Error("Uint8", fnord)
}
biggun, err := pkt.Uint32LE("biggun")
biggun, err := pkt.Uint32BE("biggun")
if err != nil {
t.Error(err)
}

View File

@ -43,8 +43,8 @@ type Stream struct {
//
// You should embed Stream into your own Application protocol stream struct.
// Use this to initialize the internal stuff netshovel needs.
func NewStream(net, transport gopacket.Flow) Stream {
return Stream{
func NewStream(net, transport gopacket.Flow) *Stream {
return &Stream{
Net: net,
Transport: transport,
conversation: make(chan Utterance, 100),