5 files changed, 28 insertions, 22 deletions
diff --git a/pkg/sentry/kernel/seqatomic_taskgoroutineschedinfo_unsafe.go b/pkg/sentry/kernel/seqatomic_taskgoroutineschedinfo_unsafe.go
index 25ad17a4e..895abb129 100755
--- a/pkg/sentry/kernel/seqatomic_taskgoroutineschedinfo_unsafe.go
+++ b/pkg/sentry/kernel/seqatomic_taskgoroutineschedinfo_unsafe.go
@@ -1,12 +1,12 @@
 package kernel
 
 import (
-	"fmt"
-	"reflect"
-	"strings"
 	"unsafe"
 
+	"fmt"
 	"gvisor.dev/gvisor/third_party/gvsync"
+	"reflect"
+	"strings"
 )
 
 // SeqAtomicLoad returns a copy of *ptr, ensuring that the read does not race
diff --git a/pkg/sentry/platform/ring0/defs_impl.go b/pkg/sentry/platform/ring0/defs_impl.go
index 5032ac56e..a36a17e37 100755
--- a/pkg/sentry/platform/ring0/defs_impl.go
+++ b/pkg/sentry/platform/ring0/defs_impl.go
@@ -1,13 +1,14 @@
 package ring0
 
 import (
-	"fmt"
 	"gvisor.dev/gvisor/pkg/cpuid"
+	"reflect"
+	"syscall"
+
+	"fmt"
 	"gvisor.dev/gvisor/pkg/sentry/platform/ring0/pagetables"
 	"gvisor.dev/gvisor/pkg/sentry/usermem"
 	"io"
-	"reflect"
-	"syscall"
 )
 
 var (
diff --git a/pkg/sentry/time/seqatomic_parameters_unsafe.go b/pkg/sentry/time/seqatomic_parameters_unsafe.go
index 89792c56d..f6560d0bb 100755
--- a/pkg/sentry/time/seqatomic_parameters_unsafe.go
+++ b/pkg/sentry/time/seqatomic_parameters_unsafe.go
@@ -1,12 +1,12 @@
 package time
 
 import (
-	"fmt"
-	"reflect"
-	"strings"
 	"unsafe"
 
+	"fmt"
 	"gvisor.dev/gvisor/third_party/gvsync"
+	"reflect"
+	"strings"
 )
 
 // SeqAtomicLoad returns a copy of *ptr, ensuring that the read does not race
diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index 89154391b..cc49c8272 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -489,8 +489,8 @@ func (e *endpoint) Readiness(mask waiter.EventMask) waiter.EventMask {
 				result |= waiter.EventIn
 			}
 		}
-
-	case StateEstablished, StateFinWait1, StateFinWait2, StateTimeWait, StateCloseWait, StateLastAck, StateClosing:
+	}
+	if e.state.connected() {
 		// Determine if the endpoint is writable if requested.
 		if (mask & waiter.EventOut) != 0 {
 			e.sndBufMu.Lock()
@@ -1323,6 +1323,17 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) (er
 		return err
 	}
 
+	if e.state.connected() {
+		// The endpoint is already connected. If caller hasn't been
+		// notified yet, return success.
+		if !e.isConnectNotified {
+			e.isConnectNotified = true
+			return nil
+		}
+		// Otherwise return that it's already connected.
+		return tcpip.ErrAlreadyConnected
+	}
+
 	nicid := addr.NIC
 	switch e.state {
 	case StateBound:
@@ -1347,15 +1358,6 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) (er
 		// yet.
 		return tcpip.ErrAlreadyConnecting
 
-	case StateEstablished:
-		// The endpoint is already connected. If caller hasn't been notified yet, return success.
-		if !e.isConnectNotified {
-			e.isConnectNotified = true
-			return nil
-		}
-		// Otherwise return that it's already connected.
-		return tcpip.ErrAlreadyConnected
-
 	case StateError:
 		return e.hardError
 
diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go
index b93959034..b3f0f6c5d 100644
--- a/pkg/tcpip/transport/tcp/endpoint_state.go
+++ b/pkg/tcpip/transport/tcp/endpoint_state.go
@@ -50,6 +50,8 @@ func (e *endpoint) beforeSave() {
 
 	switch e.state {
 	case StateInitial, StateBound:
+		// TODO(b/138137272): this enumeration duplicates
+		// EndpointState.connected. remove it.
 	case StateEstablished, StateSynSent, StateSynRecv, StateFinWait1, StateFinWait2, StateTimeWait, StateCloseWait, StateLastAck, StateClosing:
 		if e.route.Capabilities()&stack.CapabilitySaveRestore == 0 {
 			if e.route.Capabilities()&stack.CapabilityDisconnectOk == 0 {
@@ -149,9 +151,10 @@ var connectingLoading sync.WaitGroup
 func (e *endpoint) loadState(state EndpointState) {
 	// This is to ensure that the loading wait groups include all applicable
 	// endpoints before any asynchronous calls to the Wait() methods.
-	switch state {
-	case StateEstablished, StateFinWait1, StateFinWait2, StateTimeWait, StateCloseWait, StateLastAck, StateClosing:
+	if state.connected() {
 		connectedLoading.Add(1)
+	}
+	switch state {
 	case StateListen:
 		listenLoading.Add(1)
 	case StateConnecting, StateSynSent, StateSynRecv: