From 28cf71cc61c58f07554cab2b3724d2cb3c560b21 Mon Sep 17 00:00:00 2001
From: Lucas Manning <lucasmanning@google.com>
Date: Thu, 24 Mar 2022 16:53:40 -0700
Subject: [PATCH] Refactor the tcp segment implementation to use refsvfs.

This change also fixes existing refcounting correctness issues.

PiperOrigin-RevId: 437118979
---
 pkg/tcpip/stack/transport_test.go             |  2 +
 pkg/tcpip/tcpip.go                            | 15 +++
 pkg/tcpip/transport/icmp/endpoint.go          |  3 +
 pkg/tcpip/transport/internal/noop/endpoint.go |  5 +
 pkg/tcpip/transport/packet/endpoint.go        |  3 +
 pkg/tcpip/transport/raw/endpoint.go           |  3 +
 pkg/tcpip/transport/tcp/BUILD                 | 13 +++
 pkg/tcpip/transport/tcp/accept.go             |  5 +-
 pkg/tcpip/transport/tcp/connect.go            | 26 +++---
 pkg/tcpip/transport/tcp/dispatcher.go         |  4 +-
 pkg/tcpip/transport/tcp/endpoint.go           | 91 ++++++++++++++++++-
 pkg/tcpip/transport/tcp/forwarder.go          |  6 +-
 pkg/tcpip/transport/tcp/protocol.go           |  2 +-
 pkg/tcpip/transport/tcp/rack.go               |  2 +-
 pkg/tcpip/transport/tcp/rcv.go                |  7 +-
 pkg/tcpip/transport/tcp/sack_recovery.go      |  2 +-
 pkg/tcpip/transport/tcp/segment.go            | 22 ++---
 pkg/tcpip/transport/tcp/segment_queue.go      |  1 +
 pkg/tcpip/transport/tcp/segment_test.go       |  4 +-
 pkg/tcpip/transport/tcp/snd.go                | 24 +++--
 .../transport/tcp/test/e2e/forwarder_test.go  |  2 +
 pkg/tcpip/transport/tcp/test/e2e/tcp_test.go  | 36 +++++++-
 .../transport/tcp/testing/context/context.go  |  8 ++
 pkg/tcpip/transport/udp/endpoint.go           |  2 +
 runsc/boot/loader.go                          |  3 +
 25 files changed, 237 insertions(+), 54 deletions(-)

diff --git a/pkg/tcpip/stack/transport_test.go b/pkg/tcpip/stack/transport_test.go
index 51870d03f..c817e9de0 100644
--- a/pkg/tcpip/stack/transport_test.go
+++ b/pkg/tcpip/stack/transport_test.go
@@ -78,6 +78,8 @@ func (f *fakeTransportEndpoint) Abort() {
 	f.Close()
 }
 
