minor change

client/firewall/uspfilter/filter.go

@@ -50,6 +50,12 @@ const (
 
 var errNatNotSupported = errors.New("nat not supported with userspace firewall")
 
+// serviceKey represents a protocol/port combination for netstack service registry
+type serviceKey struct {
+	protocol gopacket.LayerType
+	port     uint16
+}
+
 // RuleSet is a set of rules grouped by a string key
 type RuleSet map[string]PeerRule
 
@@ -113,6 +119,9 @@ type Manager struct {
 	portDNATEnabled atomic.Bool
 	portDNATRules   []portDNATRule
 	portDNATMutex   sync.RWMutex
+
+	netstackServices     map[serviceKey]struct{}
+	netstackServiceMutex sync.RWMutex
 }
 
 // decoder for packages
@@ -203,6 +212,7 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 		localForwarding:     enableLocalForwarding,
 		dnatMappings:        make(map[netip.Addr]netip.Addr),
 		portDNATRules:       []portDNATRule{},
+		netstackServices:    make(map[serviceKey]struct{}),
 	}
 	m.routingEnabled.Store(false)
 
@@ -838,9 +848,7 @@ func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP netip.Addr, packet
 		return true
 	}
 
-	// If requested we pass local traffic to internal interfaces to the forwarder.
-	// netstack doesn't have an interface to forward packets to the native stack so we always need to use the forwarder.
-	if m.localForwarding && (m.netstack || dstIP != m.wgIface.Address().IP) {
+	if m.shouldForward(d, dstIP) {
 		return m.handleForwardedLocalTraffic(packetData)
 	}
 
@@ -1274,3 +1282,86 @@ func (m *Manager) DisableRouting() error {
 
 	return nil
 }
+
+// RegisterNetstackService registers a service as listening on the netstack for the given protocol and port
+func (m *Manager) RegisterNetstackService(protocol nftypes.Protocol, port uint16) {
+	m.netstackServiceMutex.Lock()
+	defer m.netstackServiceMutex.Unlock()
+	layerType := m.protocolToLayerType(protocol)
+	key := serviceKey{protocol: layerType, port: port}
+	m.netstackServices[key] = struct{}{}
+	m.logger.Debug3("RegisterNetstackService: registered %s:%d (layerType=%s)", protocol, port, layerType)
+	m.logger.Debug1("RegisterNetstackService: current registry size: %d", len(m.netstackServices))
+}
+
+// UnregisterNetstackService removes a service from the netstack registry
+func (m *Manager) UnregisterNetstackService(protocol nftypes.Protocol, port uint16) {
+	m.netstackServiceMutex.Lock()
+	defer m.netstackServiceMutex.Unlock()
+	layerType := m.protocolToLayerType(protocol)
+	key := serviceKey{protocol: layerType, port: port}
+	delete(m.netstackServices, key)
+	m.logger.Debug2("Unregistered netstack service on protocol %s port %d", protocol, port)
+}
+
+// protocolToLayerType converts nftypes.Protocol to gopacket.LayerType for internal use
+func (m *Manager) protocolToLayerType(protocol nftypes.Protocol) gopacket.LayerType {
+	switch protocol {
+	case nftypes.TCP:
+		return layers.LayerTypeTCP
+	case nftypes.UDP:
+		return layers.LayerTypeUDP
+	case nftypes.ICMP:
+		return layers.LayerTypeICMPv4
+	default:
+		return gopacket.LayerType(0) // Invalid/unknown
+	}
+}
+
+// shouldForward determines if a packet should be forwarded to the forwarder.
+// The forwarder handles routing packets to the native OS network stack.
+// Returns true if packet should go to the forwarder, false if it should go to netstack listeners or the native stack directly.
+func (m *Manager) shouldForward(d *decoder, dstIP netip.Addr) bool {
+	// not enabled, never forward
+	if !m.localForwarding {
+		return false
+	}
+
+	// netstack always needs to forward because it's lacking a native interface
+	// exception for registered netstack services, those should go to netstack listeners
+	if m.netstack {
+		return !m.hasMatchingNetstackService(d)
+	}
+
+	// traffic to our other local interfaces (not NetBird IP) - always forward
+	if dstIP != m.wgIface.Address().IP {
+		return true
+	}
+
+	// traffic to our NetBird IP, not netstack mode - send to netstack listeners
+	return false
+}
+
+// hasMatchingNetstackService checks if there's a registered netstack service for this packet
+func (m *Manager) hasMatchingNetstackService(d *decoder) bool {
+	if len(d.decoded) < 2 {
+		return false
+	}
+
+	var dstPort uint16
+	switch d.decoded[1] {
+	case layers.LayerTypeTCP:
+		dstPort = uint16(d.tcp.DstPort)
+	case layers.LayerTypeUDP:
+		dstPort = uint16(d.udp.DstPort)
+	default:
+		return false
+	}
+
+	key := serviceKey{protocol: d.decoded[1], port: dstPort}
+	m.netstackServiceMutex.RLock()
+	_, exists := m.netstackServices[key]
+	m.netstackServiceMutex.RUnlock()
+
+	return exists
+}

client/grpc/dialer.go

@@ -4,12 +4,15 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"errors"
+	"fmt"
 	"runtime"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/keepalive"
@@ -17,6 +20,9 @@ import (
 	"github.com/netbirdio/netbird/util/embeddedroots"
 )
 
+// ErrConnectionShutdown indicates that the connection entered shutdown state before becoming ready
+var ErrConnectionShutdown = errors.New("connection shutdown before ready")
+
 // Backoff returns a backoff configuration for gRPC calls
 func Backoff(ctx context.Context) backoff.BackOff {
 	b := backoff.NewExponentialBackOff()
@@ -25,6 +31,26 @@ func Backoff(ctx context.Context) backoff.BackOff {
 	return backoff.WithContext(b, ctx)
 }
 
+// waitForConnectionReady blocks until the connection becomes ready or fails.
+// Returns an error if the connection times out, is cancelled, or enters shutdown state.
+func waitForConnectionReady(ctx context.Context, conn *grpc.ClientConn) error {
+	conn.Connect()
+
+	state := conn.GetState()
+	for state != connectivity.Ready && state != connectivity.Shutdown {
+		if !conn.WaitForStateChange(ctx, state) {
+			return fmt.Errorf("wait state change from %s: %w", state, ctx.Err())
+		}
+		state = conn.GetState()
+	}
+
+	if state == connectivity.Shutdown {
+		return ErrConnectionShutdown
+	}
+
+	return nil
+}
+
 // CreateConnection creates a gRPC client connection with the appropriate transport options.
 // The component parameter specifies the WebSocket proxy component path (e.g., "/management", "/signal").
 func CreateConnection(ctx context.Context, addr string, tlsEnabled bool, component string) (*grpc.ClientConn, error) {
@@ -42,22 +68,24 @@ func CreateConnection(ctx context.Context, addr string, tlsEnabled bool, compone
 		}))
 	}
 
