Fix loss of state, and a thorough test

examples/simple/simple.go

@@ -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 {

netshovel.go

@@ -31,23 +31,30 @@ func Shovel(factory tcpassembly.StreamFactory) {
 	assembler := tcpassembly.NewAssembler(streamPool)
 
 	for _, fn := range flag.Args() {
-		handle, err := pcap.OpenOffline(fn)
+		ShovelFile(fn, assembler)
-		if err != nil {
-			log.Fatal(err)
-		}
-
-		packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
-		packets := packetSource.Packets()
-		for packet := range packets {
-			if packet == nil {
-				break
-			}
-			if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
-				continue
-			}
-			tcp := packet.TransportLayer().(*layers.TCP)
-			assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
-		}
 	}
+
 	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)
+	}
+
+	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
+	packets := packetSource.Packets()
+	for packet := range packets {
+		if packet == nil {
+			break
+		}
+		if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
+			continue
+		}
+		tcp := packet.TransportLayer().(*layers.TCP)
+		assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
+	}
+}

packet_test.go

@@ -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)
 	}

stream.go

@@ -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 {
+func NewStream(net, transport gopacket.Flow) *Stream {
-	return Stream{
+	return &Stream{
 		Net:          net,
 		Transport:    transport,
 		conversation: make(chan Utterance, 100),