Build PCAP file with atomic blocking writes

The previous use of non-blocking writes could result in corrupt PCAP files if a
partial write occurs. Using (*os.File).Write solves this problem by not
allowing partial writes. This change does not increase allocations (in one path
it actually reduces them), but does add additional copying.

PiperOrigin-RevId: 209652974
Change-Id: I4b1cf2eda4cfd7f237a4245aceb7391b3055a66c

3 files changed, 42 insertions, 49 deletions

diff --git a/pkg/tcpip/link/rawfile/rawfile_unsafe.go b/pkg/tcpip/link/rawfile/rawfile_unsafe.go
index 261d350d7..cea3cd6a1 100644
--- a/pkg/tcpip/link/rawfile/rawfile_unsafe.go
+++ b/pkg/tcpip/link/rawfile/rawfile_unsafe.go
@@ -94,29 +94,6 @@ func NonBlockingWrite2(fd int, b1, b2 []byte) *tcpip.Error {
 	return nil
 }
 
-// NonBlockingWriteN writes up to N byte slices to a file descriptor in a
-// single syscall. It fails if partial data is written.
-func NonBlockingWriteN(fd int, bs ...[]byte) *tcpip.Error {
-	iovec := make([]syscall.Iovec, 0, len(bs))
-
-	for _, b := range bs {
-		if len(b) == 0 {
-			continue
-		}
-		iovec = append(iovec, syscall.Iovec{
-			Base: &b[0],
-			Len:  uint64(len(b)),
-		})
-	}
-
-	_, _, e := syscall.RawSyscall(syscall.SYS_WRITEV, uintptr(fd), uintptr(unsafe.Pointer(&iovec[0])), uintptr(len(iovec)))
-	if e != 0 {
-		return TranslateErrno(e)
-	}
-
-	return nil
-}
-
 type pollEvent struct {
 	fd      int32
 	events  int16
diff --git a/pkg/tcpip/link/sniffer/BUILD b/pkg/tcpip/link/sniffer/BUILD
index 1e844f949..7155aea66 100644
--- a/pkg/tcpip/link/sniffer/BUILD
+++ b/pkg/tcpip/link/sniffer/BUILD
@@ -17,7 +17,6 @@ go_library(
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
-        "//pkg/tcpip/link/rawfile",
         "//pkg/tcpip/stack",
     ],
 )
diff --git a/pkg/tcpip/link/sniffer/sniffer.go b/pkg/tcpip/link/sniffer/sniffer.go
index 22f009751..3836cebb7 100644
--- a/pkg/tcpip/link/sniffer/sniffer.go
+++ b/pkg/tcpip/link/sniffer/sniffer.go
@@ -33,7 +33,6 @@ import (
 	"gvisor.googlesource.com/gvisor/pkg/tcpip"
 	"gvisor.googlesource.com/gvisor/pkg/tcpip/buffer"
 	"gvisor.googlesource.com/gvisor/pkg/tcpip/header"
-	"gvisor.googlesource.com/gvisor/pkg/tcpip/link/rawfile"
 	"gvisor.googlesource.com/gvisor/pkg/tcpip/stack"
 )
 
@@ -123,22 +122,28 @@ func (e *endpoint) DeliverNetworkPacket(linkEP stack.LinkEndpoint, remoteLinkAdd
 	}
 	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
 		vs := vv.Views()
-		bs := make([][]byte, 1, 1+len(vs))
-		var length int
+		length := vv.Size()
+		if length > int(e.maxPCAPLen) {
+			length = int(e.maxPCAPLen)
+		}
+
+		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen+length))
+		if err := binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(vv.Size()))); err != nil {
+			panic(err)
+		}
 		for _, v := range vs {
-			if length+len(v) > int(e.maxPCAPLen) {
-				l := int(e.maxPCAPLen) - length
-				bs = append(bs, []byte(v)[:l])
-				length += l
+			if length == 0 {
 				break
 			}
-			bs = append(bs, []byte(v))
-			length += len(v)
+			if len(v) > length {
+				v = v[:length]
+			}
+			if _, err := buf.Write([]byte(v)); err != nil {
+				panic(err)
+			}
+			length -= len(v)
 		}
-		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
-		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(vv.Size())))
-		bs[0] = buf.Bytes()
-		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
+		if _, err := e.file.Write(buf.Bytes()); err != nil {
 			panic(err)
 		}
 	}
@@ -188,21 +193,33 @@ func (e *endpoint) WritePacket(r *stack.Route, hdr *buffer.Prependable, payload
 		LogPacket("send", protocol, hdr.UsedBytes(), payload)
 	}
 	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
-		bs := [][]byte{nil, hdr.UsedBytes(), payload}
-		var length int
+		hdrBuf := hdr.UsedBytes()
+		length := len(hdrBuf) + len(payload)
+		if length > int(e.maxPCAPLen) {
+			length = int(e.maxPCAPLen)
+		}
 
-		for i, b := range bs[1:] {
-			if rem := int(e.maxPCAPLen) - length; len(b) > rem {
-				b = b[:rem]
+		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen+length))
+		if err := binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(hdr.UsedLength()+len(payload)))); err != nil {
+			panic(err)
+		}
+		if len(hdrBuf) > length {
+			hdrBuf = hdrBuf[:length]
+		}
+		if _, err := buf.Write(hdrBuf); err != nil {
+			panic(err)
+		}
+		length -= len(hdrBuf)
+		if length > 0 {
+			p := payload
+			if len(p) > length {
+				p = p[:length]
+			}
+			if _, err := buf.Write(p); err != nil {
+				panic(err)
 			}
-			bs[i+1] = b
-			length += len(b)
 		}
-
-		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
-		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(hdr.UsedLength()+len(payload))))
-		bs[0] = buf.Bytes()
-		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
+		if _, err := e.file.Write(buf.Bytes()); err != nil {
 			panic(err)
 		}
 	}