-	connCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
-
-	conn, err := grpc.DialContext(
-		connCtx,
+	conn, err := grpc.NewClient(
 		addr,
 		transportOption,
 		WithCustomDialer(tlsEnabled, component),
-		grpc.WithBlock(),
 		grpc.WithKeepaliveParams(keepalive.ClientParameters{
 			Time:    30 * time.Second,
 			Timeout: 10 * time.Second,
 		}),
 	)
 	if err != nil {
-		log.Printf("DialContext error: %v", err)
+		return nil, fmt.Errorf("new client: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
+	if err := waitForConnectionReady(ctx, conn); err != nil {
+		_ = conn.Close()
 		return nil, err
 	}
 

client/grpc/dialer_generic.go

@@ -18,7 +18,7 @@ import (
 	nbnet "github.com/netbirdio/netbird/client/net"
 )
 
-func WithCustomDialer(tlsEnabled bool, component string) grpc.DialOption {
+func WithCustomDialer(_ bool, _ string) grpc.DialOption {
 	return grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
 		if runtime.GOOS == "linux" {
 			currentUser, err := user.Current()
@@ -36,7 +36,6 @@ func WithCustomDialer(tlsEnabled bool, component string) grpc.DialOption {
 
 		conn, err := nbnet.NewDialer().DialContext(ctx, "tcp", addr)
 		if err != nil {
-			log.Errorf("Failed to dial: %s", err)
 			return nil, fmt.Errorf("nbnet.NewDialer().DialContext: %w", err)
 		}
 		return conn, nil

client/internal/dnsfwd/cache_test.go

@@ -83,4 +83,3 @@ func TestCacheMiss(t *testing.T) {
 		t.Fatalf("expected cache miss, got=%v ok=%v", got, ok)
 	}
 }
-

client/internal/dnsfwd/forwarder.go

@@ -14,6 +14,7 @@ import (
 	"github.com/hashicorp/go-multierror"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/tun/netstack"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
@@ -33,7 +34,7 @@ type firewaller interface {
 }
 
 type DNSForwarder struct {
-	listenAddress  string
+	listenAddress  netip.AddrPort
 	ttl            uint32
 	statusRecorder *peer.Status
 
@@ -47,9 +48,11 @@ type DNSForwarder struct {
 	firewall   firewaller
 	resolver   resolver
 	cache      *cache
+
+	wgIface wgIface
 }
 
-func NewDNSForwarder(listenAddress string, ttl uint32, firewall firewaller, statusRecorder *peer.Status) *DNSForwarder {
+func NewDNSForwarder(listenAddress netip.AddrPort, ttl uint32, firewall firewaller, statusRecorder *peer.Status, wgIface wgIface) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
 		listenAddress:  listenAddress,
@@ -58,30 +61,46 @@ func NewDNSForwarder(listenAddress string, ttl uint32, firewall firewaller, stat
 		statusRecorder: statusRecorder,
 		resolver:       net.DefaultResolver,
 		cache:          newCache(),
+		wgIface:        wgIface,
 	}
 }
 
 func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
-	log.Infof("starting DNS forwarder on address=%s", f.listenAddress)
+	var netstackNet *netstack.Net
+	if f.wgIface != nil {
+		netstackNet = f.wgIface.GetNet()
+	}
+
+	addrDesc := f.listenAddress.String()
+	if netstackNet != nil {
+		addrDesc = fmt.Sprintf("netstack %s", f.listenAddress)
+	}
+	log.Infof("starting DNS forwarder on address=%s", addrDesc)
+
+	udpLn, err := f.createUDPListener(netstackNet)
+	if err != nil {
+		return fmt.Errorf("create UDP listener: %w", err)
+	}
+
+	tcpLn, err := f.createTCPListener(netstackNet)
+	if err != nil {
+		return fmt.Errorf("create TCP listener: %w", err)
+	}
 
-	// UDP server
 	mux := dns.NewServeMux()
 	f.mux = mux
 	mux.HandleFunc(".", f.handleDNSQueryUDP)
 	f.dnsServer = &dns.Server{
-		Addr:    f.listenAddress,
-		Net:     "udp",
-		Handler: mux,
+		PacketConn: udpLn,
+		Handler:    mux,
 	}
 
-	// TCP server
 	tcpMux := dns.NewServeMux()
 	f.tcpMux = tcpMux
 	tcpMux.HandleFunc(".", f.handleDNSQueryTCP)
 	f.tcpServer = &dns.Server{
-		Addr:    f.listenAddress,
-		Net:     "tcp",
-		Handler: tcpMux,
+		Listener: tcpLn,
+		Handler:  tcpMux,
 	}
 
 	f.UpdateDomains(entries)
@@ -89,18 +108,33 @@ func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
 	errCh := make(chan error, 2)
 
 	go func() {
-		log.Infof("DNS UDP listener running on %s", f.listenAddress)
-		errCh <- f.dnsServer.ListenAndServe()
+		log.Infof("DNS UDP listener running on %s", addrDesc)
+		errCh <- f.dnsServer.ActivateAndServe()
 	}()
 	go func() {
-		log.Infof("DNS TCP listener running on %s", f.listenAddress)
-		errCh <- f.tcpServer.ListenAndServe()
+		log.Infof("DNS TCP listener running on %s", addrDesc)
+		errCh <- f.tcpServer.ActivateAndServe()
 	}()
 
-	// return the first error we get (e.g. bind failure or shutdown)
 	return <-errCh
 }
 
+func (f *DNSForwarder) createUDPListener(netstackNet *netstack.Net) (net.PacketConn, error) {
+	if netstackNet != nil {
+		return netstackNet.ListenUDPAddrPort(f.listenAddress)
+	}
+
+	return net.ListenUDP("udp", net.UDPAddrFromAddrPort(f.listenAddress))
+}
+
+func (f *DNSForwarder) createTCPListener(netstackNet *netstack.Net) (net.Listener, error) {
+	if netstackNet != nil {
+		return netstackNet.ListenTCPAddrPort(f.listenAddress)
+	}
+
+	return net.ListenTCP("tcp", net.TCPAddrFromAddrPort(f.listenAddress))
+}
+
 func (f *DNSForwarder) UpdateDomains(entries []*ForwarderEntry) {
 	f.mutex.Lock()
 	defer f.mutex.Unlock()

client/internal/dnsfwd/forwarder_test.go

@@ -297,7 +297,7 @@ func TestDNSForwarder_UnauthorizedDomainAccess(t *testing.T) {
 				mockResolver.On("LookupNetIP", mock.Anything, "ip4", dns.Fqdn(tt.queryDomain)).Return([]netip.Addr{fakeIP}, nil)
 			}
 
-			forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
+			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
 			forwarder.resolver = mockResolver
 
 			d, err := domain.FromString(tt.configuredDomain)
@@ -402,7 +402,7 @@ func TestDNSForwarder_FirewallSetUpdates(t *testing.T) {
 			mockResolver := &MockResolver{}
 
 			// Set up forwarder
-			forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
+			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
 			forwarder.resolver = mockResolver
 
 			// Create entries and track sets
@@ -489,7 +489,7 @@ func TestDNSForwarder_MultipleIPsInSingleUpdate(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	// Configure a single domain
@@ -584,7 +584,7 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
+			forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
 
 			d, err := domain.FromString(tt.configured)
 			require.NoError(t, err)
@@ -616,7 +616,7 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
 func TestDNSForwarder_TCPTruncation(t *testing.T) {
 	// Test that large UDP responses are truncated with TC bit set
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	d, _ := domain.FromString("example.com")
@@ -652,7 +652,7 @@ func TestDNSForwarder_TCPTruncation(t *testing.T) {
 // a subsequent upstream failure still returns a successful response from cache.
 func TestDNSForwarder_ServeFromCacheOnUpstreamFailure(t *testing.T) {
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("example.com")
@@ -696,7 +696,7 @@ func TestDNSForwarder_ServeFromCacheOnUpstreamFailure(t *testing.T) {
 // Verifies that cache normalization works across casing and trailing dot variations.
 func TestDNSForwarder_CacheNormalizationCasingAndDot(t *testing.T) {
 	mockResolver := &MockResolver{}
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("ExAmPlE.CoM")
@@ -742,7 +742,7 @@ func TestDNSForwarder_MultipleOverlappingPatterns(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	// Set up complex overlapping patterns
@@ -804,7 +804,7 @@ func TestDNSForwarder_NodataVsNxdomain(t *testing.T) {
 	mockFirewall := &MockFirewall{}
 	mockResolver := &MockResolver{}
 
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, mockFirewall, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, mockFirewall, &peer.Status{}, nil)
 	forwarder.resolver = mockResolver
 
 	d, err := domain.FromString("example.com")
@@ -925,7 +925,7 @@ func TestDNSForwarder_NodataVsNxdomain(t *testing.T) {
 
 func TestDNSForwarder_EmptyQuery(t *testing.T) {
 	// Test handling of malformed query with no questions
-	forwarder := NewDNSForwarder("127.0.0.1:0", 300, nil, &peer.Status{})
+	forwarder := NewDNSForwarder(netip.MustParseAddrPort("127.0.0.1:0"), 300, nil, &peer.Status{}, nil)
 
 	query := &dns.Msg{}
 	// Don't set any question

client/internal/dnsfwd/manager.go

@@ -10,9 +10,11 @@ import (
 
 	"github.com/hashicorp/go-multierror"
 	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/tun/netstack"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
@@ -24,6 +26,12 @@ const (
 	envServerPort = "NB_DNS_FORWARDER_PORT"
 )
 
+// wgIface defines the interface for WireGuard interface operations needed by the DNS forwarder.
+type wgIface interface {
+	GetNet() *netstack.Net
+	Address() wgaddr.Address
+}
+
 // ForwarderEntry is a mapping from a domain to a resource ID and a hash of the parent domain list.
 type ForwarderEntry struct {
 	Domain domain.Domain
@@ -34,7 +42,7 @@ type ForwarderEntry struct {
 type Manager struct {
 	firewall       firewall.Manager
 	statusRecorder *peer.Status
-	localAddr      netip.Addr
+	wgIface        wgIface
 	serverPort     uint16
 
 	fwRules      []firewall.Rule
@@ -42,7 +50,7 @@ type Manager struct {
 	dnsForwarder *DNSForwarder
 }
 
-func NewManager(fw firewall.Manager, statusRecorder *peer.Status, localAddr netip.Addr) *Manager {
+func NewManager(fw firewall.Manager, statusRecorder *peer.Status, wgIface wgIface) *Manager {
 	serverPort := nbdns.ForwarderServerPort
 	if envPort := os.Getenv(envServerPort); envPort != "" {
 		if port, err := strconv.ParseUint(envPort, 10, 16); err == nil && port > 0 {
@@ -56,7 +64,7 @@ func NewManager(fw firewall.Manager, statusRecorder *peer.Status, localAddr neti
 	return &Manager{
 		firewall:       fw,
 		statusRecorder: statusRecorder,
-		localAddr:      localAddr,
+		wgIface:        wgIface,
 		serverPort:     serverPort,
 	}
 }
@@ -71,21 +79,25 @@ func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
 		return err
 	}
 
-	if m.localAddr.IsValid() && m.firewall != nil {
-		if err := m.firewall.AddInboundDNAT(m.localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
+	localAddr := m.wgIface.Address().IP
+
+	if localAddr.IsValid() && m.firewall != nil {
+		if err := m.firewall.AddInboundDNAT(localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
 			log.Warnf("failed to add DNS UDP DNAT rule: %v", err)
 		} else {
-			log.Infof("added DNS UDP DNAT rule: %s:%d -> %s:%d", m.localAddr, nbdns.ForwarderClientPort, m.localAddr, m.serverPort)
+			log.Infof("added DNS UDP DNAT rule: %s:%d -> %s:%d", localAddr, nbdns.ForwarderClientPort, localAddr, m.serverPort)
 		}
 
-		if err := m.firewall.AddInboundDNAT(m.localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
+		if err := m.firewall.AddInboundDNAT(localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
 			log.Warnf("failed to add DNS TCP DNAT rule: %v", err)
 		} else {
-			log.Infof("added DNS TCP DNAT rule: %s:%d -> %s:%d", m.localAddr, nbdns.ForwarderClientPort, m.localAddr, m.serverPort)
+			log.Infof("added DNS TCP DNAT rule: %s:%d -> %s:%d", localAddr, nbdns.ForwarderClientPort, localAddr, m.serverPort)
 		}
 	}
 
-	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", m.serverPort), dnsTTL, m.firewall, m.statusRecorder)
+	listenAddress := netip.AddrPortFrom(localAddr, m.serverPort)
+	m.dnsForwarder = NewDNSForwarder(listenAddress, dnsTTL, m.firewall, m.statusRecorder, m.wgIface)
+
 	go func() {
 		if err := m.dnsForwarder.Listen(fwdEntries); err != nil {
 			// todo handle close error if it is exists
@@ -111,12 +123,13 @@ func (m *Manager) Stop(ctx context.Context) error {
 
 	var mErr *multierror.Error
 
-	if m.localAddr.IsValid() && m.firewall != nil {
-		if err := m.firewall.RemoveInboundDNAT(m.localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
+	localAddr := m.wgIface.Address().IP
+	if localAddr.IsValid() && m.firewall != nil {
+		if err := m.firewall.RemoveInboundDNAT(localAddr, firewall.ProtocolUDP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf("remove DNS UDP DNAT rule: %w", err))
 		}
 
-		if err := m.firewall.RemoveInboundDNAT(m.localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
+		if err := m.firewall.RemoveInboundDNAT(localAddr, firewall.ProtocolTCP, nbdns.ForwarderClientPort, m.serverPort); err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf("remove DNS TCP DNAT rule: %w", err))
 		}
 	}

client/internal/engine.go

@@ -1855,35 +1855,69 @@ func (e *Engine) updateDNSForwarder(
 	}
 
 	if !enabled {
-		if e.dnsForwardMgr == nil {
-			return
-		}
-		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
-			log.Errorf("failed to stop DNS forward: %v", err)
-		}
+		e.stopDNSForwarder()
 		return
 	}
 
 	if len(fwdEntries) > 0 {
 		if e.dnsForwardMgr == nil {
-			localAddr := e.wgInterface.Address().IP
-			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder, localAddr)
-
-			if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
-				log.Errorf("failed to start DNS forward: %v", err)
-				e.dnsForwardMgr = nil
-			}
-
-			log.Infof("started domain router service with %d entries", len(fwdEntries))
+			e.startDNSForwarder(fwdEntries)
 		} else {
 			e.dnsForwardMgr.UpdateDomains(fwdEntries)
 		}
 	} else if e.dnsForwardMgr != nil {
 		log.Infof("disable domain router service")
-		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
-			log.Errorf("failed to stop DNS forward: %v", err)
-		}
+		e.stopDNSForwarder()
+	}
+}
+
+func (e *Engine) startDNSForwarder(fwdEntries []*dnsfwd.ForwarderEntry) {
+	e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder, e.wgInterface)
+	e.registerDNSServices()
+
+	if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
+		log.Errorf("failed to start DNS forward: %v", err)
 		e.dnsForwardMgr = nil
+		return
+	}
+
+	log.Infof("started domain router service with %d entries", len(fwdEntries))
+}
+
+func (e *Engine) stopDNSForwarder() {
+	if e.dnsForwardMgr == nil {
+		return
+	}
+
+	if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
+		log.Errorf("failed to stop DNS forward: %v", err)
+	}
+
+	e.unregisterDNSServices()
+	e.dnsForwardMgr = nil
+}
+
+func (e *Engine) registerDNSServices() {
+	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
+		if registrar, ok := e.firewall.(interface {
+			RegisterNetstackService(protocol nftypes.Protocol, port uint16)
+		}); ok {
+			registrar.RegisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
+			registrar.RegisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
+			log.Debugf("registered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
+		}
+	}
+}
+
+func (e *Engine) unregisterDNSServices() {
+	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
+		if registrar, ok := e.firewall.(interface {
+			UnregisterNetstackService(protocol nftypes.Protocol, port uint16)
+		}); ok {
+			registrar.UnregisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
+			registrar.UnregisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
+			log.Debugf("unregistered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
+		}
 	}
 }
 

client/net/conn.go

@@ -17,8 +17,7 @@ type Conn struct {
 	ID hooks.ConnectionID
 }
 
-// Close overrides the net.Conn Close method to execute all registered hooks after closing the connection
-// Close overrides the net.Conn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.Conn Close method to execute all registered hooks after closing the connection.
 func (c *Conn) Close() error {
 	return closeConn(c.ID, c.Conn)
 }
@@ -29,7 +28,7 @@ type TCPConn struct {
 	ID hooks.ConnectionID
 }
 
-// Close overrides the net.TCPConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.TCPConn Close method to execute all registered hooks after closing the connection.
 func (c *TCPConn) Close() error {
 	return closeConn(c.ID, c.TCPConn)
 }
@@ -37,13 +36,16 @@ func (c *TCPConn) Close() error {
 // closeConn is a helper function to close connections and execute close hooks.
 func closeConn(id hooks.ConnectionID, conn io.Closer) error {
 	err := conn.Close()
+	cleanupConnID(id)
+	return err
+}
 
+// cleanupConnID executes close hooks for a connection ID.
+func cleanupConnID(id hooks.ConnectionID) {
 	closeHooks := hooks.GetCloseHooks()
 	for _, hook := range closeHooks {
 		if err := hook(id); err != nil {
 			log.Errorf("Error executing close hook: %v", err)
 		}
 	}
-
-	return err
 }

client/net/dial.go

@@ -74,7 +74,6 @@ func DialTCP(network string, laddr, raddr *net.TCPAddr) (transport.TCPConn, erro
 		}
 		return &TCPConn{TCPConn: tcpConn, ID: c.ID}, nil
 	}
-
 	if err := conn.Close(); err != nil {
 		log.Errorf("failed to close connection: %v", err)
 	}

client/net/dialer_dial.go

@@ -30,6 +30,7 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.
 
 	conn, err := d.Dialer.DialContext(ctx, network, address)
 	if err != nil {
+		cleanupConnID(connID)
 		return nil, fmt.Errorf("d.Dialer.DialContext: %w", err)
 	}
 
@@ -64,7 +65,7 @@ func callDialerHooks(ctx context.Context, connID hooks.ConnectionID, address str
 
 	ips, err := resolver.LookupIPAddr(ctx, host)
 	if err != nil {
-		return fmt.Errorf("failed to resolve address %s: %w", address, err)
+		return fmt.Errorf("resolve address %s: %w", address, err)
 	}
 
 	log.Debugf("Dialer resolved IPs for %s: %v", address, ips)

client/net/listener_listen.go

@@ -48,7 +48,7 @@ func (c *PacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return c.PacketConn.WriteTo(b, addr)
 }
 
-// Close overrides the net.PacketConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.PacketConn Close method to execute all registered hooks after closing the connection.
 func (c *PacketConn) Close() error {
 	defer c.seenAddrs.Clear()
 	return closeConn(c.ID, c.PacketConn)
@@ -69,7 +69,7 @@ func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return c.UDPConn.WriteTo(b, addr)
 }
 
-// Close overrides the net.UDPConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.UDPConn Close method to execute all registered hooks after closing the connection.
 func (c *UDPConn) Close() error {
 	defer c.seenAddrs.Clear()
 	return closeConn(c.ID, c.UDPConn)

client/ui/client_ui.go

@@ -833,6 +833,7 @@ func (s *serviceClient) onTrayReady() {
 
 	newProfileMenuArgs := &newProfileMenuArgs{
 		ctx:                  s.ctx,
+		serviceClient:        s,
 		profileManager:       s.profileManager,
 		eventHandler:         s.eventHandler,
 		profileMenuItem:      profileMenuItem,

go.mod

@@ -56,6 +56,7 @@ require (
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/hashicorp/go-secure-stdlib/base62 v0.1.2
 	github.com/hashicorp/go-version v1.6.0
+	github.com/jackc/pgx/v5 v5.5.5
 	github.com/libdns/route53 v1.5.0
 	github.com/libp2p/go-netroute v0.2.1
 	github.com/mdlayher/socket v0.5.1
@@ -183,7 +184,6 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
-	github.com/jackc/pgx/v5 v5.5.5 // indirect
 	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/jeandeaual/go-locale v0.0.0-20240223122105-ce5225dcaa49 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect

management/server/store/sqlstore_bench_test.go

@@ -0,0 +1,951 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/netip"
+	"sort"
+	"sync"
+	"testing"
+	"time"
+
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
+	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
+	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/posture"
+	"github.com/netbirdio/netbird/management/server/testutil"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) GetAccountSlow(ctx context.Context, accountID string) (*types.Account, error) {
+	start := time.Now()
+	defer func() {
+		elapsed := time.Since(start)
+		if elapsed > 1*time.Second {
+			log.WithContext(ctx).Tracef("GetAccount for account %s exceeded 1s, took: %v", accountID, elapsed)
+		}
+	}()
+
+	var account types.Account
+	result := s.db.Model(&account).
+		Omit("GroupsG").
+		Preload("UsersG.PATsG"). // have to be specified as this is nested reference
+		Preload(clause.Associations).
+		Take(&account, idQueryCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("error when getting account %s from the store: %s", accountID, result.Error)
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAccountNotFoundError(accountID)
+		}
+		return nil, status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	// we have to manually preload policy rules as it seems that gorm preloading doesn't do it for us
+	for i, policy := range account.Policies {
+		var rules []*types.PolicyRule
+		err := s.db.Model(&types.PolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
+		if err != nil {
+			return nil, status.Errorf(status.NotFound, "rule not found")
+		}
+		account.Policies[i].Rules = rules
+	}
+
+	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
+	for _, key := range account.SetupKeysG {
+		account.SetupKeys[key.Key] = key.Copy()
+	}
+	account.SetupKeysG = nil
+
+	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
+	for _, peer := range account.PeersG {
+		account.Peers[peer.ID] = peer.Copy()
+	}
+	account.PeersG = nil
+
+	account.Users = make(map[string]*types.User, len(account.UsersG))
+	for _, user := range account.UsersG {
+		user.PATs = make(map[string]*types.PersonalAccessToken, len(user.PATs))
+		for _, pat := range user.PATsG {
+			user.PATs[pat.ID] = pat.Copy()
+		}
+		account.Users[user.Id] = user.Copy()
+	}
+	account.UsersG = nil
+
+	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
+	for _, group := range account.GroupsG {
+		account.Groups[group.ID] = group.Copy()
+	}
+	account.GroupsG = nil
+
+	var groupPeers []types.GroupPeer
+	s.db.Model(&types.GroupPeer{}).Where("account_id = ?", accountID).
+		Find(&groupPeers)
+	for _, groupPeer := range groupPeers {
+		if group, ok := account.Groups[groupPeer.GroupID]; ok {
+			group.Peers = append(group.Peers, groupPeer.PeerID)
+		} else {
+			log.WithContext(ctx).Warnf("group %s not found for group peer %s in account %s", groupPeer.GroupID, groupPeer.PeerID, accountID)
+		}
+	}
+
+	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
+	for _, route := range account.RoutesG {
+		account.Routes[route.ID] = route.Copy()
+	}
+	account.RoutesG = nil
+
+	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
+	for _, ns := range account.NameServerGroupsG {
+		account.NameServerGroups[ns.ID] = ns.Copy()
+	}
+	account.NameServerGroupsG = nil
+
+	return &account, nil
+}
+
+func (s *SqlStore) GetAccountGormOpt(ctx context.Context, accountID string) (*types.Account, error) {
+	start := time.Now()
+	defer func() {
+		elapsed := time.Since(start)
+		if elapsed > 1*time.Second {
+			log.WithContext(ctx).Tracef("GetAccount for account %s exceeded 1s, took: %v", accountID, elapsed)
+		}
+	}()
+
+	var account types.Account
+	result := s.db.Model(&account).
+		Preload("UsersG.PATsG"). // have to be specified as this is nested reference
+		Preload("Policies.Rules").
+		Preload("SetupKeysG").
+		Preload("PeersG").
+		Preload("UsersG").
+		Preload("GroupsG.GroupPeers").
+		Preload("RoutesG").
+		Preload("NameServerGroupsG").
+		Preload("PostureChecks").
+		Preload("Networks").
+		Preload("NetworkRouters").
+		Preload("NetworkResources").
+		Preload("Onboarding").
+		Take(&account, idQueryCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("error when getting account %s from the store: %s", accountID, result.Error)
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAccountNotFoundError(accountID)
+		}
+		return nil, status.NewGetAccountFromStoreError(result.Error)
+	}
+
+	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
+	for _, key := range account.SetupKeysG {
+		if key.UpdatedAt.IsZero() {
+			key.UpdatedAt = key.CreatedAt
+		}
+		if key.AutoGroups == nil {
+			key.AutoGroups = []string{}
+		}
+		account.SetupKeys[key.Key] = &key
+	}
+	account.SetupKeysG = nil
+
+	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
+	for _, peer := range account.PeersG {
+		account.Peers[peer.ID] = &peer
+	}
+	account.PeersG = nil
+	account.Users = make(map[string]*types.User, len(account.UsersG))
+	for _, user := range account.UsersG {
+		user.PATs = make(map[string]*types.PersonalAccessToken, len(user.PATs))
+		for _, pat := range user.PATsG {
+			pat.UserID = ""
+			user.PATs[pat.ID] = &pat
+		}
+		if user.AutoGroups == nil {
+			user.AutoGroups = []string{}
+		}
+		account.Users[user.Id] = &user
+		user.PATsG = nil
+	}
+	account.UsersG = nil
+	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
+	for _, group := range account.GroupsG {
+		group.Peers = make([]string, len(group.GroupPeers))
+		for i, gp := range group.GroupPeers {
+			group.Peers[i] = gp.PeerID
+		}
+		if group.Resources == nil {
+			group.Resources = []types.Resource{}
+		}
+		account.Groups[group.ID] = group
+	}
+	account.GroupsG = nil
+
+	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
+	for _, route := range account.RoutesG {
+		account.Routes[route.ID] = &route
+	}
+	account.RoutesG = nil
+	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
+	for _, ns := range account.NameServerGroupsG {
+		ns.AccountID = ""
+		if ns.NameServers == nil {
+			ns.NameServers = []nbdns.NameServer{}
+		}
+		if ns.Groups == nil {
+			ns.Groups = []string{}
+		}
+		if ns.Domains == nil {
+			ns.Domains = []string{}
+		}
+		account.NameServerGroups[ns.ID] = &ns
+	}
+	account.NameServerGroupsG = nil
+	return &account, nil
+}
+
+func connectDBforTest(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
+	config, err := pgxpool.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse database config: %w", err)
+	}
+
+	config.MaxConns = 12
+	config.MinConns = 2
+	config.MaxConnLifetime = time.Hour
+	config.HealthCheckPeriod = time.Minute
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create connection pool: %w", err)
+	}
+
+	if err := pool.Ping(ctx); err != nil {
+		pool.Close()
+		return nil, fmt.Errorf("unable to ping database: %w", err)
+	}
+	return pool, nil
+}
+
+func setupBenchmarkDB(b testing.TB) (*SqlStore, func(), string) {
+	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+	if err != nil {
+		b.Fatalf("failed to create test container: %v", err)
+	}
+
+	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
+	if err != nil {
+		b.Fatalf("failed to connect database: %v", err)
+	}
+
+	pool, err := connectDBforTest(context.Background(), dsn)
+	if err != nil {
+		b.Fatalf("failed to connect database: %v", err)
+	}
+
+	models := []interface{}{
+		&types.Account{}, &types.SetupKey{}, &nbpeer.Peer{}, &types.User{},
+		&types.PersonalAccessToken{}, &types.Group{}, &types.GroupPeer{},
+		&types.Policy{}, &types.PolicyRule{}, &route.Route{},
+		&nbdns.NameServerGroup{}, &posture.Checks{}, &networkTypes.Network{},
+		&routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{},
+		&types.AccountOnboarding{},
+	}
+
+	for i := len(models) - 1; i >= 0; i-- {
+		err := db.Migrator().DropTable(models[i])
+		if err != nil {
+			b.Fatalf("failed to drop table: %v", err)
+		}
+	}
+
+	err = db.AutoMigrate(models...)
+	if err != nil {
+		b.Fatalf("failed to migrate database: %v", err)
+	}
+
+	store := &SqlStore{
+		db:   db,
+		pool: pool,
+	}
+
+	const (
+		accountID           = "benchmark-account-id"
+		numUsers            = 20
+		numPatsPerUser      = 3
+		numSetupKeys        = 25
+		numPeers            = 200
+		numGroups           = 30
+		numPolicies         = 50
+		numRulesPerPolicy   = 10
+		numRoutes           = 40
+		numNSGroups         = 10
+		numPostureChecks    = 15
+		numNetworks         = 5
+		numNetworkRouters   = 5
+		numNetworkResources = 10
+	)
+
+	_, ipNet, _ := net.ParseCIDR("100.64.0.0/10")
+	acc := types.Account{
+		Id:                     accountID,
+		CreatedBy:              "benchmark-user",
+		CreatedAt:              time.Now(),
+		Domain:                 "benchmark.com",
+		IsDomainPrimaryAccount: true,
+		Network: &types.Network{
+			Identifier: "benchmark-net",
+			Net:        *ipNet,
+			Serial:     1,
+		},
+		DNSSettings: types.DNSSettings{
+			DisabledManagementGroups: []string{"group-disabled-1"},
+		},
+		Settings: &types.Settings{},
+	}
+	if err := db.Create(&acc).Error; err != nil {
+		b.Fatalf("create account: %v", err)
+	}
+
+	var setupKeys []types.SetupKey
+	for i := 0; i < numSetupKeys; i++ {
+		setupKeys = append(setupKeys, types.SetupKey{
+			Id:        fmt.Sprintf("keyid-%d", i),
+			AccountID: accountID,
+			Key:       fmt.Sprintf("key-%d", i),
+			Name:      fmt.Sprintf("Benchmark Key %d", i),
+			ExpiresAt: &time.Time{},
+		})
+	}
+	if err := db.Create(&setupKeys).Error; err != nil {
+		b.Fatalf("create setup keys: %v", err)
+	}
+
+	var peers []nbpeer.Peer
+	for i := 0; i < numPeers; i++ {
+		peers = append(peers, nbpeer.Peer{
+			ID:        fmt.Sprintf("peer-%d", i),
+			AccountID: accountID,
+			Key:       fmt.Sprintf("peerkey-%d", i),
+			IP:        net.ParseIP(fmt.Sprintf("100.64.0.%d", i+1)),
+			Name:      fmt.Sprintf("peer-name-%d", i),
+			Status:    &nbpeer.PeerStatus{Connected: i%2 == 0, LastSeen: time.Now()},
+		})
+	}
+	if err := db.Create(&peers).Error; err != nil {
+		b.Fatalf("create peers: %v", err)
+	}
+
+	for i := 0; i < numUsers; i++ {
+		userID := fmt.Sprintf("user-%d", i)
+		user := types.User{Id: userID, AccountID: accountID}
+		if err := db.Create(&user).Error; err != nil {
+			b.Fatalf("create user %s: %v", userID, err)
+		}
+
+		var pats []types.PersonalAccessToken
+		for j := 0; j < numPatsPerUser; j++ {
+			pats = append(pats, types.PersonalAccessToken{
+				ID:     fmt.Sprintf("pat-%d-%d", i, j),
+				UserID: userID,
+				Name:   fmt.Sprintf("PAT %d for User %d", j, i),
+			})
+		}
+		if err := db.Create(&pats).Error; err != nil {
+			b.Fatalf("create pats for user %s: %v", userID, err)
+		}
+	}
+
+	var groups []*types.Group
+	for i := 0; i < numGroups; i++ {
+		groups = append(groups, &types.Group{
+			ID:        fmt.Sprintf("group-%d", i),
+			AccountID: accountID,
+			Name:      fmt.Sprintf("Group %d", i),
+		})
+	}
+	if err := db.Create(&groups).Error; err != nil {
+		b.Fatalf("create groups: %v", err)
+	}
+
+	for i := 0; i < numPolicies; i++ {
+		policyID := fmt.Sprintf("policy-%d", i)
+		policy := types.Policy{ID: policyID, AccountID: accountID, Name: fmt.Sprintf("Policy %d", i), Enabled: true}
+		if err := db.Create(&policy).Error; err != nil {
+			b.Fatalf("create policy %s: %v", policyID, err)
+		}
+
+		var rules []*types.PolicyRule
+		for j := 0; j < numRulesPerPolicy; j++ {
+			rules = append(rules, &types.PolicyRule{
+				ID:       fmt.Sprintf("rule-%d-%d", i, j),
+				PolicyID: policyID,
+				Name:     fmt.Sprintf("Rule %d for Policy %d", j, i),
+				Enabled:  true,
+				Protocol: "all",
+			})
+		}
+		if err := db.Create(&rules).Error; err != nil {
+			b.Fatalf("create rules for policy %s: %v", policyID, err)
+		}
+	}
+
+	var routes []route.Route
+	for i := 0; i < numRoutes; i++ {
+		routes = append(routes, route.Route{
+			ID:          route.ID(fmt.Sprintf("route-%d", i)),
+			AccountID:   accountID,
+			Description: fmt.Sprintf("Route %d", i),
+			Network:     netip.MustParsePrefix(fmt.Sprintf("192.168.%d.0/24", i)),
+			Enabled:     true,
+		})
+	}
+	if err := db.Create(&routes).Error; err != nil {
+		b.Fatalf("create routes: %v", err)
+	}
+
+	var nsGroups []nbdns.NameServerGroup
+	for i := 0; i < numNSGroups; i++ {
+		nsGroups = append(nsGroups, nbdns.NameServerGroup{
+			ID:          fmt.Sprintf("nsg-%d", i),
+			AccountID:   accountID,
+			Name:        fmt.Sprintf("NS Group %d", i),
+			Description: "Benchmark NS Group",
+			Enabled:     true,
+		})
+	}
+	if err := db.Create(&nsGroups).Error; err != nil {
+		b.Fatalf("create nsgroups: %v", err)
+	}
+
+	var postureChecks []*posture.Checks
+	for i := 0; i < numPostureChecks; i++ {
+		postureChecks = append(postureChecks, &posture.Checks{
+			ID:        fmt.Sprintf("pc-%d", i),
+			AccountID: accountID,
+			Name:      fmt.Sprintf("Posture Check %d", i),
+		})
+	}
+	if err := db.Create(&postureChecks).Error; err != nil {
+		b.Fatalf("create posture checks: %v", err)
+	}
+
+	var networks []*networkTypes.Network
+	for i := 0; i < numNetworks; i++ {
+		networks = append(networks, &networkTypes.Network{
+			ID:        fmt.Sprintf("nettype-%d", i),
+			AccountID: accountID,
+			Name:      fmt.Sprintf("Network Type %d", i),
+		})
+	}
+	if err := db.Create(&networks).Error; err != nil {
+		b.Fatalf("create networks: %v", err)
+	}
+
+	var networkRouters []*routerTypes.NetworkRouter
+	for i := 0; i < numNetworkRouters; i++ {
+		networkRouters = append(networkRouters, &routerTypes.NetworkRouter{
+			ID:        fmt.Sprintf("router-%d", i),
+			AccountID: accountID,
+			NetworkID: networks[i%numNetworks].ID,
+			Peer:      peers[i%numPeers].ID,
+		})
+	}
+	if err := db.Create(&networkRouters).Error; err != nil {
+		b.Fatalf("create network routers: %v", err)
+	}
+
+	var networkResources []*resourceTypes.NetworkResource
+	for i := 0; i < numNetworkResources; i++ {
+		networkResources = append(networkResources, &resourceTypes.NetworkResource{
+			ID:        fmt.Sprintf("resource-%d", i),
+			AccountID: accountID,
+			NetworkID: networks[i%numNetworks].ID,
+			Name:      fmt.Sprintf("Resource %d", i),
+		})
+	}
+	if err := db.Create(&networkResources).Error; err != nil {
+		b.Fatalf("create network resources: %v", err)
+	}
+
+	onboarding := types.AccountOnboarding{
+		AccountID:             accountID,
+		OnboardingFlowPending: true,
+	}
+	if err := db.Create(&onboarding).Error; err != nil {
+		b.Fatalf("create onboarding: %v", err)
+	}
+
+	return store, cleanup, accountID
+}
+
+func BenchmarkGetAccount(b *testing.B) {
+	store, cleanup, accountID := setupBenchmarkDB(b)
+	defer cleanup()
+	ctx := context.Background()
+	b.ResetTimer()
+	b.ReportAllocs()
+	b.Run("old", func(b *testing.B) {
+		for range b.N {
+			_, err := store.GetAccountSlow(ctx, accountID)
+			if err != nil {
+				b.Fatalf("GetAccountSlow failed: %v", err)
+			}
+		}
+	})
+	b.Run("gorm opt", func(b *testing.B) {
+		for range b.N {
+			_, err := store.GetAccountGormOpt(ctx, accountID)
+			if err != nil {
+				b.Fatalf("GetAccountFast failed: %v", err)
+			}
+		}
+	})
+	b.Run("raw", func(b *testing.B) {
+		for range b.N {
+			_, err := store.GetAccount(ctx, accountID)
+			if err != nil {
+				b.Fatalf("GetAccountPureSQL failed: %v", err)
+			}
+		}
+	})
+	store.pool.Close()
+}
+
+func TestAccountEquivalence(t *testing.T) {
+	store, cleanup, accountID := setupBenchmarkDB(t)
+	defer cleanup()
+	ctx := context.Background()
+
+	type getAccountFunc func(context.Context, string) (*types.Account, error)
+
+	tests := []struct {
+		name      string
+		expectedF getAccountFunc
+		actualF   getAccountFunc
+	}{
+		{"old vs new", store.GetAccountSlow, store.GetAccountGormOpt},
+		{"old vs raw", store.GetAccountSlow, store.GetAccount},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			expected, errOld := tt.expectedF(ctx, accountID)
+			assert.NoError(t, errOld, "expected function should not return an error")
+			assert.NotNil(t, expected, "expected should not be nil")
+
+			actual, errNew := tt.actualF(ctx, accountID)
+			assert.NoError(t, errNew, "actual function should not return an error")
+			assert.NotNil(t, actual, "actual should not be nil")
+			testAccountEquivalence(t, expected, actual)
+		})
+	}
+
+	expected, errOld := store.GetAccountSlow(ctx, accountID)
+	assert.NoError(t, errOld, "GetAccountSlow should not return an error")
+	assert.NotNil(t, expected, "expected should not be nil")
+
+	actual, errNew := store.GetAccount(ctx, accountID)
+	assert.NoError(t, errNew, "GetAccount (new) should not return an error")
+	assert.NotNil(t, actual, "actual should not be nil")
+}
+
+func testAccountEquivalence(t *testing.T, expected, actual *types.Account) {
+	assert.Equal(t, expected.Id, actual.Id, "Account IDs should be equal")
+	assert.Equal(t, expected.CreatedBy, actual.CreatedBy, "Account CreatedBy fields should be equal")
+	assert.WithinDuration(t, expected.CreatedAt, actual.CreatedAt, time.Second, "Account CreatedAt timestamps should be within a second")
+	assert.Equal(t, expected.Domain, actual.Domain, "Account Domains should be equal")
+	assert.Equal(t, expected.DomainCategory, actual.DomainCategory, "Account DomainCategories should be equal")
+	assert.Equal(t, expected.IsDomainPrimaryAccount, actual.IsDomainPrimaryAccount, "Account IsDomainPrimaryAccount flags should be equal")
+	assert.Equal(t, expected.Network, actual.Network, "Embedded Account Network structs should be equal")
+	assert.Equal(t, expected.DNSSettings, actual.DNSSettings, "Embedded Account DNSSettings structs should be equal")
+	assert.Equal(t, expected.Onboarding, actual.Onboarding, "Embedded Account Onboarding structs should be equal")
+
+	assert.Len(t, actual.SetupKeys, len(expected.SetupKeys), "SetupKeys maps should have the same number of elements")
+	for key, oldVal := range expected.SetupKeys {
+		newVal, ok := actual.SetupKeys[key]
+		assert.True(t, ok, "SetupKey with key '%s' should exist in new account", key)
+		assert.Equal(t, *oldVal, *newVal, "SetupKey with key '%s' should be equal", key)
+	}
+
+	assert.Len(t, actual.Peers, len(expected.Peers), "Peers maps should have the same number of elements")
+	for key, oldVal := range expected.Peers {
+		newVal, ok := actual.Peers[key]
+		assert.True(t, ok, "Peer with ID '%s' should exist in new account", key)
+		assert.Equal(t, *oldVal, *newVal, "Peer with ID '%s' should be equal", key)
+	}
+
+	assert.Len(t, actual.Users, len(expected.Users), "Users maps should have the same number of elements")
+	for key, oldUser := range expected.Users {
+		newUser, ok := actual.Users[key]
+		assert.True(t, ok, "User with ID '%s' should exist in new account", key)
+
+		assert.Len(t, newUser.PATs, len(oldUser.PATs), "PATs map for user '%s' should have the same size", key)
+		for patKey, oldPAT := range oldUser.PATs {
+			newPAT, patOk := newUser.PATs[patKey]
+			assert.True(t, patOk, "PAT with ID '%s' for user '%s' should exist in new user object", patKey, key)
+			assert.Equal(t, *oldPAT, *newPAT, "PAT with ID '%s' for user '%s' should be equal", patKey, key)
+		}
+
+		oldUser.PATs = nil
+		newUser.PATs = nil
+		assert.Equal(t, *oldUser, *newUser, "User struct for ID '%s' (without PATs) should be equal", key)
+	}
+
+	assert.Len(t, actual.Groups, len(expected.Groups), "Groups maps should have the same number of elements")
+	for key, oldVal := range expected.Groups {
+		newVal, ok := actual.Groups[key]
+		assert.True(t, ok, "Group with ID '%s' should exist in new account", key)
+		sort.Strings(oldVal.Peers)
+		sort.Strings(newVal.Peers)
+		assert.Equal(t, *oldVal, *newVal, "Group with ID '%s' should be equal", key)
+	}
+
+	assert.Len(t, actual.Routes, len(expected.Routes), "Routes maps should have the same number of elements")
+	for key, oldVal := range expected.Routes {
+		newVal, ok := actual.Routes[key]
+		assert.True(t, ok, "Route with ID '%s' should exist in new account", key)
+		assert.Equal(t, *oldVal, *newVal, "Route with ID '%s' should be equal", key)
+	}
+
+	assert.Len(t, actual.NameServerGroups, len(expected.NameServerGroups), "NameServerGroups maps should have the same number of elements")
+	for key, oldVal := range expected.NameServerGroups {
+		newVal, ok := actual.NameServerGroups[key]
+		assert.True(t, ok, "NameServerGroup with ID '%s' should exist in new account", key)
+		assert.Equal(t, *oldVal, *newVal, "NameServerGroup with ID '%s' should be equal", key)
+	}
+
+	assert.Len(t, actual.Policies, len(expected.Policies), "Policies slices should have the same number of elements")
+	sort.Slice(expected.Policies, func(i, j int) bool { return expected.Policies[i].ID < expected.Policies[j].ID })
+	sort.Slice(actual.Policies, func(i, j int) bool { return actual.Policies[i].ID < actual.Policies[j].ID })
+	for i := range expected.Policies {
+		sort.Slice(expected.Policies[i].Rules, func(j, k int) bool { return expected.Policies[i].Rules[j].ID < expected.Policies[i].Rules[k].ID })
+		sort.Slice(actual.Policies[i].Rules, func(j, k int) bool { return actual.Policies[i].Rules[j].ID < actual.Policies[i].Rules[k].ID })
+		assert.Equal(t, *expected.Policies[i], *actual.Policies[i], "Policy with ID '%s' should be equal", expected.Policies[i].ID)
+	}
+
+	assert.Len(t, actual.PostureChecks, len(expected.PostureChecks), "PostureChecks slices should have the same number of elements")
+	sort.Slice(expected.PostureChecks, func(i, j int) bool { return expected.PostureChecks[i].ID < expected.PostureChecks[j].ID })
+	sort.Slice(actual.PostureChecks, func(i, j int) bool { return actual.PostureChecks[i].ID < actual.PostureChecks[j].ID })
+	for i := range expected.PostureChecks {
+		assert.Equal(t, *expected.PostureChecks[i], *actual.PostureChecks[i], "PostureCheck with ID '%s' should be equal", expected.PostureChecks[i].ID)
+	}
+
+	assert.Len(t, actual.Networks, len(expected.Networks), "Networks slices should have the same number of elements")
+	sort.Slice(expected.Networks, func(i, j int) bool { return expected.Networks[i].ID < expected.Networks[j].ID })
+	sort.Slice(actual.Networks, func(i, j int) bool { return actual.Networks[i].ID < actual.Networks[j].ID })
+	for i := range expected.Networks {
+		assert.Equal(t, *expected.Networks[i], *actual.Networks[i], "Network with ID '%s' should be equal", expected.Networks[i].ID)
+	}
+
+	assert.Len(t, actual.NetworkRouters, len(expected.NetworkRouters), "NetworkRouters slices should have the same number of elements")
+	sort.Slice(expected.NetworkRouters, func(i, j int) bool { return expected.NetworkRouters[i].ID < expected.NetworkRouters[j].ID })
+	sort.Slice(actual.NetworkRouters, func(i, j int) bool { return actual.NetworkRouters[i].ID < actual.NetworkRouters[j].ID })
+	for i := range expected.NetworkRouters {
+		assert.Equal(t, *expected.NetworkRouters[i], *actual.NetworkRouters[i], "NetworkRouter with ID '%s' should be equal", expected.NetworkRouters[i].ID)
+	}
+
+	assert.Len(t, actual.NetworkResources, len(expected.NetworkResources), "NetworkResources slices should have the same number of elements")
+	sort.Slice(expected.NetworkResources, func(i, j int) bool { return expected.NetworkResources[i].ID < expected.NetworkResources[j].ID })
+	sort.Slice(actual.NetworkResources, func(i, j int) bool { return actual.NetworkResources[i].ID < actual.NetworkResources[j].ID })
+	for i := range expected.NetworkResources {
+		assert.Equal(t, *expected.NetworkResources[i], *actual.NetworkResources[i], "NetworkResource with ID '%s' should be equal", expected.NetworkResources[i].ID)
+	}
+}
+
+func (s *SqlStore) GetAccountPureSQL(ctx context.Context, accountID string) (*types.Account, error) {
+	account, err := s.getAccount(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	var wg sync.WaitGroup
+	errChan := make(chan error, 12)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		keys, err := s.getSetupKeys(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.SetupKeysG = keys
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		peers, err := s.getPeers(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.PeersG = peers
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		users, err := s.getUsers(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.UsersG = users
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		groups, err := s.getGroups(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.GroupsG = groups
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		policies, err := s.getPolicies(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.Policies = policies
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		routes, err := s.getRoutes(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.RoutesG = routes
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		nsgs, err := s.getNameServerGroups(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.NameServerGroupsG = nsgs
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		checks, err := s.getPostureChecks(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.PostureChecks = checks
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		networks, err := s.getNetworks(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.Networks = networks
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		routers, err := s.getNetworkRouters(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.NetworkRouters = routers
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		resources, err := s.getNetworkResources(ctx, accountID)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		account.NetworkResources = resources
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err := s.getAccountOnboarding(ctx, accountID, account)
+		if err != nil {
+			errChan <- err
+			return
+		}
+	}()
+
+	wg.Wait()
+	close(errChan)
+	for e := range errChan {
+		if e != nil {
+			return nil, e
+		}
+	}
+
+	var userIDs []string
+	for _, u := range account.UsersG {
+		userIDs = append(userIDs, u.Id)
+	}
+	var policyIDs []string
+	for _, p := range account.Policies {
+		policyIDs = append(policyIDs, p.ID)
+	}
+	var groupIDs []string
+	for _, g := range account.GroupsG {
+		groupIDs = append(groupIDs, g.ID)
+	}
+
+	wg.Add(3)
+	errChan = make(chan error, 3)
+
+	var pats []types.PersonalAccessToken
+	go func() {
+		defer wg.Done()
+		var err error
+		pats, err = s.getPersonalAccessTokens(ctx, userIDs)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	var rules []*types.PolicyRule
+	go func() {
+		defer wg.Done()
+		var err error
+		rules, err = s.getPolicyRules(ctx, policyIDs)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	var groupPeers []types.GroupPeer
+	go func() {
+		defer wg.Done()
+		var err error
+		groupPeers, err = s.getGroupPeers(ctx, groupIDs)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	wg.Wait()
+	close(errChan)
+	for e := range errChan {
+		if e != nil {
+			return nil, e
+		}
+	}
+
+	patsByUserID := make(map[string][]*types.PersonalAccessToken)
+	for i := range pats {
+		pat := &pats[i]
+		patsByUserID[pat.UserID] = append(patsByUserID[pat.UserID], pat)
+		pat.UserID = ""
+	}
+
+	rulesByPolicyID := make(map[string][]*types.PolicyRule)
+	for _, rule := range rules {
+		rulesByPolicyID[rule.PolicyID] = append(rulesByPolicyID[rule.PolicyID], rule)
+	}
+
+	peersByGroupID := make(map[string][]string)
+	for _, gp := range groupPeers {
+		peersByGroupID[gp.GroupID] = append(peersByGroupID[gp.GroupID], gp.PeerID)
+	}
+
+	account.SetupKeys = make(map[string]*types.SetupKey, len(account.SetupKeysG))
+	for i := range account.SetupKeysG {
+		key := &account.SetupKeysG[i]
+		account.SetupKeys[key.Key] = key
+	}
+
+	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
+	for i := range account.PeersG {
+		peer := &account.PeersG[i]
+		account.Peers[peer.ID] = peer
+	}
+
+	account.Users = make(map[string]*types.User, len(account.UsersG))
+	for i := range account.UsersG {
+		user := &account.UsersG[i]
+		user.PATs = make(map[string]*types.PersonalAccessToken)
+		if userPats, ok := patsByUserID[user.Id]; ok {
+			for j := range userPats {
+				pat := userPats[j]
+				user.PATs[pat.ID] = pat
+			}
+		}
+		account.Users[user.Id] = user
+	}
+
+	for i := range account.Policies {
+		policy := account.Policies[i]
+		if policyRules, ok := rulesByPolicyID[policy.ID]; ok {
+			policy.Rules = policyRules
+		}
+	}
+
+	account.Groups = make(map[string]*types.Group, len(account.GroupsG))
+	for i := range account.GroupsG {
+		group := account.GroupsG[i]
+		if peerIDs, ok := peersByGroupID[group.ID]; ok {
+			group.Peers = peerIDs
+		}
+		account.Groups[group.ID] = group
+	}
+
+	account.Routes = make(map[route.ID]*route.Route, len(account.RoutesG))
+	for i := range account.RoutesG {
+		route := &account.RoutesG[i]
+		account.Routes[route.ID] = route
+	}
+
+	account.NameServerGroups = make(map[string]*nbdns.NameServerGroup, len(account.NameServerGroupsG))
+	for i := range account.NameServerGroupsG {
+		nsg := &account.NameServerGroupsG[i]
+		nsg.AccountID = ""
+		account.NameServerGroups[nsg.ID] = nsg
+	}
+
+	account.SetupKeysG = nil
+	account.PeersG = nil
+	account.UsersG = nil
+	account.GroupsG = nil
+	account.RoutesG = nil
+	account.NameServerGroupsG = nil
+
+	return account, nil
+}

management/server/store/store.go

@@ -468,6 +468,9 @@ func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine)
 	closeConnection := func() {
 		cleanup()
 		store.Close(ctx)
+		if store.pool != nil {
+			store.pool.Close()
+		}
 	}
 
 	return store, closeConnection, nil

route/route.go

@@ -124,6 +124,7 @@ func (r *Route) EventMeta() map[string]any {
 func (r *Route) Copy() *Route {
 	route := &Route{
 		ID:                  r.ID,
+		AccountID:           r.AccountID,
 		Description:         r.Description,
 		NetID:               r.NetID,
 		Network:             r.Network,

shared/management/client/grpc.go

@@ -55,8 +55,7 @@ func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsE
 		var err error
 		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.ManagementComponent)
 		if err != nil {
-			log.Printf("createConnection error: %v", err)
-			return err
+			return fmt.Errorf("create connection: %w", err)
 		}
 		return nil
 	}

shared/signal/client/grpc.go

@@ -60,8 +60,7 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 		var err error
 		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.SignalComponent)
 		if err != nil {
-			log.Printf("createConnection error: %v", err)
-			return err
+			return fmt.Errorf("create connection: %w", err)
 		}
 		return nil
 	}