Fix more regressions

client/firewall/iptables/acl_linux.go

@@ -375,7 +375,7 @@ func (m *aclManager) updateState() {
 	}
 
 	var currentState *ShutdownState
-	if existing := m.stateManager.GetState(&ShutdownState{}); existing != nil {
+	if existing := m.stateManager.GetState(currentState); existing != nil {
 		if existingState, ok := existing.(*ShutdownState); ok {
 			currentState = existingState
 		}

client/firewall/iptables/router_linux.go

@@ -460,7 +460,7 @@ func (r *router) updateState() {
 	}
 
 	var currentState *ShutdownState
-	if existing := r.stateManager.GetState(&ShutdownState{}); existing != nil {
+	if existing := r.stateManager.GetState(currentState); existing != nil {
 		if existingState, ok := existing.(*ShutdownState); ok {
 			currentState = existingState
 		}

client/server/server.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
-	"github.com/hashicorp/go-multierror"
 	"golang.org/x/exp/maps"
 	"google.golang.org/protobuf/types/known/durationpb"
 
@@ -21,11 +20,7 @@ import (
 	gstatus "google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
-	nberrors "github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/client/internal/auth"
-	"github.com/netbirdio/netbird/client/internal/dns"
-	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
-	"github.com/netbirdio/netbird/client/internal/statemanager"
 	"github.com/netbirdio/netbird/client/system"
 
 	"github.com/netbirdio/netbird/client/internal"
@@ -848,31 +843,3 @@ func sendTerminalNotification() error {
 
 	return wallCmd.Wait()
 }
-
-// restoreResidulaConfig check if the client was not shut down in a clean way and restores residual if required.
-// Otherwise, we might not be able to connect to the management server to retrieve new config.
-func restoreResidualState(ctx context.Context) error {
-	path := statemanager.GetDefaultStatePath()
-	if path == "" {
-		return nil
-	}
-
-	mgr := statemanager.New(path)
-
-	var merr *multierror.Error
-
-	// register the states we are interested in restoring
-	// this will also allow each subsystem to record its own state
-	mgr.RegisterState(&dns.ShutdownState{})
-	mgr.RegisterState(&systemops.ShutdownState{})
-
-	if err := mgr.PerformCleanup(); err != nil {
-		merr = multierror.Append(merr, fmt.Errorf("perform cleanup: %w", err))
-	}
-
-	if err := mgr.PersistState(ctx); err != nil {
-		merr = multierror.Append(merr, fmt.Errorf("persist state: %w", err))
-	}
-
-	return nberrors.FormatErrorOrNil(merr)
-}

client/server/state.go

@@ -0,0 +1,37 @@
+package server
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/hashicorp/go-multierror"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	"github.com/netbirdio/netbird/client/internal/statemanager"
+)
+
+// restoreResidualConfig checks if the client was not shut down in a clean way and restores residual state if required.
+// Otherwise, we might not be able to connect to the management server to retrieve new config.
+func restoreResidualState(ctx context.Context) error {
+	path := statemanager.GetDefaultStatePath()
+	if path == "" {
+		return nil
+	}
+
+	mgr := statemanager.New(path)
+
+	// register the states we are interested in restoring
+	registerStates(mgr)
+
+	var merr *multierror.Error
+	if err := mgr.PerformCleanup(); err != nil {
+		merr = multierror.Append(merr, fmt.Errorf("perform cleanup: %w", err))
+	}
+
+	// persist state regardless of cleanup outcome. It could've succeeded partially
+	if err := mgr.PersistState(ctx); err != nil {
+		merr = multierror.Append(merr, fmt.Errorf("persist state: %w", err))
+	}
+
+	return nberrors.FormatErrorOrNil(merr)
+}