2 files changed, 46 insertions, 43 deletions
diff --git a/pkg/tcpip/link/sniffer/pcap.go b/pkg/tcpip/link/sniffer/pcap.go
index 3bb864ed2..d3edede63 100644
--- a/pkg/tcpip/link/sniffer/pcap.go
+++ b/pkg/tcpip/link/sniffer/pcap.go
@@ -14,7 +14,14 @@
 
 package sniffer
 
-import "time"
+import (
+	"encoding"
+	"encoding/binary"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
+)
 
 type pcapHeader struct {
 	// MagicNumber is the file magic number.
@@ -39,25 +46,38 @@ type pcapHeader struct {
 	Network uint32
 }
 
-type pcapPacketHeader struct {
-	// Seconds is the timestamp seconds.
-	Seconds uint32
-
-	// Microseconds is the timestamp microseconds.
-	Microseconds uint32
+var _ encoding.BinaryMarshaler = (*pcapPacket)(nil)
 
-	// IncludedLength is the number of octets of packet saved in file.
-	IncludedLength uint32
-
-	// OriginalLength is the actual length of packet.
-	OriginalLength uint32
+type pcapPacket struct {
+	timestamp     time.Time
+	packet        *stack.PacketBuffer
+	maxCaptureLen int
 }
 
-func newPCAPPacketHeader(now time.Time, incLen, orgLen uint32) pcapPacketHeader {
-	return pcapPacketHeader{
-		Seconds:        uint32(now.Unix()),
-		Microseconds:   uint32(now.Nanosecond() / 1000),
-		IncludedLength: incLen,
-		OriginalLength: orgLen,
+func (p *pcapPacket) MarshalBinary() ([]byte, error) {
+	packetSize := p.packet.Size()
+	captureLen := p.maxCaptureLen
+	if packetSize < captureLen {
+		captureLen = packetSize
+	}
+	b := make([]byte, 16+captureLen)
+	binary.BigEndian.PutUint32(b[0:4], uint32(p.timestamp.Unix()))
+	binary.BigEndian.PutUint32(b[4:8], uint32(p.timestamp.Nanosecond()/1000))
+	binary.BigEndian.PutUint32(b[8:12], uint32(captureLen))
+	binary.BigEndian.PutUint32(b[12:16], uint32(packetSize))
+	w := tcpip.SliceWriter(b[16:])
+	for _, v := range p.packet.Views() {
+		if captureLen == 0 {
+			break
+		}
+		if len(v) > captureLen {
+			v = v[:captureLen]
+		}
+		n, err := w.Write(v)
+		if err != nil {
+			panic(err)
+		}
+		captureLen -= n
 	}
+	return b, nil
 }
diff --git a/pkg/tcpip/link/sniffer/sniffer.go b/pkg/tcpip/link/sniffer/sniffer.go
index 3df826f3c..28a172e71 100644
--- a/pkg/tcpip/link/sniffer/sniffer.go
+++ b/pkg/tcpip/link/sniffer/sniffer.go
@@ -151,33 +151,16 @@ func (e *endpoint) dumpPacket(dir direction, protocol tcpip.NetworkProtocolNumbe
 		logPacket(e.logPrefix, dir, protocol, pkt)
 	}
 	if writer != nil && atomic.LoadUint32(&LogPacketsToPCAP) == 1 {
-		totalLength := pkt.Size()
-		length := totalLength
-		if max := int(e.maxPCAPLen); length > max {
-			length = max
+		packet := pcapPacket{
+			timestamp:     time.Now(),
+			packet:        pkt,
+			maxCaptureLen: int(e.maxPCAPLen),
 		}
-		packetHeader := newPCAPPacketHeader(time.Now(), uint32(length), uint32(totalLength))
-		packet := make([]byte, binary.Size(packetHeader)+length)
-		{
-			writer := tcpip.SliceWriter(packet)
-			if err := binary.Write(&writer, binary.BigEndian, packetHeader); err != nil {
-				panic(err)
-			}
-			for _, b := range pkt.Views() {
-				if length == 0 {
-					break
-				}
-				if len(b) > length {
-					b = b[:length]
-				}
-				n, err := writer.Write(b)
-				if err != nil {
-					panic(err)
-				}
-				length -= n
-			}
+		b, err := packet.MarshalBinary()
+		if err != nil {
+			panic(err)
 		}
-		if _, err := writer.Write(packet); err != nil {
+		if _, err := writer.Write(b); err != nil {
 			panic(err)
 		}
 	}