2 files changed, 32 insertions, 39 deletions
diff --git a/pkg/tcpip/transport/packet/endpoint.go b/pkg/tcpip/transport/packet/endpoint.go
index 0554d2f4a..2c9786175 100644
--- a/pkg/tcpip/transport/packet/endpoint.go
+++ b/pkg/tcpip/transport/packet/endpoint.go
@@ -68,31 +68,31 @@ type endpoint struct {
 	netProto    tcpip.NetworkProtocolNumber
 	waiterQueue *waiter.Queue
 	cooked      bool
-
-	// The following fields are used to manage the receive queue and are
-	// protected by rcvMu.
-	rcvMu      sync.Mutex `state:"nosave"`
-	rcvList    packetList
+	ops         tcpip.SocketOptions
+	stats       tcpip.TransportEndpointStats `state:"nosave"`
+
+	// The following fields are used to manage the receive queue.
+	rcvMu sync.Mutex `state:"nosave"`
+	// +checklocks:rcvMu
+	rcvList packetList
+	// +checklocks:rcvMu
 	rcvBufSize int
-	rcvClosed  bool
-
-	// The following fields are protected by mu.
-	mu       sync.RWMutex `state:"nosave"`
-	closed   bool
-	stats    tcpip.TransportEndpointStats `state:"nosave"`
-	bound    bool
+	// +checklocks:rcvMu
+	rcvClosed bool
+	// +checklocks:rcvMu
+	rcvDisabled bool
+
+	mu sync.RWMutex `state:"nosave"`
+	// +checklocks:mu
+	closed bool
+	// +checklocks:mu
+	bound bool
+	// +checklocks:mu
 	boundNIC tcpip.NICID
 
-	// lastErrorMu protects lastError.
 	lastErrorMu sync.Mutex `state:"nosave"`
-	lastError   tcpip.Error
-
-	// ops is used to get socket level options.
-	ops tcpip.SocketOptions
-
-	// frozen indicates if the packets should be delivered to the endpoint
-	// during restore.
-	frozen bool
+	// +checklocks:lastErrorMu
+	lastError tcpip.Error
 }
 
 // NewEndpoint returns a new packet endpoint.
@@ -414,7 +414,7 @@ func (ep *endpoint) HandlePacket(nicID tcpip.NICID, localAddr tcpip.LinkAddress,
 	}
 
 	rcvBufSize := ep.ops.GetReceiveBufferSize()
-	if ep.frozen || ep.rcvBufSize >= int(rcvBufSize) {
+	if ep.rcvDisabled || ep.rcvBufSize >= int(rcvBufSize) {
 		ep.rcvMu.Unlock()
 		ep.stack.Stats().DroppedPackets.Increment()
 		ep.stats.ReceiveErrors.ReceiveBufferOverflow.Increment()
@@ -491,18 +491,3 @@ func (*endpoint) SetOwner(tcpip.PacketOwner) {}
 func (ep *endpoint) SocketOptions() *tcpip.SocketOptions {
 	return &ep.ops
 }
-
-// freeze prevents any more packets from being delivered to the endpoint.
-func (ep *endpoint) freeze() {
-	ep.mu.Lock()
-	ep.frozen = true
-	ep.mu.Unlock()
-}
-
-// thaw unfreezes a previously frozen endpoint using endpoint.freeze() allows
-// new packets to be delivered again.
-func (ep *endpoint) thaw() {
-	ep.mu.Lock()
-	ep.frozen = false
-	ep.mu.Unlock()
-}
diff --git a/pkg/tcpip/transport/packet/endpoint_state.go b/pkg/tcpip/transport/packet/endpoint_state.go
index 5c688d286..d2768db7b 100644
--- a/pkg/tcpip/transport/packet/endpoint_state.go
+++ b/pkg/tcpip/transport/packet/endpoint_state.go
@@ -44,12 +44,16 @@ func (p *packet) loadData(data buffer.VectorisedView) {
 
 // beforeSave is invoked by stateify.
 func (ep *endpoint) beforeSave() {
-	ep.freeze()
+	ep.rcvMu.Lock()
+	defer ep.rcvMu.Unlock()
+	ep.rcvDisabled = true
 }
 
 // afterLoad is invoked by stateify.
 func (ep *endpoint) afterLoad() {
-	ep.thaw()
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+
 	ep.stack = stack.StackFromEnv
 	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
 
@@ -57,4 +61,8 @@ func (ep *endpoint) afterLoad() {
 	if err := ep.stack.RegisterPacketEndpoint(0, ep.netProto, ep); err != nil {
 		panic(err)
 	}
+
+	ep.rcvMu.Lock()
+	ep.rcvDisabled = false
+	ep.rcvMu.Unlock()
 }