From 3a6c62985c408824cc66a86d99666ece14f35760 Mon Sep 17 00:00:00 2001
From: Neale Pickett <neale@lanl.gov>
Date: Thu, 24 Sep 2020 20:23:43 -0600
Subject: [PATCH] Fix loss of state, and a thorough test

---
 examples/simple/simple.go |  9 +++++----
 netshovel.go              | 41 +++++++++++++++++++++++----------------
 packet_test.go            |  2 +-
 stream.go                 |  4 ++--
 4 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/examples/simple/simple.go b/examples/simple/simple.go
index 3afd052..699ec97 100644
--- a/examples/simple/simple.go
+++ b/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 {
diff --git a/netshovel.go b/netshovel.go
index 6309919..48a5366 100644
--- a/netshovel.go
+++ b/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)
-		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)
-		}
+		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)
+	}
+
+	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)
+	}
+}
diff --git a/packet_test.go b/packet_test.go
index 4d2465f..85b9180 100644
--- a/packet_test.go
+++ b/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)
 	}
diff --git a/stream.go b/stream.go
index 17a2533..6eb8189 100644
--- a/stream.go
+++ b/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 {
-	return Stream{
+func NewStream(net, transport gopacket.Flow) *Stream {
+	return &Stream{
 		Net:          net,
 		Transport:    transport,
 		conversation: make(chan Utterance, 100),