Merge pull request #31 from fako1024/30-support-for-gzip-compressed-f…

…iles-in-pcap-source

[feature] Support for gzip compressed files in pcap source

capture/pcap/pcap.go

@@ -1,6 +1,8 @@
 package pcap
 
 import (
+	"bufio"
+	"compress/gzip"
 	"errors"
 	"fmt"
 	"io"
@@ -15,7 +17,8 @@ import (
 
 // Source denotes a pcap file capture source
 type Source struct {
-	io.Reader
+	reader     *bufio.Reader
+	gzipReader *gzip.Reader
 
 	header        Header
 	buf           []byte
@@ -35,19 +38,25 @@ func NewSource(iface string, r io.Reader) (*Source, error) {
 	if r == nil {
 		return nil, errors.New("nil io.Reader provided")
 	}
+
 	obj := Source{
-		Reader: r,
+		reader: bufio.NewReader(r),
 		buf:    make([]byte, HeaderSize),
 	}
 
+	// Check if the source is compressed
+	if err := obj.checkCompression(); err != nil {
+		return nil, err
+	}
+
 	// Parse the main header
 	if err := obj.read(obj.buf); err != nil {
 		return nil, err
 	}
 
 	// If required, swap endianess as defined here:
 	// https://wiki.wireshark.org/Development/LibpcapFileFormat
-	obj.header = *(*Header)(unsafe.Pointer(&obj.buf[0]))
+	obj.header = *(*Header)(unsafe.Pointer(&obj.buf[0])) // #nosec G103
 	if obj.header.MagicNumber == MagicSwappedEndianess {
 		obj.header = obj.header.SwapEndianess()
 		obj.swapEndianess = true
@@ -163,8 +172,8 @@ func (s *Source) Unblock() error {
 
 // Close stops / closes the capture source
 func (s *Source) Close() error {
-	if readCloser, ok := s.Reader.(io.ReadCloser); ok {
-		return readCloser.Close()
+	if s.gzipReader != nil {
+		return s.gzipReader.Close()
 	}
 	return nil
 }
@@ -184,6 +193,30 @@ func (s *Source) PacketAddCallbackFn(fn func(payload []byte, totalLen uint32, pk
 
 ////////////////////////////////////////////////////////////////////////
 
+func (s *Source) checkCompression() error {
+
+	// Check if the first two bytes match a gzip file magic in either
+	// endianess
+	magicBytes, err := s.reader.Peek(2)
+	if err != nil {
+		return err
+	}
+	if (magicBytes[0] == 0x1f && magicBytes[1] == 0x8b) ||
+		(magicBytes[0] == 0x8b && magicBytes[1] == 0x1f) {
+
+		// Attempt to open a new gzip decompressor
+		s.gzipReader, err = gzip.NewReader(s.reader)
+		if err != nil {
+			return err
+		}
+
+		// Wrap a new bufio.Reader around the gzip decompressor
+		s.reader = bufio.NewReader(s.gzipReader)
+	}
+
+	return nil
+}
+
 func (s *Source) nextPacket() (pktHeader PacketHeader, err error) {
 	pktHeader, err = s.nextPacketHeader()
 	if err != nil {
@@ -209,7 +242,7 @@ func (s *Source) nextPacketHeader() (PacketHeader, error) {
 	if err := s.read(s.buf[:PacketHeaderSize]); err != nil {
 		return PacketHeader{}, err
 	}
-	return *(*PacketHeader)(unsafe.Pointer(&s.buf[0])), nil
+	return *(*PacketHeader)(unsafe.Pointer(&s.buf[0])), nil // #nosec G103
 }
 
 func (s *Source) nextPacketData(snapLen int) error {
@@ -222,7 +255,7 @@ func (s *Source) nextPacketData(snapLen int) error {
 }
 
 func (s *Source) read(buf []byte) error {
-	n, err := io.ReadAtLeast(s.Reader, buf, len(buf))
+	n, err := io.ReadAtLeast(s.reader, buf, len(buf))
 	if err != nil {
 		return err
 	}

capture/pcap/pcap_test.go

@@ -2,10 +2,12 @@ package pcap
 
 import (
 	"bytes"
-	"encoding/base64"
+	"embed"
 	"io"
+	"io/fs"
 	"net"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/fako1024/slimcap/capture"
@@ -14,7 +16,12 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-var testDataLE, testDataBE []byte
+const (
+	pcapTestInputNPackets = 6
+	testDataPath          = "testdata"
+)
+
+var testData []fs.DirEntry
 
 func TestInvalidInput(t *testing.T) {
 
@@ -53,12 +60,14 @@ func TestInvalidInput(t *testing.T) {
 
 func TestReader(t *testing.T) {
 
-	for _, testData := range [][]byte{
-		testDataLE,
-		testDataBE,
-	} {
+	for _, dirent := range testData {
+		file, err := os.Open(filepath.Join(testDataPath, dirent.Name()))
+		require.Nil(t, err)
+		defer func() {
+			require.Nil(t, file.Close())
+		}()
 
-		src, err := NewSource("pcap", bytes.NewBuffer(testData))
+		src, err := NewSource("pcap", file)
 		require.Nil(t, err)
 
 		require.Equal(t, &link.Link{
@@ -86,7 +95,7 @@ func TestReader(t *testing.T) {
 		require.Equal(t, capture.PacketUnknown, pktType)
 		require.Nil(t, ipLayer)
 
-		src.NextPacketFn(func(payload []byte, totalLen uint32, pktType, ipLayerOffset byte) error {
+		err = src.NextPacketFn(func(payload []byte, totalLen uint32, pktType, ipLayerOffset byte) error {
 			require.Nil(t, payload)
 			require.Zero(t, totalLen)
 			require.Equal(t, capture.PacketUnknown, pktType)
@@ -107,11 +116,14 @@ func TestReader(t *testing.T) {
 
 func TestMockPipe(t *testing.T) {
 
-	for _, testData := range [][]byte{
-		testDataLE,
-		testDataBE,
-	} {
-		src, err := NewSource("pcap", bytes.NewBuffer(testData))
+	for _, dirent := range testData {
+		file, err := os.Open(filepath.Join(testDataPath, dirent.Name()))
+		require.Nil(t, err)
+		defer func() {
+			require.Nil(t, file.Close())
+		}()
+
+		src, err := NewSource("pcap", file)
 		require.Nil(t, err)
 
 		// Setup a mock source
@@ -232,12 +244,14 @@ func TestCaptureMethods(t *testing.T) {
 }
 
 func testCaptureMethods(t *testing.T, fn func(t *testing.T, src *Source)) {
+	for _, dirent := range testData {
+		file, err := os.Open(filepath.Join(testDataPath, dirent.Name()))
+		require.Nil(t, err)
+		defer func() {
+			require.Nil(t, file.Close())
+		}()
 
-	for _, testData := range [][]byte{
-		testDataLE,
-		testDataBE,
-	} {
-		src, err := NewSource("pcap", bytes.NewBuffer(testData))
+		src, err := NewSource("pcap", file)
 		require.Nil(t, err)
 
 		for i := 0; i < pcapTestInputNPackets; i++ {
@@ -248,13 +262,13 @@ func testCaptureMethods(t *testing.T, fn func(t *testing.T, src *Source)) {
 	}
 }
 
+//go:embed testdata/*
+var pcaps embed.FS
+
 func TestMain(m *testing.M) {
 
 	var err error
-	if testDataLE, err = base64.StdEncoding.DecodeString(pcapTestInputLE); err != nil {
-		panic(err)
-	}
-	if testDataBE, err = base64.StdEncoding.DecodeString(pcapTestInputBE); err != nil {
+	if testData, err = pcaps.ReadDir(testDataPath); err != nil {
 		panic(err)
 	}