4 files changed, 23 insertions, 3 deletions
diff --git a/pkg/tcpip/link/fdbased/endpoint.go b/pkg/tcpip/link/fdbased/endpoint.go
index 8fc009160..4e20cfbf8 100644
--- a/pkg/tcpip/link/fdbased/endpoint.go
+++ b/pkg/tcpip/link/fdbased/endpoint.go
@@ -69,6 +69,8 @@ type Options struct {
 	ChecksumOffload bool
 	ClosedFunc      func(*tcpip.Error)
 	Address         tcpip.LinkAddress
+	SaveRestore     bool
+	DisconnectOk    bool
 }
 
 // New creates a new fd-based endpoint.
@@ -89,6 +91,14 @@ func New(opts *Options) tcpip.LinkEndpointID {
 		caps |= stack.CapabilityResolutionRequired
 	}
 
+	if opts.SaveRestore {
+		caps |= stack.CapabilitySaveRestore
+	}
+
+	if opts.DisconnectOk {
+		caps |= stack.CapabilityDisconnectOk
+	}
+
 	e := &endpoint{
 		fd:      opts.FD,
 		mtu:     opts.MTU,
diff --git a/pkg/tcpip/stack/registration.go b/pkg/tcpip/stack/registration.go
index c66f925a8..01a29689d 100644
--- a/pkg/tcpip/stack/registration.go
+++ b/pkg/tcpip/stack/registration.go
@@ -204,6 +204,7 @@ const (
 	CapabilityChecksumOffload LinkEndpointCapabilities = 1 << iota
 	CapabilityResolutionRequired
 	CapabilitySaveRestore
+	CapabilityDisconnectOk
 )
 
 // LinkEndpoint is the interface implemented by data link layer protocols (e.g.,
diff --git a/pkg/tcpip/transport/tcp/connect.go b/pkg/tcpip/transport/tcp/connect.go
index 8d70eb45a..b90d3fe48 100644
--- a/pkg/tcpip/transport/tcp/connect.go
+++ b/pkg/tcpip/transport/tcp/connect.go
@@ -1023,7 +1023,9 @@ func (e *endpoint) protocolMainLoop(handshake bool) *tcpip.Error {
 
 	// Mark endpoint as closed.
 	e.mu.Lock()
-	e.state = stateClosed
+	if e.state != stateError {
+		e.state = stateClosed
+	}
 	// Lock released below.
 	epilogue()
 
diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go
index 9c2e8878f..6143390b3 100644
--- a/pkg/tcpip/transport/tcp/endpoint_state.go
+++ b/pkg/tcpip/transport/tcp/endpoint_state.go
@@ -52,10 +52,17 @@ func (e *endpoint) beforeSave() {
 	case stateInitial, stateBound:
 	case stateConnected:
 		if e.route.Capabilities()&stack.CapabilitySaveRestore == 0 {
-			panic(tcpip.ErrSaveRejection{fmt.Errorf("endpoint cannot be saved in connected state: local %v:%d, remote %v:%d", e.id.LocalAddress, e.id.LocalPort, e.id.RemoteAddress, e.id.RemotePort)})
+			if e.route.Capabilities()&stack.CapabilityDisconnectOk == 0 {
+				panic(tcpip.ErrSaveRejection{fmt.Errorf("endpoint cannot be saved in connected state: local %v:%d, remote %v:%d", e.id.LocalAddress, e.id.LocalPort, e.id.RemoteAddress, e.id.RemotePort)})
+			}
+			e.resetConnectionLocked(tcpip.ErrConnectionAborted)
+			e.mu.Unlock()
+			e.Close()
+			e.mu.Lock()
 		}
 		if !e.workerRunning {
-			// The endpoint must be in acceptedChan.
+			// The endpoint must be in acceptedChan or has been just
+			// disconnected and closed.
 			break
 		}
 		fallthrough