From 01c3719c5d5a2a034e03672a5da181ce34c5d8ea Mon Sep 17 00:00:00 2001
From: Viktor Liu <17948409+lixmal@users.noreply.github.com>
Date: Thu, 1 May 2025 23:25:27 +0200
Subject: [PATCH 01/11] [client] Add debug for duration option to netbird ui
 (#3772)

---
 client/cmd/debug.go    |   7 -
 client/server/debug.go |   4 +-
 client/ui/client_ui.go |  84 ++++--
 client/ui/debug.go     | 575 ++++++++++++++++++++++++++++++++++++-----
 client/ui/network.go   |  18 +-
 5 files changed, 586 insertions(+), 102 deletions(-)

diff --git a/client/cmd/debug.go b/client/cmd/debug.go
index b4adee826..385bd95f5 100644
--- a/client/cmd/debug.go
+++ b/client/cmd/debug.go
@@ -235,13 +235,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to bundle debug: %v", status.Convert(err).Message())
 	}
 
-	// Disable network map persistence after creating the debug bundle
-	if _, err := client.SetNetworkMapPersistence(cmd.Context(), &proto.SetNetworkMapPersistenceRequest{
-		Enabled: false,
-	}); err != nil {
-		return fmt.Errorf("failed to disable network map persistence: %v", status.Convert(err).Message())
-	}
-
 	if stateWasDown {
 		if _, err := client.Down(cmd.Context(), &proto.DownRequest{}); err != nil {
 			return fmt.Errorf("failed to down: %v", status.Convert(err).Message())
diff --git a/client/server/debug.go b/client/server/debug.go
index b42b1467a..7de3e8609 100644
--- a/client/server/debug.go
+++ b/client/server/debug.go
@@ -51,14 +51,16 @@ func (s *Server) DebugBundle(_ context.Context, req *proto.DebugBundleRequest) (
 	}
 
 	if req.GetUploadURL() == "" {
-
 		return &proto.DebugBundleResponse{Path: path}, nil
 	}
 	key, err := uploadDebugBundle(context.Background(), req.GetUploadURL(), s.config.ManagementURL.String(), path)
 	if err != nil {
+		log.Errorf("failed to upload debug bundle to %s: %v", req.GetUploadURL(), err)
 		return &proto.DebugBundleResponse{Path: path, UploadFailureReason: err.Error()}, nil
 	}
 
+	log.Infof("debug bundle uploaded to %s with key %s", req.GetUploadURL(), key)
+
 	return &proto.DebugBundleResponse{Path: path, UploadedKey: key}, nil
 }
 
diff --git a/client/ui/client_ui.go b/client/ui/client_ui.go
index d8c1ee7a2..2c8023185 100644
--- a/client/ui/client_ui.go
+++ b/client/ui/client_ui.go
@@ -54,11 +54,14 @@ func main() {
 	daemonAddr, showSettings, showNetworks, showDebug, errorMsg, saveLogsInFile := parseFlags()
 
 	// Initialize file logging if needed.
+	var logFile string
 	if saveLogsInFile {
-		if err := initLogFile(); err != nil {
+		file, err := initLogFile()
+		if err != nil {
 			log.Errorf("error while initializing log: %v", err)
 			return
 		}
+		logFile = file
 	}
 
 	// Create the Fyne application.
@@ -72,7 +75,7 @@ func main() {
 	}
 
 	// Create the service client (this also builds the settings or networks UI if requested).
-	client := newServiceClient(daemonAddr, a, showSettings, showNetworks, showDebug)
+	client := newServiceClient(daemonAddr, logFile, a, showSettings, showNetworks, showDebug)
 
 	// Watch for theme/settings changes to update the icon.
 	go watchSettingsChanges(a, client)
@@ -115,9 +118,9 @@ func parseFlags() (daemonAddr string, showSettings, showNetworks, showDebug bool
 }
 
 // initLogFile initializes logging into a file.
-func initLogFile() error {
+func initLogFile() (string, error) {
 	logFile := path.Join(os.TempDir(), fmt.Sprintf("netbird-ui-%d.log", os.Getpid()))
-	return util.InitLog("trace", logFile)
+	return logFile, util.InitLog("trace", logFile)
 }
 
 // watchSettingsChanges listens for Fyne theme/settings changes and updates the client icon.
@@ -160,9 +163,10 @@ var iconConnectingMacOS []byte
 var iconErrorMacOS []byte
 
 type serviceClient struct {
-	ctx  context.Context
-	addr string
-	conn proto.DaemonServiceClient
+	ctx    context.Context
+	cancel context.CancelFunc
+	addr   string
+	conn   proto.DaemonServiceClient
 
 	icAbout              []byte
 	icConnected          []byte
@@ -224,12 +228,13 @@ type serviceClient struct {
 	updateIndicationLock sync.Mutex
 	isUpdateIconActive   bool
 	showNetworks         bool
-	wRoutes              fyne.Window
+	wNetworks            fyne.Window
 
 	eventManager *event.Manager
 
 	exitNodeMu     sync.Mutex
 	mExitNodeItems []menuHandler
+	logFile        string
 }
 
 type menuHandler struct {
@@ -240,11 +245,14 @@ type menuHandler struct {
 // newServiceClient instance constructor
 //
 // This constructor also builds the UI elements for the settings window.
-func newServiceClient(addr string, a fyne.App, showSettings bool, showNetworks bool, showDebug bool) *serviceClient {
+func newServiceClient(addr string, logFile string, a fyne.App, showSettings bool, showNetworks bool, showDebug bool) *serviceClient {
+	ctx, cancel := context.WithCancel(context.Background())
 	s := &serviceClient{
-		ctx:              context.Background(),
+		ctx:              ctx,
+		cancel:           cancel,
 		addr:             addr,
 		app:              a,
+		logFile:          logFile,
 		sendNotification: false,
 
 		showAdvancedSettings: showSettings,
@@ -256,9 +264,7 @@ func newServiceClient(addr string, a fyne.App, showSettings bool, showNetworks b
 
 	switch {
 	case showSettings:
-
 		s.showSettingsUI()
-		return s
 	case showNetworks:
 		s.showNetworksUI()
 	case showDebug:
@@ -309,6 +315,8 @@ func (s *serviceClient) updateIcon() {
 func (s *serviceClient) showSettingsUI() {
 	// add settings window UI elements.
 	s.wSettings = s.app.NewWindow("NetBird Settings")
+	s.wSettings.SetOnClosed(s.cancel)
+
 	s.iMngURL = widget.NewEntry()
 	s.iAdminURL = widget.NewEntry()
 	s.iConfigFile = widget.NewEntry()
@@ -784,7 +792,7 @@ func (s *serviceClient) onTrayReady() {
 func (s *serviceClient) runSelfCommand(command, arg string) {
 	proc, err := os.Executable()
 	if err != nil {
-		log.Errorf("show %s failed with error: %v", command, err)
+		log.Errorf("Error getting executable path: %v", err)
 		return
 	}
 
@@ -793,14 +801,48 @@ func (s *serviceClient) runSelfCommand(command, arg string) {
 		fmt.Sprintf("--daemon-addr=%s", s.addr),
 	)
 
-	out, err := cmd.CombinedOutput()
-	if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() == 1 {
-		log.Errorf("start %s UI: %v, %s", command, err, string(out))
+	if out := s.attachOutput(cmd); out != nil {
+		defer func() {
+			if err := out.Close(); err != nil {
+				log.Errorf("Error closing log file %s: %v", s.logFile, err)
+			}
+		}()
+	}
+
+	log.Printf("Running command: %s --%s=%s --daemon-addr=%s", proc, command, arg, s.addr)
+
+	err = cmd.Run()
+
+	if err != nil {
+		var exitErr *exec.ExitError
+		if errors.As(err, &exitErr) {
+			log.Printf("Command '%s %s' failed with exit code %d", command, arg, exitErr.ExitCode())
+		} else {
+			log.Printf("Failed to start/run command '%s %s': %v", command, arg, err)
+		}
 		return
 	}
-	if len(out) != 0 {
-		log.Infof("command %s executed: %s", command, string(out))
+
+	log.Printf("Command '%s %s' completed successfully.", command, arg)
+}
+
+func (s *serviceClient) attachOutput(cmd *exec.Cmd) *os.File {
+	if s.logFile == "" {
+		// attach child's streams to parent's streams
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+
+		return nil
 	}
+
+	out, err := os.OpenFile(s.logFile, os.O_WRONLY|os.O_APPEND, 0)
+	if err != nil {
+		log.Errorf("Failed to open log file %s: %v", s.logFile, err)
+		return nil
+	}
+	cmd.Stdout = out
+	cmd.Stderr = out
+	return out
 }
 
 func normalizedVersion(version string) string {
@@ -813,9 +855,7 @@ func normalizedVersion(version string) string {
 
 // onTrayExit is called when the tray icon is closed.
 func (s *serviceClient) onTrayExit() {
-	for _, item := range s.mExitNodeItems {
-		item.cancel()
-	}
+	s.cancel()
 }
 
 // getSrvClient connection to the service.
@@ -824,7 +864,7 @@ func (s *serviceClient) getSrvClient(timeout time.Duration) (proto.DaemonService
 		return s.conn, nil
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	ctx, cancel := context.WithTimeout(s.ctx, timeout)
 	defer cancel()
 
 	conn, err := grpc.DialContext(
diff --git a/client/ui/debug.go b/client/ui/debug.go
index e950e6d1e..ab7dba37a 100644
--- a/client/ui/debug.go
+++ b/client/ui/debug.go
@@ -3,8 +3,12 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"path/filepath"
+	"strconv"
+	"sync"
+	"time"
 
 	"fyne.io/fyne/v2"
 	"fyne.io/fyne/v2/container"
@@ -13,18 +17,46 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/skratchdot/open-golang/open"
 
+	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/proto"
 	nbstatus "github.com/netbirdio/netbird/client/status"
 	uptypes "github.com/netbirdio/netbird/upload-server/types"
 )
 
+// Initial state for the debug collection
+type debugInitialState struct {
+	wasDown      bool
+	logLevel     proto.LogLevel
+	isLevelTrace bool
+}
+
+// Debug collection parameters
+type debugCollectionParams struct {
+	duration          time.Duration
+	anonymize         bool
+	systemInfo        bool
+	upload            bool
+	uploadURL         string
+	enablePersistence bool
+}
+
+// UI components for progress tracking
+type progressUI struct {
+	statusLabel *widget.Label
+	progressBar *widget.ProgressBar
+	uiControls  []fyne.Disableable
+	window      fyne.Window
+}
+
 func (s *serviceClient) showDebugUI() {
 	w := s.app.NewWindow("NetBird Debug")
-	w.Resize(fyne.NewSize(600, 400))
+	w.SetOnClosed(s.cancel)
+
+	w.Resize(fyne.NewSize(600, 500))
 	w.SetFixedSize(true)
 
-	anonymizeCheck := widget.NewCheck("Anonymize sensitive information (Public IPs, domains, ...)", nil)
-	systemInfoCheck := widget.NewCheck("Include system information", nil)
+	anonymizeCheck := widget.NewCheck("Anonymize sensitive information (public IPs, domains, ...)", nil)
+	systemInfoCheck := widget.NewCheck("Include system information (routes, interfaces, ...)", nil)
 	systemInfoCheck.SetChecked(true)
 	uploadCheck := widget.NewCheck("Upload bundle automatically after creation", nil)
 	uploadCheck.SetChecked(true)
@@ -34,11 +66,6 @@ func (s *serviceClient) showDebugUI() {
 	uploadURL.SetText(uptypes.DefaultBundleURL)
 	uploadURL.SetPlaceHolder("Enter upload URL")
 
-	statusLabel := widget.NewLabel("")
-	statusLabel.Hide()
-
-	createButton := widget.NewButton("Create Debug Bundle", nil)
-
 	uploadURLContainer := container.NewVBox(
 		uploadURLLabel,
 		uploadURL,
@@ -52,7 +79,71 @@ func (s *serviceClient) showDebugUI() {
 		}
 	}
 
-	createButton.OnTapped = s.getCreateHandler(createButton, statusLabel, uploadCheck, uploadURL, anonymizeCheck, systemInfoCheck, w)
+	debugModeContainer := container.NewHBox()
+	runForDurationCheck := widget.NewCheck("Run with trace logs before creating bundle", nil)
+	runForDurationCheck.SetChecked(true)
+
+	forLabel := widget.NewLabel("for")
+
+	durationInput := widget.NewEntry()
+	durationInput.SetText("1")
+	minutesLabel := widget.NewLabel("minute")
+	durationInput.Validator = func(s string) error {
+		return validateMinute(s, minutesLabel)
+	}
+
+	noteLabel := widget.NewLabel("Note: NetBird will be brought up and down during collection")
+
+	runForDurationCheck.OnChanged = func(checked bool) {
+		if checked {
+			forLabel.Show()
+			durationInput.Show()
+			minutesLabel.Show()
+			noteLabel.Show()
+		} else {
+			forLabel.Hide()
+			durationInput.Hide()
+			minutesLabel.Hide()
+			noteLabel.Hide()
+		}
+	}
+
+	debugModeContainer.Add(runForDurationCheck)
+	debugModeContainer.Add(forLabel)
+	debugModeContainer.Add(durationInput)
+	debugModeContainer.Add(minutesLabel)
+
+	statusLabel := widget.NewLabel("")
+	statusLabel.Hide()
+
+	progressBar := widget.NewProgressBar()
+	progressBar.Hide()
+
+	createButton := widget.NewButton("Create Debug Bundle", nil)
+
+	// UI controls that should be disabled during debug collection
+	uiControls := []fyne.Disableable{
+		anonymizeCheck,
+		systemInfoCheck,
+		uploadCheck,
+		uploadURL,
+		runForDurationCheck,
+		durationInput,
+		createButton,
+	}
+
+	createButton.OnTapped = s.getCreateHandler(
+		statusLabel,
+		progressBar,
+		uploadCheck,
+		uploadURL,
+		anonymizeCheck,
+		systemInfoCheck,
+		runForDurationCheck,
+		durationInput,
+		uiControls,
+		w,
+	)
 
 	content := container.NewVBox(
 		widget.NewLabel("Create a debug bundle to help troubleshoot issues with NetBird"),
@@ -62,7 +153,11 @@ func (s *serviceClient) showDebugUI() {
 		uploadCheck,
 		uploadURLContainer,
 		widget.NewLabel(""),
+		debugModeContainer,
+		noteLabel,
+		widget.NewLabel(""),
 		statusLabel,
+		progressBar,
 		createButton,
 	)
 
@@ -72,18 +167,46 @@ func (s *serviceClient) showDebugUI() {
 	w.Show()
 }
 
+func validateMinute(s string, minutesLabel *widget.Label) error {
+	if val, err := strconv.Atoi(s); err != nil || val < 1 {
+		return fmt.Errorf("must be a number ≥ 1")
+	}
+	if s == "1" {
+		minutesLabel.SetText("minute")
+	} else {
+		minutesLabel.SetText("minutes")
+	}
+	return nil
+}
+
+// disableUIControls disables the provided UI controls
+func disableUIControls(controls []fyne.Disableable) {
+	for _, control := range controls {
+		control.Disable()
+	}
+}
+
+// enableUIControls enables the provided UI controls
+func enableUIControls(controls []fyne.Disableable) {
+	for _, control := range controls {
+		control.Enable()
+	}
+}
+
 func (s *serviceClient) getCreateHandler(
-	createButton *widget.Button,
 	statusLabel *widget.Label,
+	progressBar *widget.ProgressBar,
 	uploadCheck *widget.Check,
 	uploadURL *widget.Entry,
 	anonymizeCheck *widget.Check,
 	systemInfoCheck *widget.Check,
+	runForDurationCheck *widget.Check,
+	duration *widget.Entry,
+	uiControls []fyne.Disableable,
 	w fyne.Window,
 ) func() {
 	return func() {
-		createButton.Disable()
-		statusLabel.SetText("Creating debug bundle...")
+		disableUIControls(uiControls)
 		statusLabel.Show()
 
 		var url string
@@ -91,22 +214,329 @@ func (s *serviceClient) getCreateHandler(
 			url = uploadURL.Text
 			if url == "" {
 				statusLabel.SetText("Error: Upload URL is required when upload is enabled")
-				createButton.Enable()
+				enableUIControls(uiControls)
 				return
 			}
 		}
 
-		go s.handleDebugCreation(anonymizeCheck.Checked, systemInfoCheck.Checked, uploadCheck.Checked, url, statusLabel, createButton, w)
+		params := &debugCollectionParams{
+			anonymize:         anonymizeCheck.Checked,
+			systemInfo:        systemInfoCheck.Checked,
+			upload:            uploadCheck.Checked,
+			uploadURL:         url,
+			enablePersistence: true,
+		}
+
+		runForDuration := runForDurationCheck.Checked
+		if runForDuration {
+			minutes, err := time.ParseDuration(duration.Text + "m")
+			if err != nil {
+				statusLabel.SetText(fmt.Sprintf("Error: Invalid duration: %v", err))
+				enableUIControls(uiControls)
+				return
+			}
+			params.duration = minutes
+
+			statusLabel.SetText(fmt.Sprintf("Running in debug mode for %d minutes...", int(minutes.Minutes())))
+			progressBar.Show()
+			progressBar.SetValue(0)
+
+			go s.handleRunForDuration(
+				statusLabel,
+				progressBar,
+				uiControls,
+				w,
+				params,
+			)
+			return
+		}
+
+		statusLabel.SetText("Creating debug bundle...")
+		go s.handleDebugCreation(
+			anonymizeCheck.Checked,
+			systemInfoCheck.Checked,
+			uploadCheck.Checked,
+			url,
+			statusLabel,
+			uiControls,
+			w,
+		)
 	}
 }
 
+func (s *serviceClient) handleRunForDuration(
+	statusLabel *widget.Label,
+	progressBar *widget.ProgressBar,
+	uiControls []fyne.Disableable,
+	w fyne.Window,
+	params *debugCollectionParams,
+) {
+	progressUI := &progressUI{
+		statusLabel: statusLabel,
+		progressBar: progressBar,
+		uiControls:  uiControls,
+		window:      w,
+	}
+
+	conn, err := s.getSrvClient(failFastTimeout)
+	if err != nil {
+		handleError(progressUI, fmt.Sprintf("Failed to get client for debug: %v", err))
+		return
+	}
+
+	initialState, err := s.getInitialState(conn)
+	if err != nil {
+		handleError(progressUI, err.Error())
+		return
+	}
+
+	statusOutput, err := s.collectDebugData(conn, initialState, params, progressUI)
+	if err != nil {
+		handleError(progressUI, err.Error())
+		return
+	}
+
+	if err := s.createDebugBundleFromCollection(conn, params, statusOutput, progressUI); err != nil {
+		handleError(progressUI, err.Error())
+		return
+	}
+
+	s.restoreServiceState(conn, initialState)
+
+	progressUI.statusLabel.SetText("Bundle created successfully")
+}
+
+// Get initial state of the service
+func (s *serviceClient) getInitialState(conn proto.DaemonServiceClient) (*debugInitialState, error) {
+	statusResp, err := conn.Status(s.ctx, &proto.StatusRequest{})
+	if err != nil {
+		return nil, fmt.Errorf(" get status: %v", err)
+	}
+
+	logLevelResp, err := conn.GetLogLevel(s.ctx, &proto.GetLogLevelRequest{})
+	if err != nil {
+		return nil, fmt.Errorf("get log level: %v", err)
+	}
+
+	wasDown := statusResp.Status != string(internal.StatusConnected) &&
+		statusResp.Status != string(internal.StatusConnecting)
+
+	initialLogLevel := logLevelResp.GetLevel()
+	initialLevelTrace := initialLogLevel >= proto.LogLevel_TRACE
+
+	return &debugInitialState{
+		wasDown:      wasDown,
+		logLevel:     initialLogLevel,
+		isLevelTrace: initialLevelTrace,
+	}, nil
+}
+
+// Handle progress tracking during collection
+func startProgressTracker(ctx context.Context, wg *sync.WaitGroup, duration time.Duration, progress *progressUI) {
+	progress.progressBar.Show()
+	progress.progressBar.SetValue(0)
+
+	startTime := time.Now()
+	endTime := startTime.Add(duration)
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+		ticker := time.NewTicker(500 * time.Millisecond)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				remaining := time.Until(endTime)
+				if remaining <= 0 {
+					remaining = 0
+				}
+
+				elapsed := time.Since(startTime)
+				progressVal := float64(elapsed) / float64(duration)
+				if progressVal > 1.0 {
+					progressVal = 1.0
+				}
+
+				progress.progressBar.SetValue(progressVal)
+				progress.statusLabel.SetText(fmt.Sprintf("Running with trace logs... %s remaining", formatDuration(remaining)))
+			}
+		}
+	}()
+
+}
+
+func (s *serviceClient) configureServiceForDebug(
+	conn proto.DaemonServiceClient,
+	state *debugInitialState,
+	enablePersistence bool,
+) error {
+	if state.wasDown {
+		if _, err := conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
+			return fmt.Errorf("bring service up: %v", err)
+		}
+		log.Info("Service brought up for debug")
+		time.Sleep(time.Second * 10)
+	}
+
+	if !state.isLevelTrace {
+		if _, err := conn.SetLogLevel(s.ctx, &proto.SetLogLevelRequest{Level: proto.LogLevel_TRACE}); err != nil {
+			return fmt.Errorf("set log level to TRACE: %v", err)
+		}
+		log.Info("Log level set to TRACE for debug")
+	}
+
+	if _, err := conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
+		return fmt.Errorf("bring service down: %v", err)
+	}
+	time.Sleep(time.Second)
+
+	if enablePersistence {
+		if _, err := conn.SetNetworkMapPersistence(s.ctx, &proto.SetNetworkMapPersistenceRequest{
+			Enabled: true,
+		}); err != nil {
+			return fmt.Errorf("enable network map persistence: %v", err)
+		}
+		log.Info("Network map persistence enabled for debug")
+	}
+
+	if _, err := conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
+		return fmt.Errorf("bring service back up: %v", err)
+	}
+	time.Sleep(time.Second * 3)
+
+	return nil
+}
+
+func (s *serviceClient) collectDebugData(
+	conn proto.DaemonServiceClient,
+	state *debugInitialState,
+	params *debugCollectionParams,
+	progress *progressUI,
+) (string, error) {
+	ctx, cancel := context.WithTimeout(s.ctx, params.duration)
+	defer cancel()
+	var wg sync.WaitGroup
+	startProgressTracker(ctx, &wg, params.duration, progress)
+
+	if err := s.configureServiceForDebug(conn, state, params.enablePersistence); err != nil {
+		return "", err
+	}
+
+	postUpStatus, err := conn.Status(s.ctx, &proto.StatusRequest{GetFullPeerStatus: true})
+	if err != nil {
+		log.Warnf("Failed to get post-up status: %v", err)
+	}
+
+	var postUpStatusOutput string
+	if postUpStatus != nil {
+		overview := nbstatus.ConvertToStatusOutputOverview(postUpStatus, params.anonymize, "", nil, nil, nil)
+		postUpStatusOutput = nbstatus.ParseToFullDetailSummary(overview)
+	}
+	headerPostUp := fmt.Sprintf("----- NetBird post-up - Timestamp: %s", time.Now().Format(time.RFC3339))
+	statusOutput := fmt.Sprintf("%s\n%s", headerPostUp, postUpStatusOutput)
+
+	wg.Wait()
+	progress.progressBar.Hide()
+	progress.statusLabel.SetText("Collecting debug data...")
+
+	preDownStatus, err := conn.Status(s.ctx, &proto.StatusRequest{GetFullPeerStatus: true})
+	if err != nil {
+		log.Warnf("Failed to get pre-down status: %v", err)
+	}
+
+	var preDownStatusOutput string
+	if preDownStatus != nil {
+		overview := nbstatus.ConvertToStatusOutputOverview(preDownStatus, params.anonymize, "", nil, nil, nil)
+		preDownStatusOutput = nbstatus.ParseToFullDetailSummary(overview)
+	}
+	headerPreDown := fmt.Sprintf("----- NetBird pre-down - Timestamp: %s - Duration: %s",
+		time.Now().Format(time.RFC3339), params.duration)
+	statusOutput = fmt.Sprintf("%s\n%s\n%s", statusOutput, headerPreDown, preDownStatusOutput)
+
+	return statusOutput, nil
+}
+
+// Create the debug bundle with collected data
+func (s *serviceClient) createDebugBundleFromCollection(
+	conn proto.DaemonServiceClient,
+	params *debugCollectionParams,
+	statusOutput string,
+	progress *progressUI,
+) error {
+	progress.statusLabel.SetText("Creating debug bundle with collected logs...")
+
+	request := &proto.DebugBundleRequest{
+		Anonymize:  params.anonymize,
+		Status:     statusOutput,
+		SystemInfo: params.systemInfo,
+	}
+
+	if params.upload {
+		request.UploadURL = params.uploadURL
+	}
+
+	resp, err := conn.DebugBundle(s.ctx, request)
+	if err != nil {
+		return fmt.Errorf("create debug bundle: %v", err)
+	}
+
+	// Show appropriate dialog based on upload status
+	localPath := resp.GetPath()
+	uploadFailureReason := resp.GetUploadFailureReason()
+	uploadedKey := resp.GetUploadedKey()
+
+	if params.upload {
+		if uploadFailureReason != "" {
+			showUploadFailedDialog(progress.window, localPath, uploadFailureReason)
+		} else {
+			showUploadSuccessDialog(progress.window, localPath, uploadedKey)
+		}
+	} else {
+		showBundleCreatedDialog(progress.window, localPath)
+	}
+
+	enableUIControls(progress.uiControls)
+	return nil
+}
+
+// Restore service to original state
+func (s *serviceClient) restoreServiceState(conn proto.DaemonServiceClient, state *debugInitialState) {
+	if state.wasDown {
+		if _, err := conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
+			log.Errorf("Failed to restore down state: %v", err)
+		} else {
+			log.Info("Service state restored to down")
+		}
+	}
+
+	if !state.isLevelTrace {
+		if _, err := conn.SetLogLevel(s.ctx, &proto.SetLogLevelRequest{Level: state.logLevel}); err != nil {
+			log.Errorf("Failed to restore log level: %v", err)
+		} else {
+			log.Info("Log level restored to original setting")
+		}
+	}
+}
+
+// Handle errors during debug collection
+func handleError(progress *progressUI, errMsg string) {
+	log.Errorf("%s", errMsg)
+	progress.statusLabel.SetText(errMsg)
+	progress.progressBar.Hide()
+	enableUIControls(progress.uiControls)
+}
+
 func (s *serviceClient) handleDebugCreation(
 	anonymize bool,
 	systemInfo bool,
 	upload bool,
 	uploadURL string,
 	statusLabel *widget.Label,
-	createButton *widget.Button,
+	uiControls []fyne.Disableable,
 	w fyne.Window,
 ) {
 	log.Infof("Creating debug bundle (Anonymized: %v, System Info: %v, Upload Attempt: %v)...",
@@ -116,7 +546,7 @@ func (s *serviceClient) handleDebugCreation(
 	if err != nil {
 		log.Errorf("Failed to create debug bundle: %v", err)
 		statusLabel.SetText(fmt.Sprintf("Error creating bundle: %v", err))
-		createButton.Enable()
+		enableUIControls(uiControls)
 		return
 	}
 
@@ -134,7 +564,7 @@ func (s *serviceClient) handleDebugCreation(
 		showBundleCreatedDialog(w, localPath)
 	}
 
-	createButton.Enable()
+	enableUIControls(uiControls)
 	statusLabel.SetText("Bundle created successfully")
 }
 
@@ -173,32 +603,47 @@ func (s *serviceClient) createDebugBundle(anonymize bool, systemInfo bool, uploa
 	return resp, nil
 }
 
+// formatDuration formats a duration in HH:MM:SS format
+func formatDuration(d time.Duration) string {
+	d = d.Round(time.Second)
+	h := d / time.Hour
+	d %= time.Hour
+	m := d / time.Minute
+	d %= time.Minute
+	s := d / time.Second
+	return fmt.Sprintf("%02d:%02d:%02d", h, m, s)
+}
+
+// createButtonWithAction creates a button with the given label and action
+func createButtonWithAction(label string, action func()) *widget.Button {
+	button := widget.NewButton(label, action)
+	return button
+}
+
 // showUploadFailedDialog displays a dialog when upload fails
-func showUploadFailedDialog(parent fyne.Window, localPath, failureReason string) {
+func showUploadFailedDialog(w fyne.Window, localPath, failureReason string) {
 	content := container.NewVBox(
 		widget.NewLabel(fmt.Sprintf("Bundle upload failed:\n%s\n\n"+
 			"A local copy was saved at:\n%s", failureReason, localPath)),
 	)
 
-	customDialog := dialog.NewCustom("Upload Failed", "Cancel", content, parent)
+	customDialog := dialog.NewCustom("Upload Failed", "Cancel", content, w)
 
 	buttonBox := container.NewHBox(
-		widget.NewButton("Open File", func() {
+		createButtonWithAction("Open file", func() {
 			log.Infof("Attempting to open local file: %s", localPath)
 			if openErr := open.Start(localPath); openErr != nil {
 				log.Errorf("Failed to open local file '%s': %v", localPath, openErr)
-				dialog.ShowError(fmt.Errorf("Failed to open the local file:\n%s\n\nError: %v", localPath, openErr), parent)
+				dialog.ShowError(fmt.Errorf("open the local file:\n%s\n\nError: %v", localPath, openErr), w)
 			}
-			customDialog.Hide()
 		}),
-		widget.NewButton("Open Folder", func() {
+		createButtonWithAction("Open folder", func() {
 			folderPath := filepath.Dir(localPath)
 			log.Infof("Attempting to open local folder: %s", folderPath)
 			if openErr := open.Start(folderPath); openErr != nil {
 				log.Errorf("Failed to open local folder '%s': %v", folderPath, openErr)
-				dialog.ShowError(fmt.Errorf("Failed to open the local folder:\n%s\n\nError: %v", folderPath, openErr), parent)
+				dialog.ShowError(fmt.Errorf("open the local folder:\n%s\n\nError: %v", folderPath, openErr), w)
 			}
-			customDialog.Hide()
 		}),
 	)
 
@@ -207,7 +652,8 @@ func showUploadFailedDialog(parent fyne.Window, localPath, failureReason string)
 }
 
 // showUploadSuccessDialog displays a dialog when upload succeeds
-func showUploadSuccessDialog(parent fyne.Window, localPath, uploadedKey string) {
+func showUploadSuccessDialog(w fyne.Window, localPath, uploadedKey string) {
+	log.Infof("Upload key: %s", uploadedKey)
 	keyEntry := widget.NewEntry()
 	keyEntry.SetText(uploadedKey)
 	keyEntry.Disable()
@@ -215,62 +661,63 @@ func showUploadSuccessDialog(parent fyne.Window, localPath, uploadedKey string)
 	content := container.NewVBox(
 		widget.NewLabel("Bundle uploaded successfully!"),
 		widget.NewLabel(""),
-		widget.NewLabel("Upload Key:"),
+		widget.NewLabel("Upload key:"),
 		keyEntry,
 		widget.NewLabel(""),
 		widget.NewLabel(fmt.Sprintf("Local copy saved at:\n%s", localPath)),
 	)
 
-	customDialog := dialog.NewCustom("Upload Successful", "OK", content, parent)
+	customDialog := dialog.NewCustom("Upload Successful", "OK", content, w)
 
-	buttonBox := container.NewHBox(
-		widget.NewButton("Copy Key", func() {
-			parent.Clipboard().SetContent(uploadedKey)
-			log.Info("Upload key copied to clipboard")
-		}),
-		widget.NewButton("Open Local Folder", func() {
-			folderPath := filepath.Dir(localPath)
-			log.Infof("Attempting to open local folder: %s", folderPath)
-			if openErr := open.Start(folderPath); openErr != nil {
-				log.Errorf("Failed to open local folder '%s': %v", folderPath, openErr)
-				dialog.ShowError(fmt.Errorf("Failed to open the local folder:\n%s\n\nError: %v", folderPath, openErr), parent)
-			}
-		}),
-	)
+	copyBtn := createButtonWithAction("Copy key", func() {
+		w.Clipboard().SetContent(uploadedKey)
+		log.Info("Upload key copied to clipboard")
+	})
 
+	buttonBox := createButtonBox(localPath, w, copyBtn)
 	content.Add(buttonBox)
 	customDialog.Show()
 }
 
 // showBundleCreatedDialog displays a dialog when bundle is created without upload
-func showBundleCreatedDialog(parent fyne.Window, localPath string) {
+func showBundleCreatedDialog(w fyne.Window, localPath string) {
 	content := container.NewVBox(
 		widget.NewLabel(fmt.Sprintf("Bundle created locally at:\n%s\n\n"+
 			"Administrator privileges may be required to access the file.", localPath)),
 	)
 
-	customDialog := dialog.NewCustom("Debug Bundle Created", "Cancel", content, parent)
-
-	buttonBox := container.NewHBox(
-		widget.NewButton("Open File", func() {
-			log.Infof("Attempting to open local file: %s", localPath)
-			if openErr := open.Start(localPath); openErr != nil {
-				log.Errorf("Failed to open local file '%s': %v", localPath, openErr)
-				dialog.ShowError(fmt.Errorf("Failed to open the local file:\n%s\n\nError: %v", localPath, openErr), parent)
-			}
-			customDialog.Hide()
-		}),
-		widget.NewButton("Open Folder", func() {
-			folderPath := filepath.Dir(localPath)
-			log.Infof("Attempting to open local folder: %s", folderPath)
-			if openErr := open.Start(folderPath); openErr != nil {
-				log.Errorf("Failed to open local folder '%s': %v", folderPath, openErr)
-				dialog.ShowError(fmt.Errorf("Failed to open the local folder:\n%s\n\nError: %v", folderPath, openErr), parent)
-			}
-			customDialog.Hide()
-		}),
-	)
+	customDialog := dialog.NewCustom("Debug Bundle Created", "Cancel", content, w)
 
+	buttonBox := createButtonBox(localPath, w, nil)
 	content.Add(buttonBox)
 	customDialog.Show()
 }
+
+func createButtonBox(localPath string, w fyne.Window, elems ...fyne.Widget) *fyne.Container {
+	box := container.NewHBox()
+	for _, elem := range elems {
+		box.Add(elem)
+	}
+
+	fileBtn := createButtonWithAction("Open file", func() {
+		log.Infof("Attempting to open local file: %s", localPath)
+		if openErr := open.Start(localPath); openErr != nil {
+			log.Errorf("Failed to open local file '%s': %v", localPath, openErr)
+			dialog.ShowError(fmt.Errorf("open the local file:\n%s\n\nError: %v", localPath, openErr), w)
+		}
+	})
+
+	folderBtn := createButtonWithAction("Open folder", func() {
+		folderPath := filepath.Dir(localPath)
+		log.Infof("Attempting to open local folder: %s", folderPath)
+		if openErr := open.Start(folderPath); openErr != nil {
+			log.Errorf("Failed to open local folder '%s': %v", folderPath, openErr)
+			dialog.ShowError(fmt.Errorf("open the local folder:\n%s\n\nError: %v", folderPath, openErr), w)
+		}
+	})
+
+	box.Add(fileBtn)
+	box.Add(folderBtn)
+
+	return box
+}
diff --git a/client/ui/network.go b/client/ui/network.go
index ddd8d5000..435917f30 100644
--- a/client/ui/network.go
+++ b/client/ui/network.go
@@ -34,7 +34,8 @@ const (
 type filter string
 
 func (s *serviceClient) showNetworksUI() {
-	s.wRoutes = s.app.NewWindow("Networks")
+	s.wNetworks = s.app.NewWindow("Networks")
+	s.wNetworks.SetOnClosed(s.cancel)
 
 	allGrid := container.New(layout.NewGridLayout(3))
 	go s.updateNetworks(allGrid, allNetworks)
@@ -78,8 +79,8 @@ func (s *serviceClient) showNetworksUI() {
 
 	content := container.NewBorder(nil, buttonBox, nil, nil, scrollContainer)
 
-	s.wRoutes.SetContent(content)
-	s.wRoutes.Show()
+	s.wNetworks.SetContent(content)
+	s.wNetworks.Show()
 
 	s.startAutoRefresh(10*time.Second, tabs, allGrid, overlappingGrid, exitNodeGrid)
 }
@@ -148,7 +149,7 @@ func (s *serviceClient) updateNetworks(grid *fyne.Container, f filter) {
 		grid.Add(resolvedIPsSelector)
 	}
 
-	s.wRoutes.Content().Refresh()
+	s.wNetworks.Content().Refresh()
 	grid.Refresh()
 }
 
@@ -305,7 +306,7 @@ func (s *serviceClient) getNetworksRequest(f filter, appendRoute bool) *proto.Se
 func (s *serviceClient) showError(err error) {
 	wrappedMessage := wrapText(err.Error(), 50)
 
-	dialog.ShowError(fmt.Errorf("%s", wrappedMessage), s.wRoutes)
+	dialog.ShowError(fmt.Errorf("%s", wrappedMessage), s.wNetworks)
 }
 
 func (s *serviceClient) startAutoRefresh(interval time.Duration, tabs *container.AppTabs, allGrid, overlappingGrid, exitNodesGrid *fyne.Container) {
@@ -316,14 +317,15 @@ func (s *serviceClient) startAutoRefresh(interval time.Duration, tabs *container
 		}
 	}()
 
-	s.wRoutes.SetOnClosed(func() {
+	s.wNetworks.SetOnClosed(func() {
 		ticker.Stop()
+		s.cancel()
 	})
 }
 
 func (s *serviceClient) updateNetworksBasedOnDisplayTab(tabs *container.AppTabs, allGrid, overlappingGrid, exitNodesGrid *fyne.Container) {
 	grid, f := getGridAndFilterFromTab(tabs, allGrid, overlappingGrid, exitNodesGrid)
-	s.wRoutes.Content().Refresh()
+	s.wNetworks.Content().Refresh()
 	s.updateNetworks(grid, f)
 }
 
@@ -373,7 +375,7 @@ func (s *serviceClient) recreateExitNodeMenu(exitNodes []*proto.Network) {
 			node.Selected,
 		)
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(s.ctx)
 		s.mExitNodeItems = append(s.mExitNodeItems, menuHandler{
 			MenuItem: menuItem,
 			cancel:   cancel,

From 2abb92b0d4c0bf9aa3a2c45863c0872f31f2908f Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Fri, 2 May 2025 00:25:46 +0200
Subject: [PATCH 02/11] [management] Get account id with order (#3773)

updated log to display account id
---
 management/server/account.go         | 2 +-
 management/server/store/sql_store.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index ab1ffe8b3..869b13f59 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -721,7 +721,7 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID any)
 	if err != nil {
 		return nil, nil, err
 	}
-	log.WithContext(ctx).Debugf("%d entries received from IdP management", len(userData))
+	log.WithContext(ctx).Debugf("%d entries received from IdP management for account %s", len(userData), account.Id)
 
 	dataMap := make(map[string]*idp.UserData, len(userData))
 	for _, datum := range userData {
diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go
index 7d3b288e0..dd39cf77d 100644
--- a/management/server/store/sql_store.go
+++ b/management/server/store/sql_store.go
@@ -802,7 +802,7 @@ func (s *SqlStore) GetAccountByPeerPubKey(ctx context.Context, peerKey string) (
 
 func (s *SqlStore) GetAnyAccountID(ctx context.Context) (string, error) {
 	var account types.Account
-	result := s.db.WithContext(ctx).Select("id").Limit(1).Find(&account)
+	result := s.db.WithContext(ctx).Select("id").Order("created_at desc").Limit(1).Find(&account)
 	if result.Error != nil {
 		return "", status.NewGetAccountFromStoreError(result.Error)
 	}

From 12f883badfe00ec04ad8bd3d5dc96dca1c786289 Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Fri, 2 May 2025 00:59:41 +0200
Subject: [PATCH 03/11] [management] Optimize load account (#3774)

---
 management/server/account.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index 869b13f59..aa7cb0019 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -712,7 +712,7 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID any)
 	log.WithContext(ctx).Debugf("account %s not found in cache, reloading", accountID)
 	accountIDString := fmt.Sprintf("%v", accountID)
 
-	account, err := am.Store.GetAccount(ctx, accountIDString)
+	accountUsers, err := am.Store.GetAccountUsers(ctx, store.LockingStrengthShare, accountIDString)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -721,7 +721,7 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID any)
 	if err != nil {
 		return nil, nil, err
 	}
-	log.WithContext(ctx).Debugf("%d entries received from IdP management for account %s", len(userData), account.Id)
+	log.WithContext(ctx).Debugf("%d entries received from IdP management for account %s", len(userData), accountIDString)
 
 	dataMap := make(map[string]*idp.UserData, len(userData))
 	for _, datum := range userData {
@@ -729,7 +729,7 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID any)
 	}
 
 	matchedUserData := make([]*idp.UserData, 0)
-	for _, user := range account.Users {
+	for _, user := range accountUsers {
 		if user.IsServiceUser {
 			continue
 		}

From 055df9854c8af522919f91eed59c8f12b730cbe5 Mon Sep 17 00:00:00 2001
From: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com>
Date: Sun, 4 May 2025 20:58:04 +0200
Subject: [PATCH 04/11] [management] add gorm tag for primary key for the
 networks objects (#3758)

---
 management/server/migration/migration.go      | 21 ++++++++++++++++++
 management/server/migration/migration_test.go | 22 +++++++++++++++++++
 .../networks/resources/types/resource.go      |  2 +-
 .../server/networks/routers/types/router.go   |  2 +-
 management/server/networks/types/network.go   |  2 +-
 management/server/store/sql_store.go          |  4 ++++
 management/server/store/sql_store_test.go     |  6 ++---
 management/server/store/store.go              |  9 ++++++++
 management/server/types/group.go              |  2 +-
 9 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/management/server/migration/migration.go b/management/server/migration/migration.go
index d7abbad47..c8a852e0a 100644
--- a/management/server/migration/migration.go
+++ b/management/server/migration/migration.go
@@ -352,3 +352,24 @@ func MigrateNewField[T any](ctx context.Context, db *gorm.DB, columnName string,
 	log.WithContext(ctx).Infof("Migration of empty %s to default value in table %s completed", columnName, tableName)
 	return nil
 }
+
+func DropIndex[T any](ctx context.Context, db *gorm.DB, indexName string) error {
+	var model T
+
+	if !db.Migrator().HasTable(&model) {
+		log.WithContext(ctx).Debugf("table for %T does not exist, no migration needed", model)
+		return nil
+	}
+
+	if !db.Migrator().HasIndex(&model, indexName) {
+		log.WithContext(ctx).Debugf("index %s does not exist in table %T, no migration needed", indexName, model)
+		return nil
+	}
+
+	if err := db.Migrator().DropIndex(&model, indexName); err != nil {
+		return fmt.Errorf("failed to drop index %s: %w", indexName, err)
+	}
+
+	log.WithContext(ctx).Infof("dropped index %s from table %T", indexName, model)
+	return nil
+}
diff --git a/management/server/migration/migration_test.go b/management/server/migration/migration_test.go
index e907d6853..94377930a 100644
--- a/management/server/migration/migration_test.go
+++ b/management/server/migration/migration_test.go
@@ -227,3 +227,25 @@ func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case2(t *testing.
 
 	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
 }
+
+func TestDropIndex(t *testing.T) {
+	db := setupDatabase(t)
+
+	err := db.AutoMigrate(&types.SetupKey{})
+	require.NoError(t, err, "Failed to auto-migrate tables")
+
+	err = db.Save(&types.SetupKey{
+		Id:  "1",
+		Key: "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+	}).Error
+	require.NoError(t, err, "Failed to insert setup key")
+
+	exist := db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
+	assert.True(t, exist, "Should have the index")
+
+	err = migration.DropIndex[types.SetupKey](context.Background(), db, "idx_setup_keys_account_id")
+	require.NoError(t, err, "Migration should not fail to remove index")
+
+	exist = db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
+	assert.False(t, exist, "Should not have the index")
+}
diff --git a/management/server/networks/resources/types/resource.go b/management/server/networks/resources/types/resource.go
index ecac0a724..04c63608d 100644
--- a/management/server/networks/resources/types/resource.go
+++ b/management/server/networks/resources/types/resource.go
@@ -30,7 +30,7 @@ func (p NetworkResourceType) String() string {
 }
 
 type NetworkResource struct {
-	ID          string `gorm:"index"`
+	ID          string `gorm:"primaryKey"`
 	NetworkID   string `gorm:"index"`
 	AccountID   string `gorm:"index"`
 	Name        string
diff --git a/management/server/networks/routers/types/router.go b/management/server/networks/routers/types/router.go
index 5158ebb12..71465868f 100644
--- a/management/server/networks/routers/types/router.go
+++ b/management/server/networks/routers/types/router.go
@@ -10,7 +10,7 @@ import (
 )
 
 type NetworkRouter struct {
-	ID         string `gorm:"index"`
+	ID         string `gorm:"primaryKey"`
 	NetworkID  string `gorm:"index"`
 	AccountID  string `gorm:"index"`
 	Peer       string
diff --git a/management/server/networks/types/network.go b/management/server/networks/types/network.go
index a4ba7b821..d1c7f2b33 100644
--- a/management/server/networks/types/network.go
+++ b/management/server/networks/types/network.go
@@ -7,7 +7,7 @@ import (
 )
 
 type Network struct {
-	ID          string `gorm:"index"`
+	ID          string `gorm:"primaryKey"`
 	AccountID   string `gorm:"index"`
 	Name        string
 	Description string
diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go
index dd39cf77d..d0adad6ee 100644
--- a/management/server/store/sql_store.go
+++ b/management/server/store/sql_store.go
@@ -82,6 +82,10 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
 			log.WithContext(ctx).Warnf("setting NB_SQL_MAX_OPEN_CONNS is not supported for sqlite, using default value 1")
 		}
 		conns = 1
+		_, err = sql.Exec("PRAGMA foreign_keys = ON")
+		if err != nil {
+			return nil, fmt.Errorf("failed to set foreign keys for sqlite: %w", err)
+		}
 	}
 
 	sql.SetMaxOpenConns(conns)
diff --git a/management/server/store/sql_store_test.go b/management/server/store/sql_store_test.go
index 8bd8ce098..8e99b34e1 100644
--- a/management/server/store/sql_store_test.go
+++ b/management/server/store/sql_store_test.go
@@ -60,10 +60,10 @@ func Test_NewStore(t *testing.T) {
 
 	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
 		if store == nil {
-			t.Errorf("expected to create a new Store")
+			t.Fatalf("expected to create a new Store")
 		}
 		if len(store.GetAllAccounts(context.Background())) != 0 {
-			t.Errorf("expected to create a new empty Accounts map when creating a new FileStore")
+			t.Fatalf("expected to create a new empty Accounts map when creating a new FileStore")
 		}
 	})
 }
@@ -1115,7 +1115,7 @@ func TestSqlite_CreateAndGetObjectInTransaction(t *testing.T) {
 
 	group := &types.Group{
 		ID:        "group-id",
-		AccountID: "account-id",
+		AccountID: "bf1c8084-ba50-4ce7-9439-34653001fc3b",
 		Name:      "group-name",
 		Issued:    "api",
 		Peers:     nil,
diff --git a/management/server/store/store.go b/management/server/store/store.go
index ca332a493..6da623956 100644
--- a/management/server/store/store.go
+++ b/management/server/store/store.go
@@ -315,6 +315,15 @@ func getMigrations(ctx context.Context) []migrationFunc {
 		func(db *gorm.DB) error {
 			return migration.MigrateNewField[routerTypes.NetworkRouter](ctx, db, "enabled", true)
 		},
+		func(db *gorm.DB) error {
+			return migration.DropIndex[networkTypes.Network](ctx, db, "idx_networks_id")
+		},
+		func(db *gorm.DB) error {
+			return migration.DropIndex[resourceTypes.NetworkResource](ctx, db, "idx_network_resources_id")
+		},
+		func(db *gorm.DB) error {
+			return migration.DropIndex[routerTypes.NetworkRouter](ctx, db, "idx_network_routers_id")
+		},
 	}
 }
 
diff --git a/management/server/types/group.go b/management/server/types/group.go
index 00a28fa77..1b321387c 100644
--- a/management/server/types/group.go
+++ b/management/server/types/group.go
@@ -14,7 +14,7 @@ const (
 // Group of the peers for ACL
 type Group struct {
 	// ID of the group
-	ID string
+	ID string `gorm:"primaryKey"`
 
 	// AccountID is a reference to Account that this object belongs
 	AccountID string `json:"-" gorm:"index"`

From ffdd115ded596f53085a37afbda9bcba7ea353a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alin=20Tr=C4=83istaru?= <github@cevreau.eu>
Date: Mon, 5 May 2025 12:20:54 +0200
Subject: [PATCH 05/11] [client] set TLS ServerName for hostname-based QUIC
 connections (#3673)

* fix: set TLS ServerName for hostname-based QUIC connections

When connecting to a relay server by hostname, certificates are
validated against the IP address instead of the hostname.
This change sets ServerName in the TLS config when connecting
via hostname, ensuring proper certificate validation.

* use default port if port is missing in URL string
---
 relay/client/dialer/quic/quic.go | 37 +++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/relay/client/dialer/quic/quic.go b/relay/client/dialer/quic/quic.go
index 7fd486f87..3fd48fb19 100644
--- a/relay/client/dialer/quic/quic.go
+++ b/relay/client/dialer/quic/quic.go
@@ -28,6 +28,16 @@ func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
 		return nil, err
 	}
 
+	// Get the base TLS config
+	tlsClientConfig := quictls.ClientQUICTLSConfig()
+
+	// Set ServerName to hostname if not an IP address
+	host, _, splitErr := net.SplitHostPort(quicURL)
+	if splitErr == nil && net.ParseIP(host) == nil {
+		// It's a hostname, not an IP - modify directly
+		tlsClientConfig.ServerName = host
+	}
+
 	quicConfig := &quic.Config{
 		KeepAlivePeriod:   30 * time.Second,
 		MaxIdleTimeout:    4 * time.Minute,
@@ -47,7 +57,7 @@ func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
 		return nil, err
 	}
 
-	session, err := quic.Dial(ctx, udpConn, udpAddr, quictls.ClientQUICTLSConfig(), quicConfig)
+	session, err := quic.Dial(ctx, udpConn, udpAddr, tlsClientConfig, quicConfig)
 	if err != nil {
 		if errors.Is(err, context.Canceled) {
 			return nil, err
@@ -61,12 +71,29 @@ func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
 }
 
 func prepareURL(address string) (string, error) {
-	if !strings.HasPrefix(address, "rel://") && !strings.HasPrefix(address, "rels://") {
+	var host string
+	var defaultPort string
+
+	switch {
+	case strings.HasPrefix(address, "rels://"):
+		host = address[7:]
+		defaultPort = "443"
+	case strings.HasPrefix(address, "rel://"):
+		host = address[6:]
+		defaultPort = "80"
+	default:
 		return "", fmt.Errorf("unsupported scheme: %s", address)
 	}
 
-	if strings.HasPrefix(address, "rels://") {
-		return address[7:], nil
+	finalHost, finalPort, err := net.SplitHostPort(host)
+	if err != nil {
+		if strings.Contains(err.Error(), "missing port") {
+			return host + ":" + defaultPort, nil
+		}
+
+		// return any other split error as is
+		return "", err
 	}
-	return address[6:], nil
+
+	return finalHost + ":" + finalPort, nil
 }

From 9762b39f29e63033bfbd8a5b68aa320db1ed4584 Mon Sep 17 00:00:00 2001
From: Viktor Liu <17948409+lixmal@users.noreply.github.com>
Date: Mon, 5 May 2025 14:29:05 +0200
Subject: [PATCH 06/11] [client] Fix stale local records (#3776)

---
 client/internal/dns/handler_chain_test.go |  40 +-
 client/internal/dns/local.go              | 130 ------
 client/internal/dns/local/local.go        | 149 +++++++
 client/internal/dns/local/local_test.go   | 472 ++++++++++++++++++++++
 client/internal/dns/local_test.go         |  88 ----
 client/internal/dns/mock_test.go          |  26 --
 client/internal/dns/server.go             |  85 ++--
 client/internal/dns/server_test.go        | 160 ++++----
 client/internal/dns/test/mock.go          |  26 ++
 client/internal/dns/types/types.go        |   3 +
 client/internal/dns/upstream.go           |  11 +-
 client/internal/dns/upstream_test.go      |   6 +-
 dns/dns.go                                |   6 +-
 13 files changed, 786 insertions(+), 416 deletions(-)
 delete mode 100644 client/internal/dns/local.go
 create mode 100644 client/internal/dns/local/local.go
 create mode 100644 client/internal/dns/local/local_test.go
 delete mode 100644 client/internal/dns/local_test.go
 delete mode 100644 client/internal/dns/mock_test.go
 create mode 100644 client/internal/dns/test/mock.go
 create mode 100644 client/internal/dns/types/types.go

diff --git a/client/internal/dns/handler_chain_test.go b/client/internal/dns/handler_chain_test.go
index 4c910a95f..5f03e0758 100644
--- a/client/internal/dns/handler_chain_test.go
+++ b/client/internal/dns/handler_chain_test.go
@@ -1,7 +1,6 @@
 package dns_test
 
 import (
-	"net"
 	"testing"
 
 	"github.com/miekg/dns"
@@ -9,6 +8,7 @@ import (
 	"github.com/stretchr/testify/mock"
 
 	nbdns "github.com/netbirdio/netbird/client/internal/dns"
+	"github.com/netbirdio/netbird/client/internal/dns/test"
 )
 
 // TestHandlerChain_ServeDNS_Priorities tests that handlers are executed in priority order
@@ -30,7 +30,7 @@ func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
 	r.SetQuestion("example.com.", dns.TypeA)
 
 	// Create test writer
-	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 	// Setup expectations - only highest priority handler should be called
 	dnsRouteHandler.On("ServeDNS", mock.Anything, r).Once()
@@ -142,7 +142,7 @@ func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
 
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryDomain, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 			chain.ServeDNS(w, r)
 
@@ -259,7 +259,7 @@ func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
 			// Create and execute request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryDomain, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 			chain.ServeDNS(w, r)
 
 			// Verify expectations
@@ -316,7 +316,7 @@ func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
 	}).Once()
 
 	// Execute
-	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 	chain.ServeDNS(w, r)
 
 	// Verify all handlers were called in order
@@ -325,20 +325,6 @@ func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
 	handler3.AssertExpectations(t)
 }
 
-// mockResponseWriter implements dns.ResponseWriter for testing
-type mockResponseWriter struct {
-	mock.Mock
-}
-
-func (m *mockResponseWriter) LocalAddr() net.Addr       { return nil }
-func (m *mockResponseWriter) RemoteAddr() net.Addr      { return nil }
-func (m *mockResponseWriter) WriteMsg(*dns.Msg) error   { return nil }
-func (m *mockResponseWriter) Write([]byte) (int, error) { return 0, nil }
-func (m *mockResponseWriter) Close() error              { return nil }
-func (m *mockResponseWriter) TsigStatus() error         { return nil }
-func (m *mockResponseWriter) TsigTimersOnly(bool)       {}
-func (m *mockResponseWriter) Hijack()                   {}
-
 func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 	tests := []struct {
 		name string
@@ -425,7 +411,7 @@ func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 			// Create test request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 			// Setup expectations
 			for priority, handler := range handlers {
@@ -471,7 +457,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain)
 
 	// Test 1: Initial state
-	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 	// Highest priority handler (routeHandler) should be called
 	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 	matchHandler.On("ServeDNS", mock.Anything, r).Maybe()   // Ensure others are not expected yet
@@ -490,7 +476,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 2: Remove highest priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)
 
-	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 	// Now middle priority handler (matchHandler) should be called
 	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure default is not expected yet
@@ -506,7 +492,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 3: Remove middle priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityMatchDomain)
 
-	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 	// Now lowest priority handler (defaultHandler) should be called
 	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 
@@ -519,7 +505,7 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	// Test 4: Remove last handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
 
-	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 	chain.ServeDNS(w4, r) // Call ServeDNS on the now empty chain for this domain
 
 	for _, m := range mocks {
@@ -675,7 +661,7 @@ func TestHandlerChain_CaseSensitivity(t *testing.T) {
 			// Execute request
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			chain.ServeDNS(&mockResponseWriter{}, r)
+			chain.ServeDNS(&test.MockResponseWriter{}, r)
 
 			// Verify each handler was called exactly as expected
 			for _, h := range tt.addHandlers {
@@ -819,7 +805,7 @@ func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) {
 
 			r := new(dns.Msg)
 			r.SetQuestion(tt.query, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 			// Setup handler expectations
 			for pattern, handler := range handlers {
@@ -969,7 +955,7 @@ func TestHandlerChain_AddRemoveRoundtrip(t *testing.T) {
 			handler := &nbdns.MockHandler{}
 			r := new(dns.Msg)
 			r.SetQuestion(tt.queryPattern, dns.TypeA)
-			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 			// First verify no handler is called before adding any
 			chain.ServeDNS(w, r)
diff --git a/client/internal/dns/local.go b/client/internal/dns/local.go
deleted file mode 100644
index 76e18e3ce..000000000
--- a/client/internal/dns/local.go
+++ /dev/null
@@ -1,130 +0,0 @@
-package dns
-
-import (
-	"fmt"
-	"strings"
-	"sync"
-
-	"github.com/miekg/dns"
-	log "github.com/sirupsen/logrus"
-
-	nbdns "github.com/netbirdio/netbird/dns"
-)
-
-type registrationMap map[string]struct{}
-
-type localResolver struct {
-	registeredMap registrationMap
-	records       sync.Map // key: string (domain_class_type), value: []dns.RR
-}
-
-func (d *localResolver) MatchSubdomains() bool {
-	return true
-}
-
-func (d *localResolver) stop() {
-}
-
-// String returns a string representation of the local resolver
-func (d *localResolver) String() string {
-	return fmt.Sprintf("local resolver [%d records]", len(d.registeredMap))
-}
-
-// ID returns the unique handler ID
-func (d *localResolver) id() handlerID {
-	return "local-resolver"
-}
-
-// ServeDNS handles a DNS request
-func (d *localResolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
-	if len(r.Question) > 0 {
-		log.Tracef("received local question: domain=%s type=%v class=%v", r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
-	}
-
-	replyMessage := &dns.Msg{}
-	replyMessage.SetReply(r)
-	replyMessage.RecursionAvailable = true
-
-	// lookup all records matching the question
-	records := d.lookupRecords(r)
-	if len(records) > 0 {
-		replyMessage.Rcode = dns.RcodeSuccess
-		replyMessage.Answer = append(replyMessage.Answer, records...)
-	} else {
-		replyMessage.Rcode = dns.RcodeNameError
-	}
-
-	err := w.WriteMsg(replyMessage)
-	if err != nil {
-		log.Debugf("got an error while writing the local resolver response, error: %v", err)
-	}
-}
-
-// lookupRecords fetches *all* DNS records matching the first question in r.
-func (d *localResolver) lookupRecords(r *dns.Msg) []dns.RR {
-	if len(r.Question) == 0 {
-		return nil
-	}
-	question := r.Question[0]
-	question.Name = strings.ToLower(question.Name)
-	key := buildRecordKey(question.Name, question.Qclass, question.Qtype)
-
-	value, found := d.records.Load(key)
-	if !found {
-		// alternatively check if we have a cname
-		if question.Qtype != dns.TypeCNAME {
-			r.Question[0].Qtype = dns.TypeCNAME
-			return d.lookupRecords(r)
-		}
-
-		return nil
-	}
-
-	records, ok := value.([]dns.RR)
-	if !ok {
-		log.Errorf("failed to cast records to []dns.RR, records: %v", value)
-		return nil
-	}
-
-	// if there's more than one record, rotate them (round-robin)
-	if len(records) > 1 {
-		first := records[0]
-		records = append(records[1:], first)
-		d.records.Store(key, records)
-	}
-
-	return records
-}
-
-// registerRecord stores a new record by appending it to any existing list
-func (d *localResolver) registerRecord(record nbdns.SimpleRecord) (string, error) {
-	rr, err := dns.NewRR(record.String())
-	if err != nil {
-		return "", fmt.Errorf("register record: %w", err)
-	}
-
-	rr.Header().Rdlength = record.Len()
-	header := rr.Header()
-	key := buildRecordKey(header.Name, header.Class, header.Rrtype)
-
-	// load any existing slice of records, then append
-	existing, _ := d.records.LoadOrStore(key, []dns.RR{})
-	records := existing.([]dns.RR)
-	records = append(records, rr)
-
-	// store updated slice
-	d.records.Store(key, records)
-	return key, nil
-}
-
-// deleteRecord removes *all* records under the recordKey.
-func (d *localResolver) deleteRecord(recordKey string) {
-	d.records.Delete(dns.Fqdn(recordKey))
-}
-
-// buildRecordKey consistently generates a key: name_class_type
-func buildRecordKey(name string, class, qType uint16) string {
-	return fmt.Sprintf("%s_%d_%d", dns.Fqdn(name), class, qType)
-}
-
-func (d *localResolver) probeAvailability() {}
diff --git a/client/internal/dns/local/local.go b/client/internal/dns/local/local.go
new file mode 100644
index 000000000..de3d8514b
--- /dev/null
+++ b/client/internal/dns/local/local.go
@@ -0,0 +1,149 @@
+package local
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"sync"
+
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
+
+	"github.com/netbirdio/netbird/client/internal/dns/types"
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+type Resolver struct {
+	mu      sync.RWMutex
+	records map[dns.Question][]dns.RR
+}
+
+func NewResolver() *Resolver {
+	return &Resolver{
+		records: make(map[dns.Question][]dns.RR),
+	}
+}
+
+func (d *Resolver) MatchSubdomains() bool {
+	return true
+}
+
+// String returns a string representation of the local resolver
+func (d *Resolver) String() string {
+	return fmt.Sprintf("local resolver [%d records]", len(d.records))
+}
+
+func (d *Resolver) Stop() {}
+
+// ID returns the unique handler ID
+func (d *Resolver) ID() types.HandlerID {
+	return "local-resolver"
+}
+
+func (d *Resolver) ProbeAvailability() {}
+
+// ServeDNS handles a DNS request
+func (d *Resolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	if len(r.Question) == 0 {
+		log.Debugf("received local resolver request with no question")
+		return
+	}
+	question := r.Question[0]
+	question.Name = strings.ToLower(dns.Fqdn(question.Name))
+
+	log.Tracef("received local question: domain=%s type=%v class=%v", r.Question[0].Name, question.Qtype, question.Qclass)
+
+	replyMessage := &dns.Msg{}
+	replyMessage.SetReply(r)
+	replyMessage.RecursionAvailable = true
+
+	// lookup all records matching the question
+	records := d.lookupRecords(question)
+	if len(records) > 0 {
+		replyMessage.Rcode = dns.RcodeSuccess
+		replyMessage.Answer = append(replyMessage.Answer, records...)
+	} else {
+		// TODO: return success if we have a different record type for the same name, relevant for search domains
+		replyMessage.Rcode = dns.RcodeNameError
+	}
+
+	if err := w.WriteMsg(replyMessage); err != nil {
+		log.Warnf("failed to write the local resolver response: %v", err)
+	}
+}
+
+// lookupRecords fetches *all* DNS records matching the first question in r.
+func (d *Resolver) lookupRecords(question dns.Question) []dns.RR {
+	d.mu.RLock()
+	records, found := d.records[question]
+
+	if !found {
+		d.mu.RUnlock()
+		// alternatively check if we have a cname
+		if question.Qtype != dns.TypeCNAME {
+			question.Qtype = dns.TypeCNAME
+			return d.lookupRecords(question)
+		}
+		return nil
+	}
+
+	recordsCopy := slices.Clone(records)
+	d.mu.RUnlock()
+
+	// if there's more than one record, rotate them (round-robin)
+	if len(recordsCopy) > 1 {
+		d.mu.Lock()
+		records = d.records[question]
+		if len(records) > 1 {
+			first := records[0]
+			records = append(records[1:], first)
+			d.records[question] = records
+		}
+		d.mu.Unlock()
+	}
+
+	return recordsCopy
+}
+
+func (d *Resolver) Update(update []nbdns.SimpleRecord) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	maps.Clear(d.records)
+
+	for _, rec := range update {
+		if err := d.registerRecord(rec); err != nil {
+			log.Warnf("failed to register the record (%s): %v", rec, err)
+			continue
+		}
+	}
+}
+
+// RegisterRecord stores a new record by appending it to any existing list
+func (d *Resolver) RegisterRecord(record nbdns.SimpleRecord) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	return d.registerRecord(record)
+}
+
+// registerRecord performs the registration with the lock already held
+func (d *Resolver) registerRecord(record nbdns.SimpleRecord) error {
+	rr, err := dns.NewRR(record.String())
+	if err != nil {
+		return fmt.Errorf("register record: %w", err)
+	}
+
+	rr.Header().Rdlength = record.Len()
+	header := rr.Header()
+	q := dns.Question{
+		Name:   strings.ToLower(dns.Fqdn(header.Name)),
+		Qtype:  header.Rrtype,
+		Qclass: header.Class,
+	}
+
+	d.records[q] = append(d.records[q], rr)
+
+	return nil
+}
diff --git a/client/internal/dns/local/local_test.go b/client/internal/dns/local/local_test.go
new file mode 100644
index 000000000..1d38191e7
--- /dev/null
+++ b/client/internal/dns/local/local_test.go
@@ -0,0 +1,472 @@
+package local
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/internal/dns/test"
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+func TestLocalResolver_ServeDNS(t *testing.T) {
+	recordA := nbdns.SimpleRecord{
+		Name:  "peera.netbird.cloud.",
+		Type:  1,
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "1.2.3.4",
+	}
+
+	recordCNAME := nbdns.SimpleRecord{
+		Name:  "peerb.netbird.cloud.",
+		Type:  5,
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "www.netbird.io",
+	}
+
+	testCases := []struct {
+		name                string
+		inputRecord         nbdns.SimpleRecord
+		inputMSG            *dns.Msg
+		responseShouldBeNil bool
+	}{
+		{
+			name:        "Should Resolve A Record",
+			inputRecord: recordA,
+			inputMSG:    new(dns.Msg).SetQuestion(recordA.Name, dns.TypeA),
+		},
+		{
+			name:        "Should Resolve CNAME Record",
+			inputRecord: recordCNAME,
+			inputMSG:    new(dns.Msg).SetQuestion(recordCNAME.Name, dns.TypeCNAME),
+		},
+		{
+			name:                "Should Not Write When Not Found A Record",
+			inputRecord:         recordA,
+			inputMSG:            new(dns.Msg).SetQuestion("not.found.com", dns.TypeA),
+			responseShouldBeNil: true,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			resolver := NewResolver()
+			_ = resolver.RegisterRecord(testCase.inputRecord)
+			var responseMSG *dns.Msg
+			responseWriter := &test.MockResponseWriter{
+				WriteMsgFunc: func(m *dns.Msg) error {
+					responseMSG = m
+					return nil
+				},
+			}
+
+			resolver.ServeDNS(responseWriter, testCase.inputMSG)
+
+			if responseMSG == nil || len(responseMSG.Answer) == 0 {
+				if testCase.responseShouldBeNil {
+					return
+				}
+				t.Fatalf("should write a response message")
+			}
+
+			answerString := responseMSG.Answer[0].String()
+			if !strings.Contains(answerString, testCase.inputRecord.Name) {
+				t.Fatalf("answer doesn't contain the same domain name: \nWant: %s\nGot:%s", testCase.name, answerString)
+			}
+			if !strings.Contains(answerString, dns.Type(testCase.inputRecord.Type).String()) {
+				t.Fatalf("answer doesn't contain the correct type: \nWant: %s\nGot:%s", dns.Type(testCase.inputRecord.Type).String(), answerString)
+			}
+			if !strings.Contains(answerString, testCase.inputRecord.RData) {
+				t.Fatalf("answer doesn't contain the same address: \nWant: %s\nGot:%s", testCase.inputRecord.RData, answerString)
+			}
+		})
+	}
+}
+
+// TestLocalResolver_Update_StaleRecord verifies that updating
+// a record correctly replaces the old one, preventing stale entries.
+func TestLocalResolver_Update_StaleRecord(t *testing.T) {
+	recordName := "host.example.com."
+	recordType := dns.TypeA
+	recordClass := dns.ClassINET
+
+	record1 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "1.1.1.1",
+	}
+	record2 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "2.2.2.2",
+	}
+
+	recordKey := dns.Question{Name: recordName, Qtype: uint16(recordClass), Qclass: recordType}
+
+	resolver := NewResolver()
+
+	update1 := []nbdns.SimpleRecord{record1}
+	update2 := []nbdns.SimpleRecord{record2}
+
+	// Apply first update
+	resolver.Update(update1)
+
+	// Verify first update
+	resolver.mu.RLock()
+	rrSlice1, found1 := resolver.records[recordKey]
+	resolver.mu.RUnlock()
+
+	require.True(t, found1, "Record key %s not found after first update", recordKey)
+	require.Len(t, rrSlice1, 1, "Should have exactly 1 record after first update")
+	assert.Contains(t, rrSlice1[0].String(), record1.RData, "Record after first update should be %s", record1.RData)
+
+	// Apply second update
+	resolver.Update(update2)
+
+	// Verify second update
+	resolver.mu.RLock()
+	rrSlice2, found2 := resolver.records[recordKey]
+	resolver.mu.RUnlock()
+
+	require.True(t, found2, "Record key %s not found after second update", recordKey)
+	require.Len(t, rrSlice2, 1, "Should have exactly 1 record after update overwriting the key")
+	assert.Contains(t, rrSlice2[0].String(), record2.RData, "The single record should be the updated one (%s)", record2.RData)
+	assert.NotContains(t, rrSlice2[0].String(), record1.RData, "The stale record (%s) should not be present", record1.RData)
+}
+
+// TestLocalResolver_MultipleRecords_SameQuestion verifies that multiple records
+// with the same question are stored properly
+func TestLocalResolver_MultipleRecords_SameQuestion(t *testing.T) {
+	resolver := NewResolver()
+
+	recordName := "multi.example.com."
+	recordType := dns.TypeA
+
+	// Create two records with the same name and type but different IPs
+	record1 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1",
+	}
+	record2 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2",
+	}
+
+	update := []nbdns.SimpleRecord{record1, record2}
+
+	// Apply update with both records
+	resolver.Update(update)
+
+	// Create question that matches both records
+	question := dns.Question{
+		Name:   recordName,
+		Qtype:  recordType,
+		Qclass: dns.ClassINET,
+	}
+
+	// Verify both records are stored
+	resolver.mu.RLock()
+	records, found := resolver.records[question]
+	resolver.mu.RUnlock()
+
+	require.True(t, found, "Records for question %v not found", question)
+	require.Len(t, records, 2, "Should have exactly 2 records for the same question")
+
+	// Verify both record data values are present
+	recordStrings := []string{records[0].String(), records[1].String()}
+	assert.Contains(t, recordStrings[0]+recordStrings[1], record1.RData, "First record data should be present")
+	assert.Contains(t, recordStrings[0]+recordStrings[1], record2.RData, "Second record data should be present")
+}
+
+// TestLocalResolver_RecordRotation verifies that records are rotated in a round-robin fashion
+func TestLocalResolver_RecordRotation(t *testing.T) {
+	resolver := NewResolver()
+
+	recordName := "rotation.example.com."
+	recordType := dns.TypeA
+
+	// Create three records with the same name and type but different IPs
+	record1 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.1",
+	}
+	record2 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.2",
+	}
+	record3 := nbdns.SimpleRecord{
+		Name: recordName, Type: int(recordType), Class: nbdns.DefaultClass, TTL: 300, RData: "192.168.1.3",
+	}
+
+	update := []nbdns.SimpleRecord{record1, record2, record3}
+
+	// Apply update with all three records
+	resolver.Update(update)
+
+	msg := new(dns.Msg).SetQuestion(recordName, recordType)
+
+	// First lookup - should return the records in original order
+	var responses [3]*dns.Msg
+
+	// Perform three lookups to verify rotation
+	for i := 0; i < 3; i++ {
+		responseWriter := &test.MockResponseWriter{
+			WriteMsgFunc: func(m *dns.Msg) error {
+				responses[i] = m
+				return nil
+			},
+		}
+
+		resolver.ServeDNS(responseWriter, msg)
+	}
+
+	// Verify all three responses contain answers
+	for i, resp := range responses {
+		require.NotNil(t, resp, "Response %d should not be nil", i)
+		require.Len(t, resp.Answer, 3, "Response %d should have 3 answers", i)
+	}
+
+	// Verify the first record in each response is different due to rotation
+	firstRecordIPs := []string{
+		responses[0].Answer[0].String(),
+		responses[1].Answer[0].String(),
+		responses[2].Answer[0].String(),
+	}
+
+	// Each record should be different (rotated)
+	assert.NotEqual(t, firstRecordIPs[0], firstRecordIPs[1], "First lookup should differ from second lookup due to rotation")
+	assert.NotEqual(t, firstRecordIPs[1], firstRecordIPs[2], "Second lookup should differ from third lookup due to rotation")
+	assert.NotEqual(t, firstRecordIPs[0], firstRecordIPs[2], "First lookup should differ from third lookup due to rotation")
+
+	// After three rotations, we should have cycled through all records
+	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record1.RData)
+	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record2.RData)
+	assert.Contains(t, firstRecordIPs[0]+firstRecordIPs[1]+firstRecordIPs[2], record3.RData)
+}
+
+// TestLocalResolver_CaseInsensitiveMatching verifies that DNS record lookups are case-insensitive
+func TestLocalResolver_CaseInsensitiveMatching(t *testing.T) {
+	resolver := NewResolver()
+
+	// Create record with lowercase name
+	lowerCaseRecord := nbdns.SimpleRecord{
+		Name:  "lower.example.com.",
+		Type:  int(dns.TypeA),
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "10.10.10.10",
+	}
+
+	// Create record with mixed case name
+	mixedCaseRecord := nbdns.SimpleRecord{
+		Name:  "MiXeD.ExAmPlE.CoM.",
+		Type:  int(dns.TypeA),
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "20.20.20.20",
+	}
+
+	// Update resolver with the records
+	resolver.Update([]nbdns.SimpleRecord{lowerCaseRecord, mixedCaseRecord})
+
+	testCases := []struct {
+		name          string
+		queryName     string
+		expectedRData string
+		shouldResolve bool
+	}{
+		{
+			name:          "Query lowercase with lowercase record",
+			queryName:     "lower.example.com.",
+			expectedRData: "10.10.10.10",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query uppercase with lowercase record",
+			queryName:     "LOWER.EXAMPLE.COM.",
+			expectedRData: "10.10.10.10",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query mixed case with lowercase record",
+			queryName:     "LoWeR.eXaMpLe.CoM.",
+			expectedRData: "10.10.10.10",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query lowercase with mixed case record",
+			queryName:     "mixed.example.com.",
+			expectedRData: "20.20.20.20",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query uppercase with mixed case record",
+			queryName:     "MIXED.EXAMPLE.COM.",
+			expectedRData: "20.20.20.20",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query with different casing pattern",
+			queryName:     "mIxEd.ExaMpLe.cOm.",
+			expectedRData: "20.20.20.20",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query non-existent domain",
+			queryName:     "nonexistent.example.com.",
+			shouldResolve: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var responseMSG *dns.Msg
+
+			// Create DNS query with the test case name
+			msg := new(dns.Msg).SetQuestion(tc.queryName, dns.TypeA)
+
+			// Create mock response writer to capture the response
+			responseWriter := &test.MockResponseWriter{
+				WriteMsgFunc: func(m *dns.Msg) error {
+					responseMSG = m
+					return nil
+				},
+			}
+
+			// Perform DNS query
+			resolver.ServeDNS(responseWriter, msg)
+
+			// Check if we expect a successful resolution
+			if !tc.shouldResolve {
+				if responseMSG == nil || len(responseMSG.Answer) == 0 {
+					// Expected no answer, test passes
+					return
+				}
+				t.Fatalf("Expected no resolution for %s, but got answer: %v", tc.queryName, responseMSG.Answer)
+			}
+
+			// Verify we got a response
+			require.NotNil(t, responseMSG, "Should have received a response message")
+			require.Greater(t, len(responseMSG.Answer), 0, "Response should contain at least one answer")
+
+			// Verify the response contains the expected data
+			answerString := responseMSG.Answer[0].String()
+			assert.Contains(t, answerString, tc.expectedRData,
+				"Answer should contain the expected IP address %s, got: %s",
+				tc.expectedRData, answerString)
+		})
+	}
+}
+
+// TestLocalResolver_CNAMEFallback verifies that the resolver correctly falls back
+// to checking for CNAME records when the requested record type isn't found
+func TestLocalResolver_CNAMEFallback(t *testing.T) {
+	resolver := NewResolver()
+
+	// Create a CNAME record (but no A record for this name)
+	cnameRecord := nbdns.SimpleRecord{
+		Name:  "alias.example.com.",
+		Type:  int(dns.TypeCNAME),
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "target.example.com.",
+	}
+
+	// Create an A record for the CNAME target
+	targetRecord := nbdns.SimpleRecord{
+		Name:  "target.example.com.",
+		Type:  int(dns.TypeA),
+		Class: nbdns.DefaultClass,
+		TTL:   300,
+		RData: "192.168.100.100",
+	}
+
+	// Update resolver with both records
+	resolver.Update([]nbdns.SimpleRecord{cnameRecord, targetRecord})
+
+	testCases := []struct {
+		name          string
+		queryName     string
+		queryType     uint16
+		expectedType  string
+		expectedRData string
+		shouldResolve bool
+	}{
+		{
+			name:          "Directly query CNAME record",
+			queryName:     "alias.example.com.",
+			queryType:     dns.TypeCNAME,
+			expectedType:  "CNAME",
+			expectedRData: "target.example.com.",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query A record but get CNAME fallback",
+			queryName:     "alias.example.com.",
+			queryType:     dns.TypeA,
+			expectedType:  "CNAME",
+			expectedRData: "target.example.com.",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query AAAA record but get CNAME fallback",
+			queryName:     "alias.example.com.",
+			queryType:     dns.TypeAAAA,
+			expectedType:  "CNAME",
+			expectedRData: "target.example.com.",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query direct A record",
+			queryName:     "target.example.com.",
+			queryType:     dns.TypeA,
+			expectedType:  "A",
+			expectedRData: "192.168.100.100",
+			shouldResolve: true,
+		},
+		{
+			name:          "Query non-existent name",
+			queryName:     "nonexistent.example.com.",
+			queryType:     dns.TypeA,
+			shouldResolve: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var responseMSG *dns.Msg
+
+			// Create DNS query with the test case parameters
+			msg := new(dns.Msg).SetQuestion(tc.queryName, tc.queryType)
+
+			// Create mock response writer to capture the response
+			responseWriter := &test.MockResponseWriter{
+				WriteMsgFunc: func(m *dns.Msg) error {
+					responseMSG = m
+					return nil
+				},
+			}
+
+			// Perform DNS query
+			resolver.ServeDNS(responseWriter, msg)
+
+			// Check if we expect a successful resolution
+			if !tc.shouldResolve {
+				if responseMSG == nil || len(responseMSG.Answer) == 0 || responseMSG.Rcode != dns.RcodeSuccess {
+					// Expected no resolution, test passes
+					return
+				}
+				t.Fatalf("Expected no resolution for %s, but got answer: %v", tc.queryName, responseMSG.Answer)
+			}
+
+			// Verify we got a successful response
+			require.NotNil(t, responseMSG, "Should have received a response message")
+			require.Equal(t, dns.RcodeSuccess, responseMSG.Rcode, "Response should have success status code")
+			require.Greater(t, len(responseMSG.Answer), 0, "Response should contain at least one answer")
+
+			// Verify the response contains the expected data
+			answerString := responseMSG.Answer[0].String()
+			assert.Contains(t, answerString, tc.expectedType,
+				"Answer should be of type %s, got: %s", tc.expectedType, answerString)
+			assert.Contains(t, answerString, tc.expectedRData,
+				"Answer should contain the expected data %s, got: %s", tc.expectedRData, answerString)
+		})
+	}
+}
diff --git a/client/internal/dns/local_test.go b/client/internal/dns/local_test.go
deleted file mode 100644
index 0a42b321a..000000000
--- a/client/internal/dns/local_test.go
+++ /dev/null
@@ -1,88 +0,0 @@
-package dns
-
-import (
-	"strings"
-	"testing"
-
-	"github.com/miekg/dns"
-
-	nbdns "github.com/netbirdio/netbird/dns"
-)
-
-func TestLocalResolver_ServeDNS(t *testing.T) {
-	recordA := nbdns.SimpleRecord{
-		Name:  "peera.netbird.cloud.",
-		Type:  1,
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "1.2.3.4",
-	}
-
-	recordCNAME := nbdns.SimpleRecord{
-		Name:  "peerb.netbird.cloud.",
-		Type:  5,
-		Class: nbdns.DefaultClass,
-		TTL:   300,
-		RData: "www.netbird.io",
-	}
-
-	testCases := []struct {
-		name                string
-		inputRecord         nbdns.SimpleRecord
-		inputMSG            *dns.Msg
-		responseShouldBeNil bool
-	}{
-		{
-			name:        "Should Resolve A Record",
-			inputRecord: recordA,
-			inputMSG:    new(dns.Msg).SetQuestion(recordA.Name, dns.TypeA),
-		},
-		{
-			name:        "Should Resolve CNAME Record",
-			inputRecord: recordCNAME,
-			inputMSG:    new(dns.Msg).SetQuestion(recordCNAME.Name, dns.TypeCNAME),
-		},
-		{
-			name:                "Should Not Write When Not Found A Record",
-			inputRecord:         recordA,
-			inputMSG:            new(dns.Msg).SetQuestion("not.found.com", dns.TypeA),
-			responseShouldBeNil: true,
-		},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			resolver := &localResolver{
-				registeredMap: make(registrationMap),
-			}
-			_, _ = resolver.registerRecord(testCase.inputRecord)
-			var responseMSG *dns.Msg
-			responseWriter := &mockResponseWriter{
-				WriteMsgFunc: func(m *dns.Msg) error {
-					responseMSG = m
-					return nil
-				},
-			}
-
-			resolver.ServeDNS(responseWriter, testCase.inputMSG)
-
-			if responseMSG == nil || len(responseMSG.Answer) == 0 {
-				if testCase.responseShouldBeNil {
-					return
-				}
-				t.Fatalf("should write a response message")
-			}
-
-			answerString := responseMSG.Answer[0].String()
-			if !strings.Contains(answerString, testCase.inputRecord.Name) {
-				t.Fatalf("answer doesn't contain the same domain name: \nWant: %s\nGot:%s", testCase.name, answerString)
-			}
-			if !strings.Contains(answerString, dns.Type(testCase.inputRecord.Type).String()) {
-				t.Fatalf("answer doesn't contain the correct type: \nWant: %s\nGot:%s", dns.Type(testCase.inputRecord.Type).String(), answerString)
-			}
-			if !strings.Contains(answerString, testCase.inputRecord.RData) {
-				t.Fatalf("answer doesn't contain the same address: \nWant: %s\nGot:%s", testCase.inputRecord.RData, answerString)
-			}
-		})
-	}
-}
diff --git a/client/internal/dns/mock_test.go b/client/internal/dns/mock_test.go
deleted file mode 100644
index d52ae24da..000000000
--- a/client/internal/dns/mock_test.go
+++ /dev/null
@@ -1,26 +0,0 @@
-package dns
-
-import (
-	"net"
-
-	"github.com/miekg/dns"
-)
-
-type mockResponseWriter struct {
-	WriteMsgFunc func(m *dns.Msg) error
-}
-
-func (rw *mockResponseWriter) WriteMsg(m *dns.Msg) error {
-	if rw.WriteMsgFunc != nil {
-		return rw.WriteMsgFunc(m)
-	}
-	return nil
-}
-
-func (rw *mockResponseWriter) LocalAddr() net.Addr       { return nil }
-func (rw *mockResponseWriter) RemoteAddr() net.Addr      { return nil }
-func (rw *mockResponseWriter) Write([]byte) (int, error) { return 0, nil }
-func (rw *mockResponseWriter) Close() error              { return nil }
-func (rw *mockResponseWriter) TsigStatus() error         { return nil }
-func (rw *mockResponseWriter) TsigTimersOnly(bool)       {}
-func (rw *mockResponseWriter) Hijack()                   {}
diff --git a/client/internal/dns/server.go b/client/internal/dns/server.go
index 65b90e5f0..3f49c23fd 100644
--- a/client/internal/dns/server.go
+++ b/client/internal/dns/server.go
@@ -15,6 +15,8 @@ import (
 	"golang.org/x/exp/maps"
 
 	"github.com/netbirdio/netbird/client/iface/netstack"
+	"github.com/netbirdio/netbird/client/internal/dns/local"
+	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/listener"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -46,8 +48,6 @@ type Server interface {
 	ProbeAvailability()
 }
 
-type handlerID string
-
 type nsGroupsByDomain struct {
 	domain string
 	groups []*nbdns.NameServerGroup
@@ -61,7 +61,7 @@ type DefaultServer struct {
 	mux                sync.Mutex
 	service            service
 	dnsMuxMap          registeredHandlerMap
-	localResolver      *localResolver
+	localResolver      *local.Resolver
 	wgInterface        WGIface
 	hostManager        hostManager
 	updateSerial       uint64
@@ -84,9 +84,9 @@ type DefaultServer struct {
 
 type handlerWithStop interface {
 	dns.Handler
-	stop()
-	probeAvailability()
-	id() handlerID
+	Stop()
+	ProbeAvailability()
+	ID() types.HandlerID
 }
 
 type handlerWrapper struct {
@@ -95,7 +95,7 @@ type handlerWrapper struct {
 	priority int
 }
 
-type registeredHandlerMap map[handlerID]handlerWrapper
+type registeredHandlerMap map[types.HandlerID]handlerWrapper
 
 // NewDefaultServer returns a new dns server
 func NewDefaultServer(
@@ -171,16 +171,14 @@ func newDefaultServer(
 	handlerChain := NewHandlerChain()
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
-		ctx:          ctx,
-		ctxCancel:    stop,
-		disableSys:   disableSys,
-		service:      dnsService,
-		handlerChain: handlerChain,
-		extraDomains: make(map[domain.Domain]int),
-		dnsMuxMap:    make(registeredHandlerMap),
-		localResolver: &localResolver{
-			registeredMap: make(registrationMap),
-		},
+		ctx:            ctx,
+		ctxCancel:      stop,
+		disableSys:     disableSys,
+		service:        dnsService,
+		handlerChain:   handlerChain,
+		extraDomains:   make(map[domain.Domain]int),
+		dnsMuxMap:      make(registeredHandlerMap),
+		localResolver:  local.NewResolver(),
 		wgInterface:    wgInterface,
 		statusRecorder: statusRecorder,
 		stateManager:   stateManager,
@@ -403,7 +401,7 @@ func (s *DefaultServer) ProbeAvailability() {
 		wg.Add(1)
 		go func(mux handlerWithStop) {
 			defer wg.Done()
-			mux.probeAvailability()
+			mux.ProbeAvailability()
 		}(mux.handler)
 	}
 	wg.Wait()
@@ -420,7 +418,7 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 		s.service.Stop()
 	}
 
-	localMuxUpdates, localRecordsByDomain, err := s.buildLocalHandlerUpdate(update.CustomZones)
+	localMuxUpdates, localRecords, err := s.buildLocalHandlerUpdate(update.CustomZones)
 	if err != nil {
 		return fmt.Errorf("local handler updater: %w", err)
 	}
@@ -434,7 +432,7 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	s.updateMux(muxUpdates)
 
 	// register local records
-	s.updateLocalResolver(localRecordsByDomain)
+	s.localResolver.Update(localRecords)
 
 	s.currentConfig = dnsConfigToHostDNSConfig(update, s.service.RuntimeIP(), s.service.RuntimePort())
 
@@ -516,11 +514,9 @@ func (s *DefaultServer) handleErrNoGroupaAll(err error) {
 	)
 }
 
-func (s *DefaultServer) buildLocalHandlerUpdate(
-	customZones []nbdns.CustomZone,
-) ([]handlerWrapper, map[string][]nbdns.SimpleRecord, error) {
+func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone) ([]handlerWrapper, []nbdns.SimpleRecord, error) {
 	var muxUpdates []handlerWrapper
-	localRecords := make(map[string][]nbdns.SimpleRecord)
+	var localRecords []nbdns.SimpleRecord
 
 	for _, customZone := range customZones {
 		if len(customZone.Records) == 0 {
@@ -534,17 +530,13 @@ func (s *DefaultServer) buildLocalHandlerUpdate(
 			priority: PriorityMatchDomain,
 		})
 
-		// group all records under this domain
 		for _, record := range customZone.Records {
-			var class uint16 = dns.ClassINET
 			if record.Class != nbdns.DefaultClass {
 				log.Warnf("received an invalid class type: %s", record.Class)
 				continue
 			}
-
-			key := buildRecordKey(record.Name, class, uint16(record.Type))
-
-			localRecords[key] = append(localRecords[key], record)
+			// zone records contain the fqdn, so we can just flatten them
+			localRecords = append(localRecords, record)
 		}
 	}
 
@@ -627,7 +619,7 @@ func (s *DefaultServer) createHandlersForDomainGroup(domainGroup nsGroupsByDomai
 		}
 
 		if len(handler.upstreamServers) == 0 {
-			handler.stop()
+			handler.Stop()
 			log.Errorf("received a nameserver group with an invalid nameserver list")
 			continue
 		}
@@ -656,7 +648,7 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 	// this will introduce a short period of time when the server is not able to handle DNS requests
 	for _, existing := range s.dnsMuxMap {
 		s.deregisterHandler([]string{existing.domain}, existing.priority)
-		existing.handler.stop()
+		existing.handler.Stop()
 	}
 
 	muxUpdateMap := make(registeredHandlerMap)
@@ -667,7 +659,7 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 			containsRootUpdate = true
 		}
 		s.registerHandler([]string{update.domain}, update.handler, update.priority)
-		muxUpdateMap[update.handler.id()] = update
+		muxUpdateMap[update.handler.ID()] = update
 	}
 
 	// If there's no root update and we had a root handler, restore it
@@ -683,33 +675,6 @@ func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) {
 	s.dnsMuxMap = muxUpdateMap
 }
 
-func (s *DefaultServer) updateLocalResolver(update map[string][]nbdns.SimpleRecord) {
-	// remove old records that are no longer present
-	for key := range s.localResolver.registeredMap {
-		_, found := update[key]
-		if !found {
-			s.localResolver.deleteRecord(key)
-		}
-	}
-
-	updatedMap := make(registrationMap)
-	for _, recs := range update {
-		for _, rec := range recs {
-			// convert the record to a dns.RR and register
-			key, err := s.localResolver.registerRecord(rec)
-			if err != nil {
-				log.Warnf("got an error while registering the record (%s), error: %v",
-					rec.String(), err)
-				continue
-			}
-
-			updatedMap[key] = struct{}{}
-		}
-	}
-
-	s.localResolver.registeredMap = updatedMap
-}
-
 func getNSHostPort(ns nbdns.NameServer) string {
 	return fmt.Sprintf("%s:%d", ns.IP.String(), ns.Port)
 }
diff --git a/client/internal/dns/server_test.go b/client/internal/dns/server_test.go
index ed69b0e93..1c7c9b117 100644
--- a/client/internal/dns/server_test.go
+++ b/client/internal/dns/server_test.go
@@ -23,6 +23,9 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	pfmock "github.com/netbirdio/netbird/client/iface/mocks"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	"github.com/netbirdio/netbird/client/internal/dns/local"
+	"github.com/netbirdio/netbird/client/internal/dns/test"
+	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/netflow"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -107,6 +110,7 @@ func generateDummyHandler(domain string, servers []nbdns.NameServer) *upstreamRe
 }
 
 func TestUpdateDNSServer(t *testing.T) {
+
 	nameServers := []nbdns.NameServer{
 		{
 			IP:     netip.MustParseAddr("8.8.8.8"),
@@ -120,22 +124,21 @@ func TestUpdateDNSServer(t *testing.T) {
 		},
 	}
 
-	dummyHandler := &localResolver{}
+	dummyHandler := local.NewResolver()
 
 	testCases := []struct {
 		name                string
 		initUpstreamMap     registeredHandlerMap
-		initLocalMap        registrationMap
+		initLocalRecords    []nbdns.SimpleRecord
 		initSerial          uint64
 		inputSerial         uint64
 		inputUpdate         nbdns.Config
 		shouldFail          bool
 		expectedUpstreamMap registeredHandlerMap
-		expectedLocalMap    registrationMap
+		expectedLocalQs     []dns.Question
 	}{
 		{
 			name:            "Initial Config Should Succeed",
-			initLocalMap:    make(registrationMap),
 			initUpstreamMap: make(registeredHandlerMap),
 			initSerial:      0,
 			inputSerial:     1,
@@ -159,30 +162,30 @@ func TestUpdateDNSServer(t *testing.T) {
 				},
 			},
 			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
+				generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{
 					domain:   "netbird.io",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
-				dummyHandler.id(): handlerWrapper{
+				dummyHandler.ID(): handlerWrapper{
 					domain:   "netbird.cloud",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
-				generateDummyHandler(".", nameServers).id(): handlerWrapper{
+				generateDummyHandler(".", nameServers).ID(): handlerWrapper{
 					domain:   nbdns.RootZone,
 					handler:  dummyHandler,
 					priority: PriorityDefault,
 				},
 			},
-			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
+			expectedLocalQs: []dns.Question{{Name: "peera.netbird.cloud.", Qtype: dns.TypeA, Qclass: dns.ClassINET}},
 		},
 		{
-			name:         "New Config Should Succeed",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
+			name:             "New Config Should Succeed",
+			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
-					domain:   buildRecordKey(zoneRecords[0].Name, 1, 1),
+				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
+					domain:   "netbird.cloud",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
 				},
@@ -205,7 +208,7 @@ func TestUpdateDNSServer(t *testing.T) {
 				},
 			},
 			expectedUpstreamMap: registeredHandlerMap{
-				generateDummyHandler("netbird.io", nameServers).id(): handlerWrapper{
+				generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{
 					domain:   "netbird.io",
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -216,22 +219,22 @@ func TestUpdateDNSServer(t *testing.T) {
 					priority: PriorityMatchDomain,
 				},
 			},
-			expectedLocalMap: registrationMap{buildRecordKey(zoneRecords[0].Name, 1, 1): struct{}{}},
+			expectedLocalQs: []dns.Question{{Name: zoneRecords[0].Name, Qtype: 1, Qclass: 1}},
 		},
 		{
-			name:            "Smaller Config Serial Should Be Skipped",
-			initLocalMap:    make(registrationMap),
-			initUpstreamMap: make(registeredHandlerMap),
-			initSerial:      2,
-			inputSerial:     1,
-			shouldFail:      true,
+			name:             "Smaller Config Serial Should Be Skipped",
+			initLocalRecords: []nbdns.SimpleRecord{},
+			initUpstreamMap:  make(registeredHandlerMap),
+			initSerial:       2,
+			inputSerial:      1,
+			shouldFail:       true,
 		},
 		{
-			name:            "Empty NS Group Domain Or Not Primary Element Should Fail",
-			initLocalMap:    make(registrationMap),
-			initUpstreamMap: make(registeredHandlerMap),
-			initSerial:      0,
-			inputSerial:     1,
+			name:             "Empty NS Group Domain Or Not Primary Element Should Fail",
+			initLocalRecords: []nbdns.SimpleRecord{},
+			initUpstreamMap:  make(registeredHandlerMap),
+			initSerial:       0,
+			inputSerial:      1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -249,11 +252,11 @@ func TestUpdateDNSServer(t *testing.T) {
 			shouldFail: true,
 		},
 		{
-			name:            "Invalid NS Group Nameservers list Should Fail",
-			initLocalMap:    make(registrationMap),
-			initUpstreamMap: make(registeredHandlerMap),
-			initSerial:      0,
-			inputSerial:     1,
+			name:             "Invalid NS Group Nameservers list Should Fail",
+			initLocalRecords: []nbdns.SimpleRecord{},
+			initUpstreamMap:  make(registeredHandlerMap),
+			initSerial:       0,
+			inputSerial:      1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -271,11 +274,11 @@ func TestUpdateDNSServer(t *testing.T) {
 			shouldFail: true,
 		},
 		{
-			name:            "Invalid Custom Zone Records list Should Skip",
-			initLocalMap:    make(registrationMap),
-			initUpstreamMap: make(registeredHandlerMap),
-			initSerial:      0,
-			inputSerial:     1,
+			name:             "Invalid Custom Zone Records list Should Skip",
+			initLocalRecords: []nbdns.SimpleRecord{},
+			initUpstreamMap:  make(registeredHandlerMap),
+			initSerial:       0,
+			inputSerial:      1,
 			inputUpdate: nbdns.Config{
 				ServiceEnable: true,
 				CustomZones: []nbdns.CustomZone{
@@ -290,17 +293,17 @@ func TestUpdateDNSServer(t *testing.T) {
 					},
 				},
 			},
-			expectedUpstreamMap: registeredHandlerMap{generateDummyHandler(".", nameServers).id(): handlerWrapper{
+			expectedUpstreamMap: registeredHandlerMap{generateDummyHandler(".", nameServers).ID(): handlerWrapper{
 				domain:   ".",
 				handler:  dummyHandler,
 				priority: PriorityDefault,
 			}},
 		},
 		{
-			name:         "Empty Config Should Succeed and Clean Maps",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
+			name:             "Empty Config Should Succeed and Clean Maps",
+			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
+				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
 					domain:   zoneRecords[0].Name,
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -310,13 +313,13 @@ func TestUpdateDNSServer(t *testing.T) {
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: true},
 			expectedUpstreamMap: make(registeredHandlerMap),
-			expectedLocalMap:    make(registrationMap),
+			expectedLocalQs:     []dns.Question{},
 		},
 		{
-			name:         "Disabled Service Should clean map",
-			initLocalMap: registrationMap{"netbird.cloud": struct{}{}},
+			name:             "Disabled Service Should clean map",
+			initLocalRecords: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}},
 			initUpstreamMap: registeredHandlerMap{
-				generateDummyHandler(zoneRecords[0].Name, nameServers).id(): handlerWrapper{
+				generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{
 					domain:   zoneRecords[0].Name,
 					handler:  dummyHandler,
 					priority: PriorityMatchDomain,
@@ -326,7 +329,7 @@ func TestUpdateDNSServer(t *testing.T) {
 			inputSerial:         1,
 			inputUpdate:         nbdns.Config{ServiceEnable: false},
 			expectedUpstreamMap: make(registeredHandlerMap),
-			expectedLocalMap:    make(registrationMap),
+			expectedLocalQs:     []dns.Question{},
 		},
 	}
 
@@ -377,7 +380,7 @@ func TestUpdateDNSServer(t *testing.T) {
 			}()
 
 			dnsServer.dnsMuxMap = testCase.initUpstreamMap
-			dnsServer.localResolver.registeredMap = testCase.initLocalMap
+			dnsServer.localResolver.Update(testCase.initLocalRecords)
 			dnsServer.updateSerial = testCase.initSerial
 
 			err = dnsServer.UpdateDNSServer(testCase.inputSerial, testCase.inputUpdate)
@@ -399,15 +402,23 @@ func TestUpdateDNSServer(t *testing.T) {
 				}
 			}
 
-			if len(dnsServer.localResolver.registeredMap) != len(testCase.expectedLocalMap) {
-				t.Fatalf("update local failed, registered map size is different than expected, want %d, got %d", len(testCase.expectedLocalMap), len(dnsServer.localResolver.registeredMap))
+			var responseMSG *dns.Msg
+			responseWriter := &test.MockResponseWriter{
+				WriteMsgFunc: func(m *dns.Msg) error {
+					responseMSG = m
+					return nil
+				},
+			}
+			for _, q := range testCase.expectedLocalQs {
+				dnsServer.localResolver.ServeDNS(responseWriter, &dns.Msg{
+					Question: []dns.Question{q},
+				})
 			}
 
-			for key := range testCase.expectedLocalMap {
-				_, found := dnsServer.localResolver.registeredMap[key]
-				if !found {
-					t.Fatalf("update local failed, key %s was not found in the localResolver.registeredMap: %#v", key, dnsServer.localResolver.registeredMap)
-				}
+			if len(testCase.expectedLocalQs) > 0 {
+				assert.NotNil(t, responseMSG, "response message should not be nil")
+				assert.Equal(t, dns.RcodeSuccess, responseMSG.Rcode, "response code should be success")
+				assert.NotEmpty(t, responseMSG.Answer, "response message should have answers")
 			}
 		})
 	}
@@ -491,11 +502,12 @@ func TestDNSFakeResolverHandleUpdates(t *testing.T) {
 	dnsServer.dnsMuxMap = registeredHandlerMap{
 		"id1": handlerWrapper{
 			domain:   zoneRecords[0].Name,
-			handler:  &localResolver{},
+			handler:  &local.Resolver{},
 			priority: PriorityMatchDomain,
 		},
 	}
-	dnsServer.localResolver.registeredMap = registrationMap{"netbird.cloud": struct{}{}}
+	//dnsServer.localResolver.RegisteredMap = local.RegistrationMap{local.BuildRecordKey("netbird.cloud", dns.ClassINET, dns.TypeA): struct{}{}}
+	dnsServer.localResolver.Update([]nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}})
 	dnsServer.updateSerial = 0
 
 	nameServers := []nbdns.NameServer{
@@ -582,7 +594,7 @@ func TestDNSServerStartStop(t *testing.T) {
 			}
 			time.Sleep(100 * time.Millisecond)
 			defer dnsServer.Stop()
-			_, err = dnsServer.localResolver.registerRecord(zoneRecords[0])
+			err = dnsServer.localResolver.RegisterRecord(zoneRecords[0])
 			if err != nil {
 				t.Error(err)
 			}
@@ -630,13 +642,11 @@ func TestDNSServerStartStop(t *testing.T) {
 func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 	hostManager := &mockHostConfigurator{}
 	server := DefaultServer{
-		ctx:     context.Background(),
-		service: NewServiceViaMemory(&mocWGIface{}),
-		localResolver: &localResolver{
-			registeredMap: make(registrationMap),
-		},
-		handlerChain: NewHandlerChain(),
-		hostManager:  hostManager,
+		ctx:           context.Background(),
+		service:       NewServiceViaMemory(&mocWGIface{}),
+		localResolver: local.NewResolver(),
+		handlerChain:  NewHandlerChain(),
+		hostManager:   hostManager,
 		currentConfig: HostDNSConfig{
 			Domains: []DomainConfig{
 				{false, "domain0", false},
@@ -1004,7 +1014,7 @@ func TestHandlerChain_DomainPriorities(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			r := new(dns.Msg)
 			r.SetQuestion(tc.query, dns.TypeA)
-			w := &ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			w := &ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
 
 			if mh, ok := tc.expectedHandler.(*MockHandler); ok {
 				mh.On("ServeDNS", mock.Anything, r).Once()
@@ -1037,9 +1047,9 @@ type mockHandler struct {
 }
 
 func (m *mockHandler) ServeDNS(dns.ResponseWriter, *dns.Msg) {}
-func (m *mockHandler) stop()                                 {}
-func (m *mockHandler) probeAvailability()                    {}
-func (m *mockHandler) id() handlerID                         { return handlerID(m.Id) }
+func (m *mockHandler) Stop()                                 {}
+func (m *mockHandler) ProbeAvailability()                    {}
+func (m *mockHandler) ID() types.HandlerID                   { return types.HandlerID(m.Id) }
 
 type mockService struct{}
 
@@ -1113,7 +1123,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 		name             string
 		initialHandlers  registeredHandlerMap
 		updates          []handlerWrapper
-		expectedHandlers map[string]string // map[handlerID]domain
+		expectedHandlers map[string]string // map[HandlerID]domain
 		description      string
 	}{
 		{
@@ -1409,7 +1419,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 
 			// Check each expected handler
 			for id, expectedDomain := range tt.expectedHandlers {
-				handler, exists := server.dnsMuxMap[handlerID(id)]
+				handler, exists := server.dnsMuxMap[types.HandlerID(id)]
 				assert.True(t, exists, "Expected handler %s not found", id)
 				if exists {
 					assert.Equal(t, expectedDomain, handler.domain,
@@ -1418,9 +1428,9 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 			}
 
 			// Verify no unexpected handlers exist
-			for handlerID := range server.dnsMuxMap {
-				_, expected := tt.expectedHandlers[string(handlerID)]
-				assert.True(t, expected, "Unexpected handler found: %s", handlerID)
+			for HandlerID := range server.dnsMuxMap {
+				_, expected := tt.expectedHandlers[string(HandlerID)]
+				assert.True(t, expected, "Unexpected handler found: %s", HandlerID)
 			}
 
 			// Verify the handlerChain state and order
@@ -1696,7 +1706,7 @@ func TestExtraDomains(t *testing.T) {
 				handlerChain:   NewHandlerChain(),
 				wgInterface:    &mocWGIface{},
 				hostManager:    mockHostConfig,
-				localResolver:  &localResolver{},
+				localResolver:  &local.Resolver{},
 				service:        mockSvc,
 				statusRecorder: peer.NewRecorder("test"),
 				extraDomains:   make(map[domain.Domain]int),
@@ -1781,7 +1791,7 @@ func TestExtraDomainsRefCounting(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &localResolver{},
+		localResolver:  &local.Resolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),
@@ -1833,7 +1843,7 @@ func TestUpdateConfigWithExistingExtraDomains(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &localResolver{},
+		localResolver:  &local.Resolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),
@@ -1916,7 +1926,7 @@ func TestDomainCaseHandling(t *testing.T) {
 		ctx:            context.Background(),
 		handlerChain:   NewHandlerChain(),
 		hostManager:    mockHostConfig,
-		localResolver:  &localResolver{},
+		localResolver:  &local.Resolver{},
 		service:        mockSvc,
 		statusRecorder: peer.NewRecorder("test"),
 		extraDomains:   make(map[domain.Domain]int),
diff --git a/client/internal/dns/test/mock.go b/client/internal/dns/test/mock.go
new file mode 100644
index 000000000..1db452805
--- /dev/null
+++ b/client/internal/dns/test/mock.go
@@ -0,0 +1,26 @@
+package test
+
+import (
+	"net"
+
+	"github.com/miekg/dns"
+)
+
+type MockResponseWriter struct {
+	WriteMsgFunc func(m *dns.Msg) error
+}
+
+func (rw *MockResponseWriter) WriteMsg(m *dns.Msg) error {
+	if rw.WriteMsgFunc != nil {
+		return rw.WriteMsgFunc(m)
+	}
+	return nil
+}
+
+func (rw *MockResponseWriter) LocalAddr() net.Addr       { return nil }
+func (rw *MockResponseWriter) RemoteAddr() net.Addr      { return nil }
+func (rw *MockResponseWriter) Write([]byte) (int, error) { return 0, nil }
+func (rw *MockResponseWriter) Close() error              { return nil }
+func (rw *MockResponseWriter) TsigStatus() error         { return nil }
+func (rw *MockResponseWriter) TsigTimersOnly(bool)       {}
+func (rw *MockResponseWriter) Hijack()                   {}
diff --git a/client/internal/dns/types/types.go b/client/internal/dns/types/types.go
new file mode 100644
index 000000000..5a8be03b7
--- /dev/null
+++ b/client/internal/dns/types/types.go
@@ -0,0 +1,3 @@
+package types
+
+type HandlerID string
diff --git a/client/internal/dns/upstream.go b/client/internal/dns/upstream.go
index fa69d4934..2fbfb3b91 100644
--- a/client/internal/dns/upstream.go
+++ b/client/internal/dns/upstream.go
@@ -19,6 +19,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/iface"
+	"github.com/netbirdio/netbird/client/internal/dns/types"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/proto"
 )
@@ -81,21 +82,21 @@ func (u *upstreamResolverBase) String() string {
 }
 
 // ID returns the unique handler ID
-func (u *upstreamResolverBase) id() handlerID {
+func (u *upstreamResolverBase) ID() types.HandlerID {
 	servers := slices.Clone(u.upstreamServers)
 	slices.Sort(servers)
 
 	hash := sha256.New()
 	hash.Write([]byte(u.domain + ":"))
 	hash.Write([]byte(strings.Join(servers, ",")))
-	return handlerID("upstream-" + hex.EncodeToString(hash.Sum(nil)[:8]))
+	return types.HandlerID("upstream-" + hex.EncodeToString(hash.Sum(nil)[:8]))
 }
 
 func (u *upstreamResolverBase) MatchSubdomains() bool {
 	return true
 }
 
-func (u *upstreamResolverBase) stop() {
+func (u *upstreamResolverBase) Stop() {
 	log.Debugf("stopping serving DNS for upstreams %s", u.upstreamServers)
 	u.cancel()
 }
@@ -198,9 +199,9 @@ func (u *upstreamResolverBase) checkUpstreamFails(err error) {
 	)
 }
 
-// probeAvailability tests all upstream servers simultaneously and
+// ProbeAvailability tests all upstream servers simultaneously and
 // disables the resolver if none work
-func (u *upstreamResolverBase) probeAvailability() {
+func (u *upstreamResolverBase) ProbeAvailability() {
 	u.mutex.Lock()
 	defer u.mutex.Unlock()
 
diff --git a/client/internal/dns/upstream_test.go b/client/internal/dns/upstream_test.go
index 5dbcc9f79..13bc91a37 100644
--- a/client/internal/dns/upstream_test.go
+++ b/client/internal/dns/upstream_test.go
@@ -8,6 +8,8 @@ import (
 	"time"
 
 	"github.com/miekg/dns"
+
+	"github.com/netbirdio/netbird/client/internal/dns/test"
 )
 
 func TestUpstreamResolver_ServeDNS(t *testing.T) {
@@ -66,7 +68,7 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			}
 
 			var responseMSG *dns.Msg
-			responseWriter := &mockResponseWriter{
+			responseWriter := &test.MockResponseWriter{
 				WriteMsgFunc: func(m *dns.Msg) error {
 					responseMSG = m
 					return nil
@@ -130,7 +132,7 @@ func TestUpstreamResolver_DeactivationReactivation(t *testing.T) {
 	resolver.failsTillDeact = 0
 	resolver.reactivatePeriod = time.Microsecond * 100
 
-	responseWriter := &mockResponseWriter{
+	responseWriter := &test.MockResponseWriter{
 		WriteMsgFunc: func(m *dns.Msg) error { return nil },
 	}
 
diff --git a/dns/dns.go b/dns/dns.go
index 3a1c76e56..f889a32ec 100644
--- a/dns/dns.go
+++ b/dns/dns.go
@@ -66,17 +66,17 @@ func (s SimpleRecord) String() string {
 func (s SimpleRecord) Len() uint16 {
 	emptyString := s.RData == ""
 	switch s.Type {
-	case 1:
+	case int(dns.TypeA):
 		if emptyString {
 			return 0
 		}
 		return net.IPv4len
-	case 5:
+	case int(dns.TypeCNAME):
 		if emptyString || s.RData == "." {
 			return 1
 		}
 		return uint16(len(s.RData) + 1)
-	case 28:
+	case int(dns.TypeAAAA):
 		if emptyString {
 			return 0
 		}

From 59faaa99f691d5ed9f743e7aca298ecc9af621d1 Mon Sep 17 00:00:00 2001
From: hakansa <43675540+hakansa@users.noreply.github.com>
Date: Mon, 5 May 2025 17:05:01 +0300
Subject: [PATCH 07/11] [client] Improve NetBird installation script to handle
 daemon connection timeout (#3761)

[client] Improve NetBird installation script to handle daemon connection timeout
---
 release_files/install.sh | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/release_files/install.sh b/release_files/install.sh
index 49e313f2f..da5c613d5 100755
--- a/release_files/install.sh
+++ b/release_files/install.sh
@@ -224,16 +224,22 @@ check_use_bin_variable() {
 
 install_netbird() {
     if [ -x "$(command -v netbird)" ]; then
-        status_output=$(netbird status)
-        if echo "$status_output" | grep -q 'Management: Connected' && echo "$status_output" | grep -q 'Signal: Connected'; then
-            echo "NetBird service is running, please stop it before proceeding"
-            exit 1
-        fi
+      status_output="$(netbird status 2>&1 || true)"
 
-        if [ -n "$status_output" ]; then
-            echo "NetBird seems to be installed already, please remove it before proceeding"
-            exit 1
-        fi
+      if echo "$status_output" | grep -q 'failed to connect to daemon error: context deadline exceeded'; then
+          echo "Warning: could not reach NetBird daemon (timeout), proceeding anyway"
+      else
+          if echo "$status_output" | grep -q 'Management: Connected' && \
+              echo "$status_output" | grep -q 'Signal: Connected'; then
+              echo "NetBird service is running, please stop it before proceeding"
+              exit 1
+          fi
+
+          if [ -n "$status_output" ]; then
+              echo "NetBird seems to be installed already, please remove it before proceeding"
+              exit 1
+          fi
+      fi
     fi
 
     # Run the installation, if a desktop environment is not detected

From 25faf9283da3361f77b1892fe52847235d340305 Mon Sep 17 00:00:00 2001
From: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com>
Date: Mon, 5 May 2025 18:21:48 +0200
Subject: [PATCH 08/11] [management] removal of foreign key constraint
 enforcement on sqlite (#3786)

---
 management/server/store/sql_store.go | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go
index d0adad6ee..dd39cf77d 100644
--- a/management/server/store/sql_store.go
+++ b/management/server/store/sql_store.go
@@ -82,10 +82,6 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
 			log.WithContext(ctx).Warnf("setting NB_SQL_MAX_OPEN_CONNS is not supported for sqlite, using default value 1")
 		}
 		conns = 1
-		_, err = sql.Exec("PRAGMA foreign_keys = ON")
-		if err != nil {
-			return nil, fmt.Errorf("failed to set foreign keys for sqlite: %w", err)
-		}
 	}
 
 	sql.SetMaxOpenConns(conns)

From ac135ab11dd4001b40bb47db97fcedbdde259f88 Mon Sep 17 00:00:00 2001
From: "M. Essam" <github@messam.xyz>
Date: Mon, 5 May 2025 19:54:47 +0300
Subject: [PATCH 09/11] [management/client/rest] fix panic when body is nil
 (#3714)

Fixes panic occurring when body is nil (this usually happens when connections is refused) due to lack of nil check by centralizing response.Body.Close() behavior.
---
 management/client/rest/accounts.go      | 12 +++--
 management/client/rest/dns.go           | 28 +++++++++---
 management/client/rest/events.go        |  4 +-
 management/client/rest/geo.go           |  8 +++-
 management/client/rest/groups.go        | 20 ++++++---
 management/client/rest/networks.go      | 60 ++++++++++++++++++-------
 management/client/rest/peers.go         | 20 ++++++---
 management/client/rest/policies.go      | 20 ++++++---
 management/client/rest/posturechecks.go | 20 ++++++---
 management/client/rest/routes.go        | 20 ++++++---
 management/client/rest/setupkeys.go     | 20 ++++++---
 management/client/rest/tokens.go        | 16 +++++--
 management/client/rest/users.go         | 24 +++++++---
 13 files changed, 204 insertions(+), 68 deletions(-)

diff --git a/management/client/rest/accounts.go b/management/client/rest/accounts.go
index f38b19f70..29d4ac79d 100644
--- a/management/client/rest/accounts.go
+++ b/management/client/rest/accounts.go
@@ -20,7 +20,9 @@ func (a *AccountsAPI) List(ctx context.Context) ([]api.Account, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Account](resp)
 	return ret, err
 }
@@ -36,7 +38,9 @@ func (a *AccountsAPI) Update(ctx context.Context, accountID string, request api.
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Account](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *AccountsAPI) Delete(ctx context.Context, accountID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/dns.go b/management/client/rest/dns.go
index ef9923b1f..0e2d15842 100644
--- a/management/client/rest/dns.go
+++ b/management/client/rest/dns.go
@@ -20,7 +20,9 @@ func (a *DNSAPI) ListNameserverGroups(ctx context.Context) ([]api.NameserverGrou
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.NameserverGroup](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *DNSAPI) GetNameserverGroup(ctx context.Context, nameserverGroupID strin
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NameserverGroup](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *DNSAPI) CreateNameserverGroup(ctx context.Context, request api.PostApiD
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NameserverGroup](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *DNSAPI) UpdateNameserverGroup(ctx context.Context, nameserverGroupID st
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NameserverGroup](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *DNSAPI) DeleteNameserverGroup(ctx context.Context, nameserverGroupID st
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -88,7 +98,9 @@ func (a *DNSAPI) GetSettings(ctx context.Context) (*api.DNSSettings, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.DNSSettings](resp)
 	return &ret, err
 }
@@ -104,7 +116,9 @@ func (a *DNSAPI) UpdateSettings(ctx context.Context, request api.PutApiDnsSettin
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.DNSSettings](resp)
 	return &ret, err
 }
diff --git a/management/client/rest/events.go b/management/client/rest/events.go
index 1157700ff..ed74fae39 100644
--- a/management/client/rest/events.go
+++ b/management/client/rest/events.go
@@ -18,7 +18,9 @@ func (a *EventsAPI) List(ctx context.Context) ([]api.Event, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Event](resp)
 	return ret, err
 }
diff --git a/management/client/rest/geo.go b/management/client/rest/geo.go
index ed9090fe2..0bdcc0a22 100644
--- a/management/client/rest/geo.go
+++ b/management/client/rest/geo.go
@@ -18,7 +18,9 @@ func (a *GeoLocationAPI) ListCountries(ctx context.Context) ([]api.Country, erro
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Country](resp)
 	return ret, err
 }
@@ -30,7 +32,9 @@ func (a *GeoLocationAPI) ListCountryCities(ctx context.Context, countryCode stri
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.City](resp)
 	return ret, err
 }
diff --git a/management/client/rest/groups.go b/management/client/rest/groups.go
index feb664273..aac453b93 100644
--- a/management/client/rest/groups.go
+++ b/management/client/rest/groups.go
@@ -20,7 +20,9 @@ func (a *GroupsAPI) List(ctx context.Context) ([]api.Group, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Group](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *GroupsAPI) Get(ctx context.Context, groupID string) (*api.Group, error)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Group](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *GroupsAPI) Create(ctx context.Context, request api.PostApiGroupsJSONReq
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Group](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *GroupsAPI) Update(ctx context.Context, groupID string, request api.PutA
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Group](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *GroupsAPI) Delete(ctx context.Context, groupID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/networks.go b/management/client/rest/networks.go
index 2cdd6d73d..b211312c9 100644
--- a/management/client/rest/networks.go
+++ b/management/client/rest/networks.go
@@ -20,7 +20,9 @@ func (a *NetworksAPI) List(ctx context.Context) ([]api.Network, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Network](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *NetworksAPI) Get(ctx context.Context, networkID string) (*api.Network,
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Network](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *NetworksAPI) Create(ctx context.Context, request api.PostApiNetworksJSO
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Network](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *NetworksAPI) Update(ctx context.Context, networkID string, request api.
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Network](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *NetworksAPI) Delete(ctx context.Context, networkID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -102,7 +112,9 @@ func (a *NetworkResourcesAPI) List(ctx context.Context) ([]api.NetworkResource,
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.NetworkResource](resp)
 	return ret, err
 }
@@ -114,7 +126,9 @@ func (a *NetworkResourcesAPI) Get(ctx context.Context, networkResourceID string)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkResource](resp)
 	return &ret, err
 }
@@ -130,7 +144,9 @@ func (a *NetworkResourcesAPI) Create(ctx context.Context, request api.PostApiNet
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkResource](resp)
 	return &ret, err
 }
@@ -146,7 +162,9 @@ func (a *NetworkResourcesAPI) Update(ctx context.Context, networkResourceID stri
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkResource](resp)
 	return &ret, err
 }
@@ -158,7 +176,9 @@ func (a *NetworkResourcesAPI) Delete(ctx context.Context, networkResourceID stri
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -184,7 +204,9 @@ func (a *NetworkRoutersAPI) List(ctx context.Context) ([]api.NetworkRouter, erro
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.NetworkRouter](resp)
 	return ret, err
 }
@@ -196,7 +218,9 @@ func (a *NetworkRoutersAPI) Get(ctx context.Context, networkRouterID string) (*a
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkRouter](resp)
 	return &ret, err
 }
@@ -212,7 +236,9 @@ func (a *NetworkRoutersAPI) Create(ctx context.Context, request api.PostApiNetwo
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkRouter](resp)
 	return &ret, err
 }
@@ -228,7 +254,9 @@ func (a *NetworkRoutersAPI) Update(ctx context.Context, networkRouterID string,
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.NetworkRouter](resp)
 	return &ret, err
 }
@@ -240,7 +268,9 @@ func (a *NetworkRoutersAPI) Delete(ctx context.Context, networkRouterID string)
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/peers.go b/management/client/rest/peers.go
index 9d35f013c..2b1a65b4c 100644
--- a/management/client/rest/peers.go
+++ b/management/client/rest/peers.go
@@ -20,7 +20,9 @@ func (a *PeersAPI) List(ctx context.Context) ([]api.Peer, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Peer](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *PeersAPI) Get(ctx context.Context, peerID string) (*api.Peer, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Peer](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *PeersAPI) Update(ctx context.Context, peerID string, request api.PutApi
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Peer](resp)
 	return &ret, err
 }
@@ -60,7 +66,9 @@ func (a *PeersAPI) Delete(ctx context.Context, peerID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -72,7 +80,9 @@ func (a *PeersAPI) ListAccessiblePeers(ctx context.Context, peerID string) ([]ap
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Peer](resp)
 	return ret, err
 }
diff --git a/management/client/rest/policies.go b/management/client/rest/policies.go
index be6abafaf..975a95440 100644
--- a/management/client/rest/policies.go
+++ b/management/client/rest/policies.go
@@ -20,7 +20,9 @@ func (a *PoliciesAPI) List(ctx context.Context) ([]api.Policy, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Policy](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *PoliciesAPI) Get(ctx context.Context, policyID string) (*api.Policy, er
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Policy](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *PoliciesAPI) Create(ctx context.Context, request api.PostApiPoliciesJSO
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Policy](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *PoliciesAPI) Update(ctx context.Context, policyID string, request api.P
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Policy](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *PoliciesAPI) Delete(ctx context.Context, policyID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/posturechecks.go b/management/client/rest/posturechecks.go
index 950d17ba0..7343957a5 100644
--- a/management/client/rest/posturechecks.go
+++ b/management/client/rest/posturechecks.go
@@ -20,7 +20,9 @@ func (a *PostureChecksAPI) List(ctx context.Context) ([]api.PostureCheck, error)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.PostureCheck](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *PostureChecksAPI) Get(ctx context.Context, postureCheckID string) (*api
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.PostureCheck](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *PostureChecksAPI) Create(ctx context.Context, request api.PostApiPostur
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.PostureCheck](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *PostureChecksAPI) Update(ctx context.Context, postureCheckID string, re
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.PostureCheck](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *PostureChecksAPI) Delete(ctx context.Context, postureCheckID string) er
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/routes.go b/management/client/rest/routes.go
index bccbb8847..6ca4be2c5 100644
--- a/management/client/rest/routes.go
+++ b/management/client/rest/routes.go
@@ -20,7 +20,9 @@ func (a *RoutesAPI) List(ctx context.Context) ([]api.Route, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.Route](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *RoutesAPI) Get(ctx context.Context, routeID string) (*api.Route, error)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Route](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *RoutesAPI) Create(ctx context.Context, request api.PostApiRoutesJSONReq
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Route](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *RoutesAPI) Update(ctx context.Context, routeID string, request api.PutA
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.Route](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *RoutesAPI) Delete(ctx context.Context, routeID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/setupkeys.go b/management/client/rest/setupkeys.go
index 645614fcf..91f370663 100644
--- a/management/client/rest/setupkeys.go
+++ b/management/client/rest/setupkeys.go
@@ -20,7 +20,9 @@ func (a *SetupKeysAPI) List(ctx context.Context) ([]api.SetupKey, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.SetupKey](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *SetupKeysAPI) Get(ctx context.Context, setupKeyID string) (*api.SetupKe
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.SetupKey](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *SetupKeysAPI) Create(ctx context.Context, request api.PostApiSetupKeysJ
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.SetupKeyClear](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *SetupKeysAPI) Update(ctx context.Context, setupKeyID string, request ap
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.SetupKey](resp)
 	return &ret, err
 }
@@ -76,7 +84,9 @@ func (a *SetupKeysAPI) Delete(ctx context.Context, setupKeyID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/tokens.go b/management/client/rest/tokens.go
index 3275bea81..7e5004147 100644
--- a/management/client/rest/tokens.go
+++ b/management/client/rest/tokens.go
@@ -20,7 +20,9 @@ func (a *TokensAPI) List(ctx context.Context, userID string) ([]api.PersonalAcce
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.PersonalAccessToken](resp)
 	return ret, err
 }
@@ -32,7 +34,9 @@ func (a *TokensAPI) Get(ctx context.Context, userID, tokenID string) (*api.Perso
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.PersonalAccessToken](resp)
 	return &ret, err
 }
@@ -48,7 +52,9 @@ func (a *TokensAPI) Create(ctx context.Context, userID string, request api.PostA
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.PersonalAccessTokenGenerated](resp)
 	return &ret, err
 }
@@ -60,7 +66,9 @@ func (a *TokensAPI) Delete(ctx context.Context, userID, tokenID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
diff --git a/management/client/rest/users.go b/management/client/rest/users.go
index 31ffad051..bb81796c0 100644
--- a/management/client/rest/users.go
+++ b/management/client/rest/users.go
@@ -20,7 +20,9 @@ func (a *UsersAPI) List(ctx context.Context) ([]api.User, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[[]api.User](resp)
 	return ret, err
 }
@@ -36,7 +38,9 @@ func (a *UsersAPI) Create(ctx context.Context, request api.PostApiUsersJSONReque
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.User](resp)
 	return &ret, err
 }
@@ -52,7 +56,9 @@ func (a *UsersAPI) Update(ctx context.Context, userID string, request api.PutApi
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 	ret, err := parseResponse[api.User](resp)
 	return &ret, err
 }
@@ -64,7 +70,9 @@ func (a *UsersAPI) Delete(ctx context.Context, userID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -76,7 +84,9 @@ func (a *UsersAPI) ResendInvitation(ctx context.Context, userID string) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	return nil
 }
@@ -88,7 +98,9 @@ func (a *UsersAPI) Current(ctx context.Context) (*api.User, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
 
 	ret, err := parseResponse[api.User](resp)
 	return &ret, err

From ebda0fc5385d26ebdbd83904f6d72e36f11ce156 Mon Sep 17 00:00:00 2001
From: Bethuel Mmbaga <bethuelmbaga12@gmail.com>
Date: Tue, 6 May 2025 18:31:03 +0300
Subject: [PATCH 10/11] [management] Delete service users with account manager
 (#3793)

---
 management/server/account.go      |  8 ++++--
 management/server/account_test.go | 44 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index aa7cb0019..5c474a343 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -603,11 +603,15 @@ func (am *DefaultAccountManager) DeleteAccount(ctx context.Context, accountID, u
 	}
 
 	for _, otherUser := range account.Users {
-		if otherUser.IsServiceUser {
+		if otherUser.Id == userID {
 			continue
 		}
 
-		if otherUser.Id == userID {
+		if otherUser.IsServiceUser {
+			err = am.deleteServiceUser(ctx, accountID, userID, otherUser)
+			if err != nil {
+				return err
+			}
 			continue
 		}
 
diff --git a/management/server/account_test.go b/management/server/account_test.go
index fe082d9a0..c5583d226 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -853,6 +853,42 @@ func TestAccountManager_DeleteAccount(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	account.Users["service-user-1"] = &types.User{
+		Id:            "service-user-1",
+		Role:          types.UserRoleAdmin,
+		IsServiceUser: true,
+		Issued:        types.UserIssuedAPI,
+		PATs: map[string]*types.PersonalAccessToken{
+			"pat-1": {
+				ID:          "pat-1",
+				UserID:      "service-user-1",
+				Name:        "service-user-1",
+				HashedToken: "hashedToken",
+				CreatedAt:   time.Now(),
+			},
+		},
+	}
+	account.Users[userId] = &types.User{
+		Id:            "service-user-2",
+		Role:          types.UserRoleUser,
+		IsServiceUser: true,
+		Issued:        types.UserIssuedAPI,
+		PATs: map[string]*types.PersonalAccessToken{
+			"pat-2": {
+				ID:          "pat-2",
+				UserID:      userId,
+				Name:        userId,
+				HashedToken: "hashedToken",
+				CreatedAt:   time.Now(),
+			},
+		},
+	}
+
+	err = manager.Store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	err = manager.DeleteAccount(context.Background(), account.Id, userId)
 	if err != nil {
 		t.Fatal(err)
@@ -862,6 +898,14 @@ func TestAccountManager_DeleteAccount(t *testing.T) {
 	if err == nil {
 		t.Fatal(fmt.Errorf("expected to get an error when trying to get deleted account, got %v", getAccount))
 	}
+
+	pats, err := manager.Store.GetUserPATs(context.Background(), store.LockingStrengthShare, "service-user-1")
+	require.NoError(t, err)
+	assert.Len(t, pats, 0)
+
+	pats, err = manager.Store.GetUserPATs(context.Background(), store.LockingStrengthShare, userId)
+	require.NoError(t, err)
+	assert.Len(t, pats, 0)
 }
 
 func BenchmarkTest_GetAccountWithclaims(b *testing.B) {

From fcd2c15a37320064ef29aa036ebf32a8d5b69714 Mon Sep 17 00:00:00 2001
From: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com>
Date: Wed, 7 May 2025 07:25:25 +0200
Subject: [PATCH 11/11] [management] policy delete cleans policy rules (#3788)

---
 management/server/store/sql_store.go | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go
index dd39cf77d..d568460f9 100644
--- a/management/server/store/sql_store.go
+++ b/management/server/store/sql_store.go
@@ -1683,18 +1683,26 @@ func (s *SqlStore) SavePolicy(ctx context.Context, lockStrength LockingStrength,
 }
 
 func (s *SqlStore) DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&types.Policy{}, accountAndIDQueryCondition, accountID, policyID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
-		return status.Errorf(status.Internal, "failed to delete policy from store")
-	}
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if err := tx.Where("policy_id = ?", policyID).Delete(&types.PolicyRule{}).Error; err != nil {
+			return fmt.Errorf("delete policy rules: %w", err)
+		}
 
-	if result.RowsAffected == 0 {
-		return status.NewPolicyNotFoundError(policyID)
-	}
+		result := tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Where(accountAndIDQueryCondition, accountID, policyID).
+			Delete(&types.Policy{})
 
-	return nil
+		if err := result.Error; err != nil {
+			log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
+			return status.Errorf(status.Internal, "failed to delete policy from store")
+		}
+
+		if result.RowsAffected == 0 {
+			return status.NewPolicyNotFoundError(policyID)
+		}
+
+		return nil
+	})
 }
 
 // GetAccountPostureChecks retrieves posture checks for an account.