Add ability to concat finite packet sources.

Author: gconnell

packet.go

@@ -684,6 +684,31 @@ type PacketDataSource interface {
 	ReadPacketData() (data []byte, ci CaptureInfo, err error)
 }
 
+// ConcatFinitePacketDataSources returns a PacketDataSource that wraps a set
+// of internal PacketDataSources, each of which will stop with io.EOF after
+// reading a finite number of packets.  The returned PacketDataSource will
+// return all packets from the first finite source, followed by all packets from
+// the second, etc.  Once all finite sources have returned io.EOF, the returned
+// source will as well.
+func ConcatFinitePacketDataSources(pds ...PacketDataSource) PacketDataSource {
+	c := concat(pds)
+	return &c
+}
+
+type concat []PacketDataSource
+
+func (c *concat) ReadPacketData() (data []byte, ci CaptureInfo, err error) {
+	for len(*c) > 0 {
+		data, ci, err = (*c)[0].ReadPacketData()
+		if err == io.EOF {
+			*c = (*c)[1:]
+			continue
+		}
+		return
+	}
+	return nil, CaptureInfo{}, io.EOF
+}
+
 // ZeroCopyPacketDataSource is an interface to pull packet data from sources
 // that allow data to be returned without copying to a user-controlled buffer.
 // It's very similar to PacketDataSource, except that the caller must be more

packet_test.go

@@ -7,6 +7,7 @@
 package gopacket
 
 import (
+	"io"
 	"reflect"
 	"testing"
 )
@@ -26,3 +27,36 @@ func TestDumpEmbedded(t *testing.T) {
 		t.Errorf("embedded dump mismatch:\n   got: %v\n  want: %v", got, want)
 	}
 }
+
+type singlePacketSource [1][]byte
+
+func (s *singlePacketSource) ReadPacketData() ([]byte, CaptureInfo, error) {
+	if (*s)[0] == nil {
+		return nil, CaptureInfo{}, io.EOF
+	}
+	out := (*s)[0]
+	(*s)[0] = nil
+	return out, CaptureInfo{}, nil
+}
+
+func TestConcatPacketSources(t *testing.T) {
+	sourceA := &singlePacketSource{[]byte{1}}
+	sourceB := &singlePacketSource{[]byte{2}}
+	sourceC := &singlePacketSource{[]byte{3}}
+	concat := ConcatFinitePacketDataSources(sourceA, sourceB, sourceC)
+	a, _, err := concat.ReadPacketData()
+	if err != nil || len(a) != 1 || a[0] != 1 {
+		t.Errorf("expected [1], got %v/%v", a, err)
+	}
+	b, _, err := concat.ReadPacketData()
+	if err != nil || len(b) != 1 || b[0] != 2 {
+		t.Errorf("expected [2], got %v/%v", b, err)
+	}
+	c, _, err := concat.ReadPacketData()
+	if err != nil || len(c) != 1 || c[0] != 3 {
+		t.Errorf("expected [3], got %v/%v", c, err)
+	}
+	if _, _, err := concat.ReadPacketData(); err != io.EOF {
+		t.Errorf("expected io.EOF, got %v", err)
+	}
+}