+func (*fakeTransportEndpoint) Release() {}
+
 func (f *fakeTransportEndpoint) Close() {
 	// TODO(gvisor.dev/issue/5153): Consider retaining the route.
 	f.route.Release()
diff --git a/pkg/tcpip/tcpip.go b/pkg/tcpip/tcpip.go
index da6546a07..ede0f5b9e 100644
--- a/pkg/tcpip/tcpip.go
+++ b/pkg/tcpip/tcpip.go
@@ -674,6 +674,9 @@ type Endpoint interface {
 	// SocketOptions returns the structure which contains all the socket
 	// level options.
 	SocketOptions() *SocketOptions
+
+	// Release releases all reference counted objects held by the endpoint.
+	Release()
 }
 
 // LinkPacketInfo holds Link layer information for a received packet.
@@ -2490,6 +2493,18 @@ func GetDanglingEndpoints() []Endpoint {
 	return es
 }
 
+// ReleaseDanglingEndpoints clears out all all reference counted objects held by
+// dangling endpoints.
+func ReleaseDanglingEndpoints() {
+	// Get the dangling endpoints first to avoid locking around Release(), which
+	// can cause a lock inversion with endpoint.mu and danglingEndpointsMu.
+	// Calling Release on a dangling endpoint that has been deleted is a noop.
+	eps := GetDanglingEndpoints()
+	for _, ep := range eps {
+		ep.Release()
+	}
+}
+
 // AddDanglingEndpoint adds a dangling endpoint.
 func AddDanglingEndpoint(e Endpoint) {
 	danglingEndpointsMu.Lock()
diff --git a/pkg/tcpip/transport/icmp/endpoint.go b/pkg/tcpip/transport/icmp/endpoint.go
index 1fcc4d0ac..f86fcc672 100644
--- a/pkg/tcpip/transport/icmp/endpoint.go
+++ b/pkg/tcpip/transport/icmp/endpoint.go
@@ -767,6 +767,9 @@ func (e *endpoint) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketB
 // HandleError implements stack.TransportEndpoint.
 func (*endpoint) HandleError(stack.TransportError, *stack.PacketBuffer) {}
 
+// Release implements stack.TransportEndpoint.
+func (*endpoint) Release() {}
+
 // State implements tcpip.Endpoint.State. The ICMP endpoint currently doesn't
 // expose internal socket state.
 func (e *endpoint) State() uint32 {
diff --git a/pkg/tcpip/transport/internal/noop/endpoint.go b/pkg/tcpip/transport/internal/noop/endpoint.go
index 443b4e416..be2adae1c 100644
--- a/pkg/tcpip/transport/internal/noop/endpoint.go
+++ b/pkg/tcpip/transport/internal/noop/endpoint.go
@@ -151,6 +151,11 @@ func (*endpoint) Wait() {
 	// No-op.
 }
 
+// Release implements stack.TransportEndpoint.Release.
+func (*endpoint) Release() {
+	// No-op.
+}
+
 // LastError implements tcpip.Endpoint.LastError.
 func (*endpoint) LastError() tcpip.Error {
 	return nil
diff --git a/pkg/tcpip/transport/packet/endpoint.go b/pkg/tcpip/transport/packet/endpoint.go
index 7afc88c1a..7fe039a46 100644
--- a/pkg/tcpip/transport/packet/endpoint.go
+++ b/pkg/tcpip/transport/packet/endpoint.go
@@ -492,6 +492,9 @@ func (ep *endpoint) Stats() tcpip.EndpointStats {
 // SetOwner implements tcpip.Endpoint.SetOwner.
 func (*endpoint) SetOwner(tcpip.PacketOwner) {}
 
+// Release implements tcpip.Release.
+func (*endpoint) Release() {}
+
 // SocketOptions implements tcpip.Endpoint.SocketOptions.
 func (ep *endpoint) SocketOptions() *tcpip.SocketOptions {
 	return &ep.ops
diff --git a/pkg/tcpip/transport/raw/endpoint.go b/pkg/tcpip/transport/raw/endpoint.go
index c13919955..15eb6a32e 100644
--- a/pkg/tcpip/transport/raw/endpoint.go
+++ b/pkg/tcpip/transport/raw/endpoint.go
@@ -735,6 +735,9 @@ func (*endpoint) LastError() tcpip.Error {
 	return nil
 }
 
+// Release implements stack.TransportEndpoint.Release.
+func (*endpoint) Release() {}
+
 // SocketOptions implements tcpip.Endpoint.SocketOptions.
 func (e *endpoint) SocketOptions() *tcpip.SocketOptions {
 	return &e.ops
diff --git a/pkg/tcpip/transport/tcp/BUILD b/pkg/tcpip/transport/tcp/BUILD
index 58aeb748c..524e744f2 100644
--- a/pkg/tcpip/transport/tcp/BUILD
+++ b/pkg/tcpip/transport/tcp/BUILD
@@ -15,6 +15,17 @@ go_template_instance(
     },
 )
 
+go_template_instance(
+    name = "tcp_segment_refs",
+    out = "tcp_segment_refs.go",
+    package = "tcp",
+    prefix = "segment",
+    template = "//pkg/refsvfs2:refs_template",
+    types = {
+        "T": "segment",
+    },
+)
+
 go_template_instance(
     name = "tcp_endpoint_list",
     out = "tcp_endpoint_list.go",
@@ -54,6 +65,7 @@ go_library(
         "snd.go",
         "tcp_endpoint_list.go",
         "tcp_segment_list.go",
+        "tcp_segment_refs.go",
         "timer.go",
     ],
     imports = ["gvisor.dev/gvisor/pkg/tcpip/buffer"],
@@ -61,6 +73,7 @@ go_library(
     deps = [
         "//pkg/log",
         "//pkg/rand",
+        "//pkg/refsvfs2",
         "//pkg/sleep",
         "//pkg/sync",
         "//pkg/tcpip",
diff --git a/pkg/tcpip/transport/tcp/accept.go b/pkg/tcpip/transport/tcp/accept.go
index 39373ea13..ab61f0ed4 100644
--- a/pkg/tcpip/transport/tcp/accept.go
+++ b/pkg/tcpip/transport/tcp/accept.go
@@ -715,7 +715,6 @@ func (e *endpoint) handleListenSegment(ctx *listenContext, s *segment) tcpip.Err
 		// Requeue the segment if the ACK completing the handshake has more info
 		// to be procesed by the newly established endpoint.
 		if (s.flags.Contains(header.TCPFlagFin) || s.data.Size() > 0) && n.enqueueSegment(s) {
-			s.incRef()
 			n.newSegmentWaker.Assert()
 		}
 
@@ -780,7 +779,7 @@ func (e *endpoint) protocolListenLoop(rcvWnd seqnum.Size) {
 					// TODO(gvisor.dev/issue/4690): Better handle errors instead of
 					// silently dropping.
 					_ = e.handleListenSegment(ctx, s)
-					s.decRef()
+					s.DecRef()
 				}
 				close(e.drainDone)
 				e.mu.Unlock()
@@ -801,7 +800,7 @@ func (e *endpoint) protocolListenLoop(rcvWnd seqnum.Size) {
 				// TODO(gvisor.dev/issue/4690): Better handle errors instead of
 				// silently dropping.
 				_ = e.handleListenSegment(ctx, s)
-				s.decRef()
+				s.DecRef()
 			}
 
 			// If the queue is not empty, make sure we'll wake up
diff --git a/pkg/tcpip/transport/tcp/connect.go b/pkg/tcpip/transport/tcp/connect.go
index 0532adb7a..22e37f1a4 100644
--- a/pkg/tcpip/transport/tcp/connect.go
+++ b/pkg/tcpip/transport/tcp/connect.go
@@ -423,7 +423,6 @@ func (h *handshake) synRcvdState(s *segment) tcpip.Error {
 		// Requeue the segment if the ACK completing the handshake has more info
 		// to be procesed by the newly established endpoint.
 		if (s.flags.Contains(header.TCPFlagFin) || s.data.Size() > 0) && h.ep.enqueueSegment(s) {
-			s.incRef()
 			h.ep.newSegmentWaker.Assert()
 		}
 		return nil
@@ -459,7 +458,7 @@ func (h *handshake) processSegments() tcpip.Error {
 		}
 
 		err := h.handleSegment(s)
-		s.decRef()
+		s.DecRef()
 		if err != nil {
 			return err
 		}
@@ -590,7 +589,7 @@ func (h *handshake) complete() tcpip.Error {
 				for !h.ep.segmentQueue.empty() {
 					s := h.ep.segmentQueue.dequeue()
 					err := h.handleSegment(s)
-					s.decRef()
+					s.DecRef()
 					if err != nil {
 						return err
 					}
@@ -969,7 +968,7 @@ func (e *endpoint) sendData(next *segment) {
 		if next == nil {
 			return
 		}
-		e.snd.writeNext = next
+		e.snd.updateWriteNext(next)
 	}
 
 	// Push out any new packets.
@@ -1003,6 +1002,10 @@ func (e *endpoint) resetConnectionLocked(err tcpip.Error) {
 		}
 		e.sendRaw(buffer.VectorisedView{}, header.TCPFlagAck|header.TCPFlagRst, resetSeqNum, e.rcv.RcvNxt, 0)
 	}
+	// Don't purge read queues here. If there's buffered data, it's still allowed
+	// to be read.
+	e.purgeWriteQueue()
+	e.purgePendingRcvQueue()
 }
 
 // completeWorkerLocked is called by the worker goroutine when it's about to
@@ -1058,7 +1061,6 @@ func (e *endpoint) tryDeliverSegmentFromClosedEndpoint(s *segment) {
 	}
 	if ep == nil {
 		replyWithReset(e.stack, s, stack.DefaultTOS, tcpip.UseDefaultIPv4TTL, tcpip.UseDefaultIPv6HopLimit)
-		s.decRef()
 		return
 	}
 
@@ -1082,6 +1084,7 @@ func (e *endpoint) drainClosingSegmentQueue() {
 		}
 
 		e.tryDeliverSegmentFromClosedEndpoint(s)
+		s.DecRef()
 	}
 }
 
@@ -1150,7 +1153,7 @@ func (e *endpoint) handleSegmentsLocked(fastPath bool) tcpip.Error {
 		}
 
 		cont, err := e.handleSegmentLocked(s)
-		s.decRef()
+		s.DecRef()
 		if err != nil {
 			return err
 		}
@@ -1553,6 +1556,7 @@ func (e *endpoint) protocolMainLoop(handshake bool, wakerInitDone chan<- struct{
 		e.workerCleanup = true
 		if err != nil {
 			e.resetConnectionLocked(err)
+			e.releaseLocked()
 		}
 	}
 
@@ -1656,15 +1660,13 @@ func (e *endpoint) handleTimeWaitSegments() (extendTimeWait bool, reuseTW func()
 					if EndpointState(tcpEP.State()) == StateListen {
 						reuseTW = func() {
 							if !tcpEP.enqueueSegment(s) {
-								s.decRef()
 								return
 							}
 							tcpEP.newSegmentWaker.Assert()
+							s.DecRef()
 						}
-						// We explicitly do not decRef
-						// the segment as it's still
-						// valid and being reflected to
-						// a listening endpoint.
+						// We explicitly do not DecRef the segment as it's still valid and
+						// being reflected to a listening endpoint.
 						return false, reuseTW
 					}
 				}
@@ -1673,7 +1675,7 @@ func (e *endpoint) handleTimeWaitSegments() (extendTimeWait bool, reuseTW func()
 		if extTW {
 			extendTimeWait = true
 		}
-		s.decRef()
+		s.DecRef()
 	}
 	if checkRequeue && !e.segmentQueue.empty() {
 		e.newSegmentWaker.Assert()
diff --git a/pkg/tcpip/transport/tcp/dispatcher.go b/pkg/tcpip/transport/tcp/dispatcher.go
index 01d9275d9..a69e0e104 100644
--- a/pkg/tcpip/transport/tcp/dispatcher.go
+++ b/pkg/tcpip/transport/tcp/dispatcher.go
@@ -179,17 +179,16 @@ func (d *dispatcher) queuePacket(stackEP stack.TransportEndpoint, id stack.Trans
 	ep := stackEP.(*endpoint)
 
 	s := newIncomingSegment(id, clock, pkt)
+	defer s.DecRef()
 	if !s.parse(pkt.RXTransportChecksumValidated) {
 		ep.stack.Stats().TCP.InvalidSegmentsReceived.Increment()
 		ep.stats.ReceiveErrors.MalformedPacketsReceived.Increment()
-		s.decRef()
 		return
 	}
 
 	if !s.csumValid {
 		ep.stack.Stats().TCP.ChecksumErrors.Increment()
 		ep.stats.ReceiveErrors.ChecksumErrors.Increment()
-		s.decRef()
 		return
 	}
 
@@ -200,7 +199,6 @@ func (d *dispatcher) queuePacket(stackEP stack.TransportEndpoint, id stack.Trans
 	}
 
 	if !ep.enqueueSegment(s) {
-		s.decRef()
 		return
 	}
 
diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index cb93af58e..11fcd95d4 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -15,6 +15,7 @@
 package tcp
 
 import (
+	"container/heap"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -999,6 +1000,71 @@ func (e *endpoint) notifyProtocolGoroutine(n uint32) {
 	}
 }
 
+func (e *endpoint) Release() {
+	e.LockUser()
+	defer e.UnlockUser()
+	e.releaseLocked()
+}
+
+// +checklocks:e.mu
+func (e *endpoint) releaseLocked() {
+	e.purgeReadQueue()
+	e.purgeWriteQueue()
+	for {
+		s := e.segmentQueue.dequeue()
+		if s == nil {
+			break
+		}
+		s.DecRef()
+	}
+}
+
+// Purging pending rcv segments is only necessary on RST.
+func (e *endpoint) purgePendingRcvQueue() {
+	if e.rcv != nil {
+		for e.rcv.pendingRcvdSegments.Len() > 0 {
+			s := heap.Pop(&e.rcv.pendingRcvdSegments).(*segment)
+			s.DecRef()
+		}
+	}
+}
+
+// +checklocks:e.mu
+func (e *endpoint) purgeReadQueue() {
+	if e.rcv != nil {
+		e.rcvQueueInfo.rcvQueueMu.Lock()
+		defer e.rcvQueueInfo.rcvQueueMu.Unlock()
+		for {
+			s := e.rcvQueueInfo.rcvQueue.Front()
+			if s == nil {
+				break
+			}
+			e.rcvQueueInfo.rcvQueue.Remove(s)
+			s.DecRef()
+		}
+		e.rcvQueueInfo.RcvBufUsed = 0
+	}
+}
+
+// +checklocks:e.mu
+func (e *endpoint) purgeWriteQueue() {
+	if e.snd != nil {
+		e.sndQueueInfo.sndQueueMu.Lock()
+		defer e.sndQueueInfo.sndQueueMu.Unlock()
+		e.snd.updateWriteNext(nil)
+		for {
+			s := e.snd.writeList.Front()
+			if s == nil {
+				break
+			}
+			e.snd.writeList.Remove(s)
+			s.DecRef()
+		}
+		e.sndQueueInfo.SndBufUsed = 0
+		e.sndQueueInfo.SndClosed = true
+	}
+}
+
 // Abort implements stack.TransportEndpoint.Abort.
 func (e *endpoint) Abort() {
 	// The abort notification is not processed synchronously, so no
@@ -1041,6 +1107,9 @@ func (e *endpoint) Close() {
 		return
 	}
 
+	// We always want to purge the read queue, but do so after the checks in
+	// shutdownLocked.
+	defer e.purgeReadQueue()
 	linger := e.SocketOptions().GetLinger()
 	if linger.Enabled && linger.Timeout == 0 {
 		s := e.EndpointState()
@@ -1066,6 +1135,10 @@ func (e *endpoint) Close() {
 	// if we're connected, or stop accepting if we're listening.
 	e.shutdownLocked(tcpip.ShutdownWrite | tcpip.ShutdownRead)
 	e.closeNoShutdownLocked()
+	switch e.EndpointState() {
+	case StateClose, StateError:
+		e.releaseLocked()
+	}
 }
 
 // closeNoShutdown closes the endpoint without doing a full shutdown.
@@ -1187,6 +1260,8 @@ func (e *endpoint) cleanupLocked() {
 		e.route = nil
 	}
 
+	// It's not safe to purge the read queues yet, there could be unread data.
+	e.purgeWriteQueue()
 	e.stack.CompleteTransportEndpointCleanup(e)
 	tcpip.DeleteDanglingEndpoint(e)
 }
@@ -1470,10 +1545,16 @@ func (e *endpoint) commitRead(done int) *segment {
 		// Memory is only considered released when the whole segment has been
 		// read.
 		memDelta += s.segMemSize()
-		s.decRef()
+		s.DecRef()
 		s = e.rcvQueueInfo.rcvQueue.Front()
 	}
-	e.rcvQueueInfo.RcvBufUsed -= done
+	// Concurrent calls to Close() and Read() could cause RcvBufUsed to be
+	// negative because Read() unlocks between startRead() and commitRead(). In
+	// this case the read is allowed, but we refrain from subtracting from
+	// RcvBufUsed since it should already be zero.
+	if e.rcvQueueInfo.RcvBufUsed != 0 {
+		e.rcvQueueInfo.RcvBufUsed -= done
+	}
 
 	if memDelta > 0 {
 		// If the window was small before this read and if the read freed up
@@ -1599,6 +1680,7 @@ func (e *endpoint) queueSegment(p tcpip.Payloader, opts tcpip.WriteOptions) (*se
 	// Add data to the send queue.
 	s := newOutgoingSegment(e.TransportEndpointInfo.ID, e.stack.Clock(), v)
 	e.sndQueueInfo.SndBufUsed += len(v)
+	s.IncRef()
 	e.snd.writeList.PushBack(s)
 
 	return s, len(v), nil
@@ -1616,6 +1698,9 @@ func (e *endpoint) Write(p tcpip.Payloader, opts tcpip.WriteOptions) (int64, tcp
 	// Return if either we didn't queue anything or if an error occurred while
 	// attempting to queue data.
 	nextSeg, n, err := e.queueSegment(p, opts)
+	if nextSeg != nil {
+		defer nextSeg.DecRef()
+	}
 	if n == 0 || err != nil {
 		return 0, err
 	}
@@ -2892,7 +2977,7 @@ func (e *endpoint) readyToRead(s *segment) {
 	e.rcvQueueInfo.rcvQueueMu.Lock()
 	if s != nil {
 		e.rcvQueueInfo.RcvBufUsed += s.payloadSize()
-		s.incRef()
+		s.IncRef()
 		e.rcvQueueInfo.rcvQueue.PushBack(s)
 	} else {
 		e.rcvQueueInfo.RcvClosed = true
diff --git a/pkg/tcpip/transport/tcp/forwarder.go b/pkg/tcpip/transport/tcp/forwarder.go
index dc06c4335..cbe2d3859 100644
--- a/pkg/tcpip/transport/tcp/forwarder.go
+++ b/pkg/tcpip/transport/tcp/forwarder.go
@@ -66,7 +66,7 @@ func NewForwarder(s *stack.Stack, rcvWnd, maxInFlight int, handler func(*Forward
 // stack.SetTransportProtocolHandler function.
 func (f *Forwarder) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
 	s := newIncomingSegment(id, f.stack.Clock(), pkt)
-	defer s.decRef()
+	defer s.DecRef()
 
 	// We only care about well-formed SYN packets (not SYN-ACK) packets.
 	if !s.parse(pkt.RXTransportChecksumValidated) || !s.csumValid || !s.flags.Contains(header.TCPFlagSyn) || s.flags.Contains(header.TCPFlagAck) {
@@ -90,7 +90,7 @@ func (f *Forwarder) HandlePacket(id stack.TransportEndpointID, pkt *stack.Packet
 
 	// Launch a new goroutine to handle the request.
 	f.inFlight[id] = struct{}{}
-	s.incRef()
+	s.IncRef()
 	go f.handler(&ForwarderRequest{ // S/R-SAFE: not used by Sentry.
 		forwarder:  f,
 		segment:    s,
@@ -136,7 +136,7 @@ func (r *ForwarderRequest) Complete(sendReset bool) {
 	}
 
 	// Release all resources.
-	r.segment.decRef()
+	r.segment.DecRef()
 	r.segment = nil
 	r.forwarder = nil
 }
diff --git a/pkg/tcpip/transport/tcp/protocol.go b/pkg/tcpip/transport/tcp/protocol.go
index 8d3ab0dc9..af9edb12e 100644
--- a/pkg/tcpip/transport/tcp/protocol.go
+++ b/pkg/tcpip/transport/tcp/protocol.go
@@ -158,7 +158,7 @@ func (p *protocol) QueuePacket(ep stack.TransportEndpoint, id stack.TransportEnd
 // means."
 func (p *protocol) HandleUnknownDestinationPacket(id stack.TransportEndpointID, pkt *stack.PacketBuffer) stack.UnknownDestinationPacketDisposition {
 	s := newIncomingSegment(id, p.stack.Clock(), pkt)
-	defer s.decRef()
+	defer s.DecRef()
 
 	if !s.parse(pkt.RXTransportChecksumValidated) || !s.csumValid {
 		return stack.UnknownDestinationPacketMalformed
diff --git a/pkg/tcpip/transport/tcp/rack.go b/pkg/tcpip/transport/tcp/rack.go
index bd2d855c2..fe0a47e13 100644
--- a/pkg/tcpip/transport/tcp/rack.go
+++ b/pkg/tcpip/transport/tcp/rack.go
@@ -197,7 +197,7 @@ func (s *sender) probeTimerExpired() tcpip.Error {
 		dataSent = s.maybeSendSegment(s.writeNext, int(s.ep.scoreboard.SMSS()), s.SndUna.Add(s.SndWnd))
 		if dataSent {
 			s.Outstanding += s.pCount(s.writeNext, s.MaxPayloadSize)
-			s.writeNext = s.writeNext.Next()
+			s.updateWriteNext(s.writeNext.Next())
 		}
 	}
 
diff --git a/pkg/tcpip/transport/tcp/rcv.go b/pkg/tcpip/transport/tcp/rcv.go
index e9bbcc436..0cccb9d8a 100644
--- a/pkg/tcpip/transport/tcp/rcv.go
+++ b/pkg/tcpip/transport/tcp/rcv.go
@@ -286,8 +286,7 @@ func (r *receiver) consumeSegment(s *segment, segSeq seqnum.Value, segLen seqnum
 
 		for i := first; i < len(r.pendingRcvdSegments); i++ {
 			r.PendingBufUsed -= r.pendingRcvdSegments[i].segMemSize()
-			r.pendingRcvdSegments[i].decRef()
-
+			r.pendingRcvdSegments[i].DecRef()
 			// Note that slice truncation does not allow garbage collection of
 			// truncated items, thus truncated items must be set to nil to avoid
 			// memory leaks.
@@ -490,7 +489,7 @@ func (r *receiver) handleRcvdSegment(s *segment) (drop bool, err tcpip.Error) {
 				r.ep.rcvQueueInfo.rcvQueueMu.Lock()
 				r.PendingBufUsed += s.segMemSize()
 				r.ep.rcvQueueInfo.rcvQueueMu.Unlock()
-				s.incRef()
+				s.IncRef()
 				heap.Push(&r.pendingRcvdSegments, s)
 				UpdateSACKBlocks(&r.ep.sack, segSeq, segSeq.Add(segLen), r.RcvNxt)
 			}
@@ -526,7 +525,7 @@ func (r *receiver) handleRcvdSegment(s *segment) (drop bool, err tcpip.Error) {
 		r.ep.rcvQueueInfo.rcvQueueMu.Lock()
 		r.PendingBufUsed -= s.segMemSize()
 		r.ep.rcvQueueInfo.rcvQueueMu.Unlock()
-		s.decRef()
+		s.DecRef()
 	}
 	return false, nil
 }
diff --git a/pkg/tcpip/transport/tcp/sack_recovery.go b/pkg/tcpip/transport/tcp/sack_recovery.go
index a3a7e904b..74f1698d7 100644
--- a/pkg/tcpip/transport/tcp/sack_recovery.go
+++ b/pkg/tcpip/transport/tcp/sack_recovery.go
@@ -69,7 +69,7 @@ func (sr *sackRecovery) handleSACKRecovery(limit int, end seqnum.Value) (dataSen
 			}
 			dataSent = true
 			snd.Outstanding++
-			snd.writeNext = nextSeg.Next()
+			snd.updateWriteNext(nextSeg.Next())
 			continue
 		}
 
diff --git a/pkg/tcpip/transport/tcp/segment.go b/pkg/tcpip/transport/tcp/segment.go
index ca78c96f2..43968704b 100644
--- a/pkg/tcpip/transport/tcp/segment.go
+++ b/pkg/tcpip/transport/tcp/segment.go
@@ -16,7 +16,6 @@ package tcp
 
 import (
 	"fmt"
-	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
@@ -41,7 +40,8 @@ const (
 // +stateify savable
 type segment struct {
 	segmentEntry
-	refCnt int32
+	segmentRefs
+
 	ep     *endpoint
 	qFlags queueFlags
 	id     stack.TransportEndpointID `state:"manual"`
@@ -90,13 +90,13 @@ type segment struct {
 func newIncomingSegment(id stack.TransportEndpointID, clock tcpip.Clock, pkt *stack.PacketBuffer) *segment {
 	netHdr := pkt.Network()
 	s := &segment{
-		refCnt:   1,
 		id:       id,
 		srcAddr:  netHdr.SourceAddress(),
 		dstAddr:  netHdr.DestinationAddress(),
 		netProto: pkt.NetworkProtocolNumber,
 		nicID:    pkt.NICID,
 	}
+	s.InitRefs()
 	s.data = pkt.Data().ExtractVV().Clone(s.views[:])
 	s.hdr = header.TCP(pkt.TransportHeader().View())
 	s.rcvdTime = clock.NowMonotonic()
@@ -106,9 +106,9 @@ func newIncomingSegment(id stack.TransportEndpointID, clock tcpip.Clock, pkt *st
 
 func newOutgoingSegment(id stack.TransportEndpointID, clock tcpip.Clock, v buffer.View) *segment {
 	s := &segment{
-		refCnt: 1,
-		id:     id,
+		id: id,
 	}
+	s.InitRefs()
 	s.rcvdTime = clock.NowMonotonic()
 	if len(v) != 0 {
 		s.views[0] = v
@@ -120,7 +120,6 @@ func newOutgoingSegment(id stack.TransportEndpointID, clock tcpip.Clock, v buffe
 
 func (s *segment) clone() *segment {
 	t := &segment{
-		refCnt:         1,
 		id:             s.id,
 		sequenceNumber: s.sequenceNumber,
 		ackNumber:      s.ackNumber,
@@ -135,6 +134,7 @@ func (s *segment) clone() *segment {
 		qFlags:         s.qFlags,
 		dataMemSize:    s.dataMemSize,
 	}
+	t.InitRefs()
 	t.data = s.data.Clone(t.views[:])
 	return t
 }
@@ -164,8 +164,8 @@ func (s *segment) setOwner(ep *endpoint, qFlags queueFlags) {
 	s.qFlags = qFlags
 }
 
-func (s *segment) decRef() {
-	if atomic.AddInt32(&s.refCnt, -1) == 0 {
+func (s *segment) DecRef() {
+	s.segmentRefs.DecRef(func() {
 		if s.ep != nil {
 			switch s.qFlags {
 			case recvQ:
@@ -176,11 +176,7 @@ func (s *segment) decRef() {
 				panic(fmt.Sprintf("unexpected queue flag %b set for segment", s.qFlags))
 			}
 		}
-	}
-}
-
-func (s *segment) incRef() {
-	atomic.AddInt32(&s.refCnt, 1)
+	})
 }
 
 // logicalLen is the segment length in the sequence number space. It's defined
diff --git a/pkg/tcpip/transport/tcp/segment_queue.go b/pkg/tcpip/transport/tcp/segment_queue.go
index d0d1b0b8a..568c7d358 100644
--- a/pkg/tcpip/transport/tcp/segment_queue.go
+++ b/pkg/tcpip/transport/tcp/segment_queue.go
@@ -60,6 +60,7 @@ func (q *segmentQueue) enqueue(s *segment) bool {
 	allow := (used <= int(bufSz) || s.payloadSize() == 0) && !q.frozen
 
 	if allow {
+		s.IncRef()
 		q.list.PushBack(s)
 		// Set the owner now that the endpoint owns the segment.
 		s.setOwner(q.ep, recvQ)
diff --git a/pkg/tcpip/transport/tcp/segment_test.go b/pkg/tcpip/transport/tcp/segment_test.go
index 2d5fdda19..3af4e8a73 100644
--- a/pkg/tcpip/transport/tcp/segment_test.go
+++ b/pkg/tcpip/transport/tcp/segment_test.go
@@ -43,9 +43,9 @@ func TestSegmentMerge(t *testing.T) {
 	var clock faketime.NullClock
 	id := stack.TransportEndpointID{}
 	seg1 := newOutgoingSegment(id, &clock, buffer.NewView(10))
-	defer seg1.decRef()
+	defer seg1.DecRef()
 	seg2 := newOutgoingSegment(id, &clock, buffer.NewView(20))
-	defer seg2.decRef()
+	defer seg2.DecRef()
 
 	checkSegmentSize(t, "seg1", seg1, segmentSizeWants{
 		DataSize:   10,
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index 1c5481b21..4e22faf0a 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -333,7 +333,7 @@ func (s *sender) updateMaxPayloadSize(mtu, count int) {
 
 	// Since we likely reduced the number of outstanding packets, we may be
 	// ready to send some more.
-	s.writeNext = nextSeg
+	s.updateWriteNext(nextSeg)
 	s.sendData()
 }
 
@@ -545,7 +545,7 @@ func (s *sender) retransmitTimerExpired() bool {
 	// information as we lack more rigorous checks to validate if the SACK
 	// information is usable after an RTO.
 	s.ep.scoreboard.Reset()
-	s.writeNext = s.writeList.Front()
+	s.updateWriteNext(s.writeList.Front())
 
 	// RFC 1122 4.2.2.17: Start sending zero window probes when we still see a
 	// zero receive window after retransmission interval and we have data to
@@ -752,7 +752,7 @@ func (s *sender) maybeSendSegment(seg *segment, limit int, end seqnum.Value) (se
 				}
 				seg.merge(nSeg)
 				s.writeList.Remove(nSeg)
-				nSeg.decRef()
+				nSeg.DecRef()
 			}
 			if !nextTooBig && seg.data.Size() < available {
 				// Segment is not full.
@@ -972,7 +972,7 @@ func (s *sender) sendData() {
 		if s.isAssignedSequenceNumber(seg) && s.ep.SACKPermitted && s.ep.scoreboard.IsSACKED(seg.sackBlock()) {
 			// Move writeNext along so that we don't try and scan data that
 			// has already been SACKED.
-			s.writeNext = seg.Next()
+			s.updateWriteNext(seg.Next())
 			continue
 		}
 		if sent := s.maybeSendSegment(seg, limit, end); !sent {
@@ -980,7 +980,7 @@ func (s *sender) sendData() {
 		}
 		dataSent = true
 		s.Outstanding += s.pCount(seg, s.MaxPayloadSize)
-		s.writeNext = seg.Next()
+		s.updateWriteNext(seg.Next())
 	}
 
 	s.postXmit(dataSent, true /* shouldScheduleProbe */)
@@ -1526,7 +1526,7 @@ func (s *sender) handleRcvdSegment(rcvdSeg *segment) {
 			}
 
 			if s.writeNext == seg {
-				s.writeNext = seg.Next()
+				s.updateWriteNext(seg.Next())
 			}
 
 			// Update the RACK fields if SACK is enabled.
@@ -1545,7 +1545,7 @@ func (s *sender) handleRcvdSegment(rcvdSeg *segment) {
 			} else {
 				s.SackedOut -= s.pCount(seg, s.MaxPayloadSize)
 			}
-			seg.decRef()
+			seg.DecRef()
 			ackLeft -= datalen
 		}
 
@@ -1699,3 +1699,13 @@ func (s *sender) maybeSendOutOfWindowAck(seg *segment) {
 		s.sendAck()
 	}
 }
+
+func (s *sender) updateWriteNext(seg *segment) {
+	if s.writeNext != nil {
+		s.writeNext.DecRef()
+	}
+	if seg != nil {
+		seg.IncRef()
+	}
+	s.writeNext = seg
+}
diff --git a/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go b/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
index dd8a5e137..ceaee440f 100644
--- a/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
+++ b/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
@@ -41,6 +41,7 @@ func TestForwarderSendMSSLessThanMTU(t *testing.T) {
 		c.EP, err = r.CreateEndpoint(&c.WQ)
 		ch <- err
 		close(ch)
+		r.Complete(false)
 	})
 	s.SetTransportProtocolHandler(tcp.ProtocolNumber, f.HandlePacket)
 
@@ -85,6 +86,7 @@ func TestForwarderDoesNotRejectECNFlags(t *testing.T) {
 				c.EP, err = r.CreateEndpoint(&c.WQ)
 				ch <- err
 				close(ch)
+				r.Complete(false)
 			})
 			s.SetTransportProtocolHandler(tcp.ProtocolNumber, f.HandlePacket)
 
diff --git a/pkg/tcpip/transport/tcp/test/e2e/tcp_test.go b/pkg/tcpip/transport/tcp/test/e2e/tcp_test.go
index 40f41a5d9..401af4bb0 100644
--- a/pkg/tcpip/transport/tcp/test/e2e/tcp_test.go
+++ b/pkg/tcpip/transport/tcp/test/e2e/tcp_test.go
@@ -855,7 +855,7 @@ func TestClosingWithEnqueuedSegments(t *testing.T) {
 		t.Errorf("got stats.TCP.CurrentEstablished.Value() = %d, want = 0", got)
 	}
 
-	// Check if the endpoint was moved to CLOSED and netstack a reset in
+	// Check if the endpoint was moved to CLOSED and netstack sent a reset in
 	// response to the ACK packet that we sent after last-ACK.
 	checker.IPv4(t, c.GetPacket(),
 		checker.TCP(
@@ -8628,9 +8628,43 @@ func TestECNFlagsAccept(t *testing.T) {
 	}
 }
 
+func TestReadAfterCloseWithBufferedData(t *testing.T) {
+	c := context.New(t, e2e.DefaultMTU)
+	defer c.Cleanup()
+	con := c.CreateConnectedWithOptionsNoDelay(header.TCPSynOptions{})
+	// Fill up the receive queue.
+	for i := 0; i < 300; i++ {
+		con.SendPacket([]byte{1, 2, 3, 4}, nil)
+	}
+
+	timeout := time.After(5 * time.Second)
+	// If the receive queue is not properly drained, the endpoint will never
+	// return ErrClosedForReceive.
+	c.EP.Close()
+	for {
+		select {
+		case <-timeout:
+			t.Fatalf("timed out waiting for read to return error %q", &tcpip.ErrClosedForReceive{})
+			return
+		default:
+			if _, err := c.EP.Read(ioutil.Discard, tcpip.ReadOptions{}); cmp.Equal(err, &tcpip.ErrClosedForReceive{}) {
+				return
+			}
+		}
+	}
+}
+
+func TestReleaseAfterClose(t *testing.T) {
+	c := context.New(t, e2e.DefaultMTU)
+	c.CreateConnectedWithOptionsNoDelay(header.TCPSynOptions{})
+	c.CloseNoWait()
+	c.EP.Release()
+}
+
 func TestMain(m *testing.M) {
 	refs.SetLeakMode(refs.LeaksPanic)
 	code := m.Run()
+	tcpip.ReleaseDanglingEndpoints()
 	// Allow TCP async work to complete to avoid false reports of leaks.
 	// TODO(gvisor.dev/issue/5940): Use fake clock in tests.
 	time.Sleep(1 * time.Second)
diff --git a/pkg/tcpip/transport/tcp/testing/context/context.go b/pkg/tcpip/transport/tcp/testing/context/context.go
index e7cbaf7a3..0f0b49d6d 100644
--- a/pkg/tcpip/transport/tcp/testing/context/context.go
+++ b/pkg/tcpip/transport/tcp/testing/context/context.go
@@ -280,12 +280,20 @@ func NewWithOpts(t *testing.T, opts Options) *Context {
 func (c *Context) Cleanup() {
 	if c.EP != nil {
 		c.EP.Close()
+		c.EP.Release()
 	}
 	c.Stack().Close()
 	c.Stack().Wait()
 	c.linkEP.Close()
 }
 
+// CloseNoWait closes the link and transport endpoints and does not wait for
+// them to finish processing.
+func (c *Context) CloseNoWait() {
+	c.EP.Close()
+	c.linkEP.Close()
+}
+
 // Stack returns a reference to the stack in the Context.
 func (c *Context) Stack() *stack.Stack {
 	return c.s
diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go
index 17fd766cd..0e86c2541 100644
--- a/pkg/tcpip/transport/udp/endpoint.go
+++ b/pkg/tcpip/transport/udp/endpoint.go
@@ -157,6 +157,8 @@ func (e *endpoint) Abort() {
 	e.Close()
 }
 
+func (*endpoint) Release() {}
+
 // Close puts the endpoint in a closed state and frees all resources
 // associated with it.
 func (e *endpoint) Close() {
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index eddc01a6f..afcf0887f 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -515,6 +515,9 @@ func (l *Loader) Destroy() {
 	// save/restore.
 	l.k.Release()
 
+	// Release any dangling tcp connections.
+	tcpip.ReleaseDanglingEndpoints()
+
 	// In the success case, stdioFDs and goferFDs will only contain
 	// released/closed FDs that ownership has been passed over to host FDs and
 	// gofer sessions. Close them here in case of failure.