streaming

commit f7c6ae15d84416d1c6a4d34c1c5a35a35ecbb612
parent cb9f53840ce84710f2da783c0f4efc3a96479ba2
Author: Oliver Lowe <o@olowe.co>
Date:   Sat, 17 Aug 2024 11:06:12 +1000

pcap: export Encode, Decode functions

Use Go's io.Reader and io.Writer so that we don't have to keep
everything in memory, and it works in a way that Go programmers are
familiar with.

Diffstat:
M
pcap/pcap.go
 | 
35
+++++++++++++++++++++++------------
M
pcap/pcap_test.go
 | 
14
+++++++-------

2 files changed, 30 insertions(+), 19 deletions(-)
diff --git a/pcap/pcap.go b/pcap/pcap.go
@@ -48,9 +48,10 @@ type File struct {
 	Packets []Packet
 }
 
-func decode(reader io.Reader) (*File, error) {
+// Decode returns a decoded savefile packet capture from rd.
+func Decode(rd io.Reader) (*File, error) {
 	var magic uint32
-	if err := binary.Read(reader, binary.NativeEndian, &magic); err != nil {
+	if err := binary.Read(rd, binary.NativeEndian, &magic); err != nil {
 		return nil, fmt.Errorf("read magic number: %w", err)
 	}
 	if magic != magicLittleEndian && magic != magicBigEndian {
@@ -58,7 +59,7 @@ func decode(reader io.Reader) (*File, error) {
 	}
 
 	var v [2]uint16
-	if err := binary.Read(reader, binary.LittleEndian, &v); err != nil {
+	if err := binary.Read(rd, binary.LittleEndian, &v); err != nil {
 		return nil, fmt.Errorf("read pcap version: %w", err)
 	}
 	if v != version {
@@ -66,14 +67,14 @@ func decode(reader io.Reader) (*File, error) {
 	}
 
 	var gheader GlobalHeader
-	if err := binary.Read(reader, binary.LittleEndian, &gheader); err != nil {
+	if err := binary.Read(rd, binary.LittleEndian, &gheader); err != nil {
 		return nil, fmt.Errorf("read pcap version: %w", err)
 	}
 
 	var packets []Packet
 	for i := 1; ; i++ {
 		var h header
-		err := binary.Read(reader, binary.LittleEndian, &h)
+		err := binary.Read(rd, binary.LittleEndian, &h)
 		if err == io.EOF {
 			break
 		} else if err != nil {
@@ -85,7 +86,7 @@ func decode(reader io.Reader) (*File, error) {
 		}
 
 		data := make([]byte, h.InclLen)
-		if _, err = io.ReadFull(reader, data); err != nil {
+		if _, err = io.ReadFull(rd, data); err != nil {
 			return nil, fmt.Errorf("packet %d: read data: %w", i, err)
 		}
 
@@ -101,16 +102,22 @@ func decode(reader io.Reader) (*File, error) {
 	}, nil
 }
 
-func encode(file *File) ([]byte, error) {
+// Encode writes a savefile-encoded representation of file to w.
+func Encode(w io.Writer, file *File) (n int64, err error) {
 	b := make([]byte, 4+4) // magic + version
 	binary.NativeEndian.PutUint32(b, magicLittleEndian)
 	binary.NativeEndian.PutUint16(b[4:6], version[0])
 	binary.NativeEndian.PutUint16(b[6:8], version[1])
 	buf := bytes.NewBuffer(b)
 	if err := binary.Write(buf, binary.LittleEndian, &file.Header); err != nil {
-		return nil, fmt.Errorf("global header: %v", err)
+		return n, fmt.Errorf("global header: %v", err)
+	}
+	n, err = io.Copy(w, buf)
+	if err != nil {
+		return n, fmt.Errorf("global header: %w", err)
 	}
 
+	buf.Reset()
 	for i, p := range file.Packets {
 		sec, nsec := timestamp(p.Header.Time)
 		h := header{
@@ -120,14 +127,18 @@ func encode(file *File) ([]byte, error) {
 			OrigLen:    p.Header.OrigLen,
 		}
 		if err := binary.Write(buf, binary.LittleEndian, h); err != nil {
-			return nil, fmt.Errorf("packet %d: header: %v", i, err)
+			return n, fmt.Errorf("packet %d: header: %v", i, err)
 		}
 		if _, err := buf.Write(p.Data); err != nil {
-			return nil, fmt.Errorf("packet %d: data: %v", i, err)
+			return n, fmt.Errorf("packet %d: data: %v", i, err)
+		}
+		nn, err := io.Copy(w, buf)
+		n += nn
+		if err != nil {
+			return n, fmt.Errorf("packet %d: %w", i, err)
 		}
 	}
-
-	return buf.Bytes(), nil
+	return n, nil
 }
 
 func timestamp(t time.Time) (seconds, nanoSeconds uint32) {
diff --git a/pcap/pcap_test.go b/pcap/pcap_test.go
@@ -23,7 +23,7 @@ func TestDecode(t *testing.T) {
 		OrigLen: 45,
 	}
 
-	capture, err := decode(f)
+	capture, err := Decode(f)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -44,16 +44,16 @@ func TestEncode(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	capture, err := decode(bytes.NewReader(want))
+	capture, err := Decode(bytes.NewReader(want))
 	if err != nil {
 		t.Fatal(err)
 	}
-	got, err := encode(capture)
-	if err != nil {
-		t.Fatalf("encode: %v", err)
+	buf := &bytes.Buffer{}
+	if _, err := Encode(buf, capture); err != nil {
+		t.Fatal(err)
 	}
-	if !bytes.Equal(got, want) {
-		t.Errorf("encode(%v) = %x, want %x", capture, got, want)
+	if !bytes.Equal(buf.Bytes(), want) {
+		t.Errorf("Encode(%v) = %x, want %x", capture, buf.Bytes(), want)
 	}
 }