From c6c3c2d98c058a5af5e561e889e9409d652e090b Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Fri, 22 Nov 2024 15:40:21 -0800
Subject: [PATCH] netstack: check the TCP write list for validity before
 increasing SND.NXT

Writing a segment not actually in the write list could be causing the
SND.NXT/writeList out of sync bug. This will panic early with a useful stack
trace if that's the case.

PiperOrigin-RevId: 699308767
---
 pkg/tcpip/transport/tcp/endpoint.go |  2 +-
 pkg/tcpip/transport/tcp/snd.go      | 60 ++++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index 9752cd86a..5a66ec0fe 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -2462,7 +2462,7 @@ func (e *Endpoint) connect(addr tcpip.FullAddress, handshake bool) tcpip.Error {
 	// connection setting here.
 	if !handshake {
 		e.segmentQueue.mu.Lock()
-		for _, l := range []segmentList{e.segmentQueue.list, e.snd.writeList} {
+		for _, l := range []segmentList{e.segmentQueue.list, e.snd.writeList.writeList} {
 			for s := l.Front(); s != nil; s = s.Next() {
 				s.id = e.TransportEndpointInfo.ID
 			}
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index f9fe69625..5b75af69d 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -124,7 +124,7 @@ type sender struct {
 	// writeList holds all writable data: both unsent data and
 	// sent-but-unacknowledged data. Alternatively: it holds all bytes
 	// starting from SND.UNA.
-	writeList segmentList
+	writeList protectedWriteList
 
 	// resendTimer is used for RTOs.
 	resendTimer timer `state:"nosave"`
@@ -180,6 +180,54 @@ type sender struct {
 	corkTimer timer `state:"nosave"`
 }
 
+// protectedWriteList wraps the write list, checking for invalid state when
+// segments are added or removed.
+//
+// TODO(b/339664055): Revert once bug is fixed.
+//
+// +stateify savable
+type protectedWriteList struct {
+	writeList segmentList
+	set       map[*segment]struct{}
+}
+
+// Front returns the front of the write list.
+func (wl *protectedWriteList) Front() *segment {
+	return wl.writeList.Front()
+}
+
+// Back returns the back of the write list.
+func (wl *protectedWriteList) Back() *segment {
+	return wl.writeList.Back()
+}
+
+// Remove removes seg from the write list.
+func (wl *protectedWriteList) Remove(seg *segment) {
+	if _, ok := wl.set[seg]; !ok {
+		panic("segment not found write list")
+	}
+	wl.writeList.Remove(seg)
+	delete(wl.set, seg)
+}
+
+// PushBack pushes seg onto the back of the write list.
+func (wl *protectedWriteList) PushBack(seg *segment) {
+	if _, ok := wl.set[seg]; ok {
+		panic("segment already in write list")
+	}
+	wl.writeList.PushBack(seg)
+	wl.set[seg] = struct{}{}
+}
+
+// InsertAfter inserts seg after before.
+func (wl *protectedWriteList) InsertAfter(before, seg *segment) {
+	if _, ok := wl.set[seg]; ok {
+		panic("segment already in write list")
+	}
+	wl.writeList.InsertAfter(before, seg)
+	wl.set[seg] = struct{}{}
+}
+
 // rtt is a synchronization wrapper used to appease stateify. See the comment
 // in sender, where it is used.
 //
@@ -216,6 +264,9 @@ func newSender(ep *Endpoint, iss, irs seqnum.Value, sndWnd seqnum.Size, mss uint
 			RTO: 1 * time.Second,
 		},
 		gso: ep.gso.Type != stack.GSONone,
+		writeList: protectedWriteList{
+			set: make(map[*segment]struct{}),
+		},
 	}
 
 	if s.gso {
@@ -912,6 +963,13 @@ func (s *sender) maybeSendSegment(seg *segment, limit int, end seqnum.Value) (se
 		segEnd = seg.sequenceNumber.Add(seqnum.Size(seg.payloadSize()))
 	}
 
+	// TODO(b/379932042): Below is the only place we update SND.NXT besides
+	// initialization. It's possible that we're increasing SND.NXT by
+	// trying to write a segment that isn't in the write list.
+	if _, ok := s.writeList.set[seg]; !ok {
+		panic("attempted to send segment not in write list")
+	}
+
 	s.sendSegment(seg)
 
 	// Update sndNxt if we actually sent new data (as opposed to