From a3e58874158ec5978e748aee4aa00998d4bc8b9f Mon Sep 17 00:00:00 2001 From: Nayana Bidari Date: Wed, 18 Dec 2024 12:48:08 -0800 Subject: [PATCH] Changes to support netstack save restore. - Added a new Stats() method in inet.Stack to get the saved stats during restore. - Mark stack.nic, tcpip.Route and stack.addressState structs as "nosave". These fields should not be saved because the IP addresses and routes can change during restore and new configuration of routes and IP addresses will be extracted from the restore spec and initialized in the saved stack. - Changes in Restore() method in icmp, udp, tcp, packet and raw endpoint files to support save restore of these endpoints. These changes are flag guarded by the TESTONLY-save-restore-netstack flag. PiperOrigin-RevId: 707639274 --- pkg/sentry/inet/inet.go | 3 + pkg/sentry/inet/test_stack.go | 6 + pkg/sentry/kernel/timekeeper_state.go | 5 + pkg/sentry/socket/hostinet/stack.go | 5 + pkg/sentry/socket/netstack/stack.go | 16 ++- pkg/tcpip/stack/addressable_endpoint_state.go | 4 +- pkg/tcpip/stack/stack.go | 1 + pkg/tcpip/transport/icmp/endpoint.go | 2 +- pkg/tcpip/transport/icmp/endpoint_state.go | 10 +- .../transport/internal/network/endpoint.go | 4 +- pkg/tcpip/transport/packet/endpoint.go | 2 +- pkg/tcpip/transport/packet/endpoint_state.go | 12 +- pkg/tcpip/transport/raw/endpoint.go | 2 +- pkg/tcpip/transport/raw/endpoint_state.go | 17 ++- pkg/tcpip/transport/tcp/endpoint.go | 2 +- pkg/tcpip/transport/tcp/endpoint_state.go | 108 ++++++++++-------- pkg/tcpip/transport/udp/endpoint.go | 2 +- pkg/tcpip/transport/udp/endpoint_state.go | 16 ++- 18 files changed, 143 insertions(+), 74 deletions(-) diff --git a/pkg/sentry/inet/inet.go b/pkg/sentry/inet/inet.go index 5f12113c4..8b4ef8848 100644 --- a/pkg/sentry/inet/inet.go +++ b/pkg/sentry/inet/inet.go @@ -137,6 +137,9 @@ type Stack interface { // IsSaveRestoreEnabled returns true when netstack s/r is enabled. IsSaveRestoreEnabled() bool + + // Stats returns the network stats. + Stats() tcpip.Stats } // Interface contains information about a network interface. diff --git a/pkg/sentry/inet/test_stack.go b/pkg/sentry/inet/test_stack.go index 119e6c78b..26d3f6165 100644 --- a/pkg/sentry/inet/test_stack.go +++ b/pkg/sentry/inet/test_stack.go @@ -235,3 +235,9 @@ func (*TestStack) IsSaveRestoreEnabled() bool { // No-op. return false } + +// Stats implements Stack. +func (*TestStack) Stats() tcpip.Stats { + // No-op. + return tcpip.Stats{} +} diff --git a/pkg/sentry/kernel/timekeeper_state.go b/pkg/sentry/kernel/timekeeper_state.go index 38ad158f5..9845312dc 100644 --- a/pkg/sentry/kernel/timekeeper_state.go +++ b/pkg/sentry/kernel/timekeeper_state.go @@ -26,6 +26,11 @@ func (t *Timekeeper) beforeSave() { panic("pauseUpdates must be called before Save") } + if t.clocks == nil { + t.restored = nil + return + } + // N.B. we want the *offset* monotonic time. var err error if t.saveMonotonic, err = t.GetTime(time.Monotonic); err != nil { diff --git a/pkg/sentry/socket/hostinet/stack.go b/pkg/sentry/socket/hostinet/stack.go index 32fa99a04..4d1facb0a 100644 --- a/pkg/sentry/socket/hostinet/stack.go +++ b/pkg/sentry/socket/hostinet/stack.go @@ -438,3 +438,8 @@ func (*Stack) EnableSaveRestore() error { func (s *Stack) IsSaveRestoreEnabled() bool { return false } + +// Stats implements inet.Stack.Stats. +func (s *Stack) Stats() tcpip.Stats { + return tcpip.Stats{} +} diff --git a/pkg/sentry/socket/netstack/stack.go b/pkg/sentry/socket/netstack/stack.go index c579f5a2f..fc21771bd 100644 --- a/pkg/sentry/socket/netstack/stack.go +++ b/pkg/sentry/socket/netstack/stack.go @@ -592,6 +592,7 @@ func (s *Stack) SetTCPRecovery(recovery inet.TCPLossRecovery) error { // Statistics implements inet.Stack.Statistics. func (s *Stack) Statistics(stat any, arg string) error { + netStats := s.Stats() switch stats := stat.(type) { case *inet.StatDev: for _, ni := range s.Stack.NICInfo() { @@ -622,7 +623,7 @@ func (s *Stack) Statistics(stat any, arg string) error { break } case *inet.StatSNMPIP: - ip := Metrics.IP + ip := netStats.IP // TODO(gvisor.dev/issue/969) Support stubbed stats. *stats = inet.StatSNMPIP{ 0, // Ip/Forwarding. @@ -646,8 +647,8 @@ func (s *Stack) Statistics(stat any, arg string) error { 0, // Support Ip/FragCreates. } case *inet.StatSNMPICMP: - in := Metrics.ICMP.V4.PacketsReceived.ICMPv4PacketStats - out := Metrics.ICMP.V4.PacketsSent.ICMPv4PacketStats + in := netStats.ICMP.V4.PacketsReceived.ICMPv4PacketStats + out := netStats.ICMP.V4.PacketsSent.ICMPv4PacketStats // TODO(gvisor.dev/issue/969) Support stubbed stats. *stats = inet.StatSNMPICMP{ 0, // Icmp/InMsgs. @@ -679,7 +680,7 @@ func (s *Stack) Statistics(stat any, arg string) error { out.InfoReply.Value(), // OutAddrMaskReps. } case *inet.StatSNMPTCP: - tcp := Metrics.TCP + tcp := netStats.TCP // RFC 2012 (updates 1213): SNMPv2-MIB-TCP. *stats = inet.StatSNMPTCP{ 1, // RtoAlgorithm. @@ -699,7 +700,7 @@ func (s *Stack) Statistics(stat any, arg string) error { tcp.ChecksumErrors.Value(), // InCsumErrors. } case *inet.StatSNMPUDP: - udp := Metrics.UDP + udp := netStats.UDP // TODO(gvisor.dev/issue/969) Support stubbed stats. *stats = inet.StatSNMPUDP{ udp.PacketsReceived.Value(), // InDatagrams. @@ -717,6 +718,11 @@ func (s *Stack) Statistics(stat any, arg string) error { return nil } +// Stats implements inet.Stack.Stats. +func (s *Stack) Stats() tcpip.Stats { + return s.Stack.Stats() +} + // RouteTable implements inet.Stack.RouteTable. func (s *Stack) RouteTable() []inet.Route { var routeTable []inet.Route diff --git a/pkg/tcpip/stack/addressable_endpoint_state.go b/pkg/tcpip/stack/addressable_endpoint_state.go index bb2e0faf0..e08909c81 100644 --- a/pkg/tcpip/stack/addressable_endpoint_state.go +++ b/pkg/tcpip/stack/addressable_endpoint_state.go @@ -738,8 +738,6 @@ func (a *AddressableEndpointState) Cleanup() { var _ AddressEndpoint = (*addressState)(nil) // addressState holds state for an address. -// -// +stateify savable type addressState struct { addressableEndpointState *AddressableEndpointState addr tcpip.AddressWithPrefix @@ -750,7 +748,7 @@ type addressState struct { // // AddressableEndpointState.mu // addressState.mu - mu addressStateRWMutex `state:"nosave"` + mu addressStateRWMutex refs addressStateRefs // checklocks:mu kind AddressKind diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go index 0232a45ee..645f52f78 100644 --- a/pkg/tcpip/stack/stack.go +++ b/pkg/tcpip/stack/stack.go @@ -1986,6 +1986,7 @@ func (s *Stack) ReplaceConfig(st *Stack) { s.nics[id] = nic _ = s.NextNICID() } + s.tables = st.tables } // Restore restarts the stack after a restore. This must be called after the diff --git a/pkg/tcpip/transport/icmp/endpoint.go b/pkg/tcpip/transport/icmp/endpoint.go index 988604fcd..c7924d534 100644 --- a/pkg/tcpip/transport/icmp/endpoint.go +++ b/pkg/tcpip/transport/icmp/endpoint.go @@ -57,7 +57,7 @@ type endpoint struct { // The following fields are initialized at creation time and are // immutable. - stack *stack.Stack `state:"manual"` + stack *stack.Stack transProto tcpip.TransportProtocolNumber waiterQueue *waiter.Queue net network.Endpoint diff --git a/pkg/tcpip/transport/icmp/endpoint_state.go b/pkg/tcpip/transport/icmp/endpoint_state.go index 134797e8b..9ee7b5215 100644 --- a/pkg/tcpip/transport/icmp/endpoint_state.go +++ b/pkg/tcpip/transport/icmp/endpoint_state.go @@ -36,7 +36,11 @@ func (p *icmpPacket) loadReceivedAt(_ context.Context, nsec int64) { // afterLoad is invoked by stateify. func (e *endpoint) afterLoad(ctx context.Context) { - stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + if e.stack.IsSaveRestoreEnabled() { + e.stack.RegisterRestoredEndpoint(e) + } else { + stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + } } // beforeSave is invoked by stateify. @@ -50,6 +54,10 @@ func (e *endpoint) Restore(s *stack.Stack) { e.thaw() e.net.Resume(s) + if e.stack.IsSaveRestoreEnabled() { + e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) + return + } e.stack = s e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) diff --git a/pkg/tcpip/transport/internal/network/endpoint.go b/pkg/tcpip/transport/internal/network/endpoint.go index 6a0523963..8eefeb7b1 100644 --- a/pkg/tcpip/transport/internal/network/endpoint.go +++ b/pkg/tcpip/transport/internal/network/endpoint.go @@ -35,7 +35,7 @@ import ( // +stateify savable type Endpoint struct { // The following fields must only be set once then never changed. - stack *stack.Stack `state:"manual"` + stack *stack.Stack ops *tcpip.SocketOptions netProto tcpip.NetworkProtocolNumber transProto tcpip.TransportProtocolNumber @@ -53,7 +53,7 @@ type Endpoint struct { // +checklocks:mu effectiveNetProto tcpip.NetworkProtocolNumber // +checklocks:mu - connectedRoute *stack.Route `state:"manual"` + connectedRoute *stack.Route `state:"nosave"` // +checklocks:mu multicastMemberships map[multicastMembership]struct{} // +checklocks:mu diff --git a/pkg/tcpip/transport/packet/endpoint.go b/pkg/tcpip/transport/packet/endpoint.go index 9166bca6c..0cf7c0158 100644 --- a/pkg/tcpip/transport/packet/endpoint.go +++ b/pkg/tcpip/transport/packet/endpoint.go @@ -63,7 +63,7 @@ type endpoint struct { // The following fields are initialized at creation time and are // immutable. - stack *stack.Stack `state:"manual"` + stack *stack.Stack waiterQueue *waiter.Queue cooked bool ops tcpip.SocketOptions diff --git a/pkg/tcpip/transport/packet/endpoint_state.go b/pkg/tcpip/transport/packet/endpoint_state.go index 16be7d6b3..355381ff3 100644 --- a/pkg/tcpip/transport/packet/endpoint_state.go +++ b/pkg/tcpip/transport/packet/endpoint_state.go @@ -43,12 +43,20 @@ func (ep *endpoint) beforeSave() { // afterLoad is invoked by stateify. func (ep *endpoint) afterLoad(ctx context.Context) { + if !ep.stack.IsSaveRestoreEnabled() { + ep.mu.Lock() + ep.stack = stack.RestoreStackFromContext(ctx) + ep.mu.Unlock() + } + ep.stack.RegisterRestoredEndpoint(ep) +} + +// Restore implements tcpip.RestoredEndpoint.Restore. +func (ep *endpoint) Restore(_ *stack.Stack) { ep.mu.Lock() defer ep.mu.Unlock() - ep.stack = stack.RestoreStackFromContext(ctx) ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) - if err := ep.stack.RegisterPacketEndpoint(ep.boundNIC, ep.boundNetProto, ep); err != nil { panic(fmt.Sprintf("RegisterPacketEndpoint(%d, %d, _): %s", ep.boundNIC, ep.boundNetProto, err)) } diff --git a/pkg/tcpip/transport/raw/endpoint.go b/pkg/tcpip/transport/raw/endpoint.go index 1eaedc197..5203c4dd9 100644 --- a/pkg/tcpip/transport/raw/endpoint.go +++ b/pkg/tcpip/transport/raw/endpoint.go @@ -73,7 +73,7 @@ type endpoint struct { // The following fields are initialized at creation time and are // immutable. - stack *stack.Stack `state:"manual"` + stack *stack.Stack transProto tcpip.TransportProtocolNumber waiterQueue *waiter.Queue associated bool diff --git a/pkg/tcpip/transport/raw/endpoint_state.go b/pkg/tcpip/transport/raw/endpoint_state.go index d915ade2e..f2e8896ff 100644 --- a/pkg/tcpip/transport/raw/endpoint_state.go +++ b/pkg/tcpip/transport/raw/endpoint_state.go @@ -16,7 +16,6 @@ package raw import ( "context" - "fmt" "time" "gvisor.dev/gvisor/pkg/tcpip" @@ -35,7 +34,11 @@ func (p *rawPacket) loadReceivedAt(_ context.Context, nsec int64) { // afterLoad is invoked by stateify. func (e *endpoint) afterLoad(ctx context.Context) { - stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + if e.stack.IsSaveRestoreEnabled() { + e.stack.RegisterRestoredEndpoint(e) + } else { + stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + } } // beforeSave is invoked by stateify. @@ -46,16 +49,20 @@ func (e *endpoint) beforeSave() { // Restore implements tcpip.RestoredEndpoint.Restore. func (e *endpoint) Restore(s *stack.Stack) { - e.net.Resume(s) - e.setReceiveDisabled(false) + e.net.Resume(s) + if e.stack.IsSaveRestoreEnabled() { + e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) + return + } + e.stack = s e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) if e.associated { netProto := e.net.NetProto() if err := e.stack.RegisterRawTransportEndpoint(netProto, e.transProto, e); err != nil { - panic(fmt.Sprintf("e.stack.RegisterRawTransportEndpoint(%d, %d, _): %s", netProto, e.transProto, err)) + panic("RegisterRawTransportEndpoint failed during restore") } } } diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go index 6ef254ff4..49c2474c0 100644 --- a/pkg/tcpip/transport/tcp/endpoint.go +++ b/pkg/tcpip/transport/tcp/endpoint.go @@ -418,7 +418,7 @@ type Endpoint struct { isPortReserved bool isRegistered bool boundNICID tcpip.NICID - route *stack.Route `state:"manual"` + route *stack.Route `state:"nosave"` ipv4TTL uint8 ipv6HopLimit int16 isConnectNotified bool diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go index 98f9ca652..0c8e7e260 100644 --- a/pkg/tcpip/transport/tcp/endpoint_state.go +++ b/pkg/tcpip/transport/tcp/endpoint_state.go @@ -107,8 +107,9 @@ var connectingLoading sync.WaitGroup func (e *Endpoint) loadState(_ context.Context, epState EndpointState) { // This is to ensure that the loading wait groups include all applicable // endpoints before any asynchronous calls to the Wait() methods. - // For restore purposes we treat TimeWait like a connected endpoint. - if epState.connected() || epState == StateTimeWait { + // For restore purposes we treat all endpoints with state after + // StateEstablished and before StateClosed like connected endpoint. + if epState.connected() { connectedLoading.Add(1) } switch { @@ -159,21 +160,23 @@ func (e *Endpoint) Restore(s *stack.Stack) { bind := func() { e.mu.Lock() defer e.mu.Unlock() - addr, _, err := e.checkV4MappedLocked(tcpip.FullAddress{Addr: e.BindAddr, Port: e.TransportEndpointInfo.ID.LocalPort}, true /* bind */) - if err != nil { - panic("unable to parse BindAddr: " + err.String()) - } - portRes := ports.Reservation{ - Networks: e.effectiveNetProtos, - Transport: ProtocolNumber, - Addr: addr.Addr, - Port: addr.Port, - Flags: e.boundPortFlags, - BindToDevice: e.boundBindToDevice, - Dest: e.boundDest, - } - if ok := e.stack.ReserveTuple(portRes); !ok { - panic(fmt.Sprintf("unable to re-reserve tuple (%v, %q, %d, %+v, %d, %v)", e.effectiveNetProtos, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest)) + if !saveRestoreEnabled { + addr, _, err := e.checkV4MappedLocked(tcpip.FullAddress{Addr: e.BindAddr, Port: e.TransportEndpointInfo.ID.LocalPort}, true /* bind */) + if err != nil { + panic("unable to parse BindAddr: " + err.String()) + } + portRes := ports.Reservation{ + Networks: e.effectiveNetProtos, + Transport: ProtocolNumber, + Addr: addr.Addr, + Port: addr.Port, + Flags: e.boundPortFlags, + BindToDevice: e.boundBindToDevice, + Dest: e.boundDest, + } + if ok := e.stack.ReserveTuple(portRes); !ok { + panic(fmt.Sprintf("unable to re-reserve tuple (%v, %q, %d, %+v, %d, %v)", e.effectiveNetProtos, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest)) + } } e.isPortReserved = true @@ -201,6 +204,10 @@ func (e *Endpoint) Restore(s *stack.Stack) { // Reset the scoreboard to reinitialize the sack information as // we do not restore SACK information. e.scoreboard.Reset() + if saveRestoreEnabled { + // Unregister the endpoint before registering again during Connect. + e.stack.UnregisterTransportEndpoint(e.effectiveNetProtos, header.TCPProtocolNumber, e.TransportEndpointInfo.ID, e, e.boundPortFlags, e.boundBindToDevice) + } e.mu.Lock() err := e.connect(tcpip.FullAddress{NIC: e.boundNICID, Addr: e.connectingAddress, Port: e.TransportEndpointInfo.ID.RemotePort}, false /* handshake */) if _, ok := err.(*tcpip.ErrConnectStarted); !ok { @@ -224,8 +231,8 @@ func (e *Endpoint) Restore(s *stack.Stack) { e.mu.Unlock() connectedLoading.Done() case epState == StateListen: + tcpip.AsyncLoading.Add(1) if !saveRestoreEnabled { - tcpip.AsyncLoading.Add(1) go func() { connectedLoading.Wait() bind() @@ -244,14 +251,19 @@ func (e *Endpoint) Restore(s *stack.Stack) { tcpip.AsyncLoading.Done() }() } else { - e.LockUser() - // All endpoints will be moved to initial state after - // restore. Set endpoint to its originial listen state. - e.setEndpointState(StateListen) - // Initialize the listening context. - rcvWnd := seqnum.Size(e.receiveBufferAvailable()) - e.listenCtx = newListenContext(e.stack, e.protocol, e, rcvWnd, e.ops.GetV6Only(), e.NetProto) - e.UnlockUser() + go func() { + connectedLoading.Wait() + e.LockUser() + // All endpoints will be moved to initial state after + // restore. Set endpoint to its originial listen state. + e.setEndpointState(StateListen) + // Initialize the listening context. + rcvWnd := seqnum.Size(e.receiveBufferAvailable()) + e.listenCtx = newListenContext(e.stack, e.protocol, e, rcvWnd, e.ops.GetV6Only(), e.NetProto) + e.UnlockUser() + listenLoading.Done() + tcpip.AsyncLoading.Done() + }() } case epState == StateConnecting: // Initial SYN hasn't been sent yet so initiate a connect. @@ -268,26 +280,30 @@ func (e *Endpoint) Restore(s *stack.Stack) { tcpip.AsyncLoading.Done() }() case epState == StateSynSent || epState == StateSynRecv: - connectedLoading.Wait() - listenLoading.Wait() - // Initial SYN has been sent/received so we should bind the - // ports start the retransmit timer for the SYNs and let it - // naturally complete the connection. - bind() - e.mu.Lock() - defer e.mu.Unlock() - e.setEndpointState(epState) - r, err := e.stack.FindRoute(e.boundNICID, e.TransportEndpointInfo.ID.LocalAddress, e.TransportEndpointInfo.ID.RemoteAddress, e.effectiveNetProtos[0], false /* multicastLoop */) - if err != nil { - panic(fmt.Sprintf("FindRoute failed when restoring endpoint w/ ID: %+v", e.ID)) - } - e.route = r - timer, err := newBackoffTimer(e.stack.Clock(), InitialRTO, MaxRTO, timerHandler(e, e.h.retransmitHandlerLocked)) - if err != nil { - panic(fmt.Sprintf("newBackOffTimer(_, %s, %s, _) failed: %s", InitialRTO, MaxRTO, err)) - } - e.h.retransmitTimer = timer - connectingLoading.Done() + tcpip.AsyncLoading.Add(1) + go func() { + connectedLoading.Wait() + listenLoading.Wait() + // Initial SYN has been sent/received so we should bind the + // ports start the retransmit timer for the SYNs and let it + // naturally complete the connection. + bind() + e.mu.Lock() + defer e.mu.Unlock() + e.setEndpointState(epState) + r, err := e.stack.FindRoute(e.boundNICID, e.TransportEndpointInfo.ID.LocalAddress, e.TransportEndpointInfo.ID.RemoteAddress, e.effectiveNetProtos[0], false /* multicastLoop */) + if err != nil { + panic(fmt.Sprintf("FindRoute failed when restoring endpoint w/ ID: %+v", e.ID)) + } + e.route = r + timer, err := newBackoffTimer(e.stack.Clock(), InitialRTO, MaxRTO, timerHandler(e, e.h.retransmitHandlerLocked)) + if err != nil { + panic(fmt.Sprintf("newBackOffTimer(_, %s, %s, _) failed: %s", InitialRTO, MaxRTO, err)) + } + e.h.retransmitTimer = timer + connectingLoading.Done() + tcpip.AsyncLoading.Done() + }() case epState == StateBound: tcpip.AsyncLoading.Add(1) go func() { diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go index 0bd7db350..b5ad77a62 100644 --- a/pkg/tcpip/transport/udp/endpoint.go +++ b/pkg/tcpip/transport/udp/endpoint.go @@ -61,7 +61,7 @@ type endpoint struct { // The following fields are initialized at creation time and do not // change throughout the lifetime of the endpoint. - stack *stack.Stack `state:"manual"` + stack *stack.Stack waiterQueue *waiter.Queue net network.Endpoint stats tcpip.TransportEndpointStats diff --git a/pkg/tcpip/transport/udp/endpoint_state.go b/pkg/tcpip/transport/udp/endpoint_state.go index 488e46600..c8bc4b1df 100644 --- a/pkg/tcpip/transport/udp/endpoint_state.go +++ b/pkg/tcpip/transport/udp/endpoint_state.go @@ -16,7 +16,6 @@ package udp import ( "context" - "fmt" "time" "gvisor.dev/gvisor/pkg/tcpip" @@ -36,7 +35,11 @@ func (p *udpPacket) loadReceivedAt(_ context.Context, nsec int64) { // afterLoad is invoked by stateify. func (e *endpoint) afterLoad(ctx context.Context) { - stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + if e.stack.IsSaveRestoreEnabled() { + e.stack.RegisterRestoredEndpoint(e) + } else { + stack.RestoreStackFromContext(ctx).RegisterRestoredEndpoint(e) + } } // beforeSave is invoked by stateify. @@ -53,7 +56,10 @@ func (e *endpoint) Restore(s *stack.Stack) { defer e.mu.Unlock() e.net.Resume(s) - + if e.stack.IsSaveRestoreEnabled() { + e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) + return + } e.stack = s e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits) @@ -69,12 +75,12 @@ func (e *endpoint) Restore(s *stack.Stack) { id.RemotePort = e.remotePort id, e.boundBindToDevice, err = e.registerWithStack(e.effectiveNetProtos, id) if err != nil { - panic(err) + panic("registering udp endpoint with the stack failed during restore") } e.localPort = id.LocalPort e.remotePort = id.RemotePort default: - panic(fmt.Sprintf("unhandled state = %s", state)) + panic("unhandled state") } }