3 files changed, 52 insertions, 59 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()
 }
diff --git a/pkg/tcpip/transport/packet/packet_state_autogen.go b/pkg/tcpip/transport/packet/packet_state_autogen.go
index 2640a55d0..6017e9e8b 100644
--- a/pkg/tcpip/transport/packet/packet_state_autogen.go
+++ b/pkg/tcpip/transport/packet/packet_state_autogen.go
@@ -59,15 +59,15 @@ func (ep *endpoint) StateFields() []string {
 		"netProto",
 		"waiterQueue",
 		"cooked",
+		"ops",
 		"rcvList",
 		"rcvBufSize",
 		"rcvClosed",
+		"rcvDisabled",
 		"closed",
 		"bound",
 		"boundNIC",
 		"lastError",
-		"ops",
-		"frozen",
 	}
 }
 
@@ -79,15 +79,15 @@ func (ep *endpoint) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(2, &ep.netProto)
 	stateSinkObject.Save(3, &ep.waiterQueue)
 	stateSinkObject.Save(4, &ep.cooked)
-	stateSinkObject.Save(5, &ep.rcvList)
-	stateSinkObject.Save(6, &ep.rcvBufSize)
-	stateSinkObject.Save(7, &ep.rcvClosed)
-	stateSinkObject.Save(8, &ep.closed)
-	stateSinkObject.Save(9, &ep.bound)
-	stateSinkObject.Save(10, &ep.boundNIC)
-	stateSinkObject.Save(11, &ep.lastError)
-	stateSinkObject.Save(12, &ep.ops)
-	stateSinkObject.Save(13, &ep.frozen)
+	stateSinkObject.Save(5, &ep.ops)
+	stateSinkObject.Save(6, &ep.rcvList)
+	stateSinkObject.Save(7, &ep.rcvBufSize)
+	stateSinkObject.Save(8, &ep.rcvClosed)
+	stateSinkObject.Save(9, &ep.rcvDisabled)
+	stateSinkObject.Save(10, &ep.closed)
+	stateSinkObject.Save(11, &ep.bound)
+	stateSinkObject.Save(12, &ep.boundNIC)
+	stateSinkObject.Save(13, &ep.lastError)
 }
 
 // +checklocksignore
@@ -97,15 +97,15 @@ func (ep *endpoint) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(2, &ep.netProto)
 	stateSourceObject.Load(3, &ep.waiterQueue)
 	stateSourceObject.Load(4, &ep.cooked)
-	stateSourceObject.Load(5, &ep.rcvList)
-	stateSourceObject.Load(6, &ep.rcvBufSize)
-	stateSourceObject.Load(7, &ep.rcvClosed)
-	stateSourceObject.Load(8, &ep.closed)
-	stateSourceObject.Load(9, &ep.bound)
-	stateSourceObject.Load(10, &ep.boundNIC)
-	stateSourceObject.Load(11, &ep.lastError)
-	stateSourceObject.Load(12, &ep.ops)
-	stateSourceObject.Load(13, &ep.frozen)
+	stateSourceObject.Load(5, &ep.ops)
+	stateSourceObject.Load(6, &ep.rcvList)
+	stateSourceObject.Load(7, &ep.rcvBufSize)
+	stateSourceObject.Load(8, &ep.rcvClosed)
+	stateSourceObject.Load(9, &ep.rcvDisabled)
+	stateSourceObject.Load(10, &ep.closed)
+	stateSourceObject.Load(11, &ep.bound)
+	stateSourceObject.Load(12, &ep.boundNIC)
+	stateSourceObject.Load(13, &ep.lastError)
 	stateSourceObject.AfterLoad(ep.afterLoad)
 }