From 588735f2bcfe33324ce2e9b8313a4c9b41721e97 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 14 Jul 2026 19:29:22 +0200 Subject: [PATCH] [client] warm lazy connections from the DNS resolver --- client/internal/dns/local/local.go | 66 +++++++++++++++++++ client/internal/dns/local/warmup_test.go | 81 ++++++++++++++++++++++++ client/internal/dns/server.go | 7 ++ client/internal/dns_peer_activator.go | 71 +++++++++++++++++++++ client/internal/engine.go | 12 ++++ 5 files changed, 237 insertions(+) create mode 100644 client/internal/dns/local/warmup_test.go create mode 100644 client/internal/dns_peer_activator.go diff --git a/client/internal/dns/local/local.go b/client/internal/dns/local/local.go index d0268186c..667bd0cc2 100644 --- a/client/internal/dns/local/local.go +++ b/client/internal/dns/local/local.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "net/netip" + "os" "slices" "strings" "sync" @@ -39,6 +40,30 @@ type PeerConnectivity interface { IsConnectedByIP(ip string) (known, connected bool) } +// PeerActivator wakes lazy-connection peers on demand. The local resolver calls +// it with the tunnel IPs an answer points at, so a peer that is idle (lazily +// disconnected) begins its WireGuard handshake at DNS-resolution time rather +// than racing the client's first request packet. nil disables warm-up. +type PeerActivator interface { + // ActivatePeersByIP triggers wake-up for the peer(s) owning ips and blocks + // until one is connected or ctx (a short per-query budget) expires. It is a + // fast no-op for unknown or already-connected IPs. + ActivatePeersByIP(ctx context.Context, ips []string) +} + +const defaultLazyWarmupTimeout = 2 * time.Second + +// lazyWarmupTimeout is the per-query budget for waking a lazy-connection peer a +// DNS answer points at. Tunable via NB_DNS_LAZY_WARMUP_TIMEOUT (a Go duration). +func lazyWarmupTimeout() time.Duration { + if v := os.Getenv("NB_DNS_LAZY_WARMUP_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + return d + } + } + return defaultLazyWarmupTimeout +} + type Resolver struct { mu sync.RWMutex records map[dns.Question][]dns.RR @@ -51,6 +76,9 @@ type Resolver struct { // filter and preserves the legacy "return whatever is registered" // behaviour for callers that never wire a status source. peerConn PeerConnectivity + // peerActivator, when non-nil, is called at resolution time to warm the + // lazy connection to the peer(s) an answer points at. nil disables warm-up. + peerActivator PeerActivator ctx context.Context cancel context.CancelFunc @@ -76,6 +104,14 @@ func (d *Resolver) SetPeerConnectivity(p PeerConnectivity) { d.peerConn = p } +// SetPeerActivator wires the DNS-time lazy-connection warm-up. Pass nil to +// disable. Safe to call multiple times; the latest value wins. +func (d *Resolver) SetPeerActivator(a PeerActivator) { + d.mu.Lock() + defer d.mu.Unlock() + d.peerActivator = a +} + func (d *Resolver) MatchSubdomains() bool { return true } @@ -122,6 +158,9 @@ func (d *Resolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) { replyMessage.RecursionAvailable = true result := d.lookupRecords(logger, question) + // Warm before filtering: activation flips a lazily-idle target to connected, + // which then lets it survive the disconnected-peer filter below. + d.warmLazyPeers(result.records) result.records = d.filterDisconnectedPeerAnswers(logger, question, result.records) replyMessage.Authoritative = !result.hasExternalData replyMessage.Answer = result.records @@ -518,6 +557,33 @@ func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns return kept } +// warmLazyPeers triggers lazy-connection wake-up for the peers a resolved +// answer points at and waits briefly for one to connect, so the caller's first +// request doesn't race the WireGuard handshake. No-op when no activator is +// wired (lazy connections disabled) or the answer carries no peer IPs. +func (d *Resolver) warmLazyPeers(records []dns.RR) { + d.mu.RLock() + activator := d.peerActivator + d.mu.RUnlock() + if activator == nil { + return + } + + var ips []string + for _, rr := range records { + if ip := extractRecordIP(rr); ip != "" { + ips = append(ips, ip) + } + } + if len(ips) == 0 { + return + } + + ctx, cancel := context.WithTimeout(d.ctx, lazyWarmupTimeout()) + defer cancel() + activator.ActivatePeersByIP(ctx, ips) +} + // extractRecordIP returns the dotted-decimal / colon-hex IP carried by // an A or AAAA record, or "" for any other record type. func extractRecordIP(rr dns.RR) string { diff --git a/client/internal/dns/local/warmup_test.go b/client/internal/dns/local/warmup_test.go new file mode 100644 index 000000000..5c3408ec4 --- /dev/null +++ b/client/internal/dns/local/warmup_test.go @@ -0,0 +1,81 @@ +package local + +import ( + "context" + "sync" + "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" +) + +// recordingActivator records the IPs it was asked to warm and returns +// immediately, so ServeDNS is not blocked by the test. +type recordingActivator struct { + mu sync.Mutex + called bool + ips []string +} + +func (r *recordingActivator) ActivatePeersByIP(_ context.Context, ips []string) { + r.mu.Lock() + defer r.mu.Unlock() + r.called = true + r.ips = append(r.ips, ips...) +} + +func serveA(t *testing.T, resolver *Resolver, name string) *dns.Msg { + t.Helper() + var resp *dns.Msg + w := &test.MockResponseWriter{WriteMsgFunc: func(m *dns.Msg) error { resp = m; return nil }} + resolver.ServeDNS(w, new(dns.Msg).SetQuestion(name, dns.TypeA)) + return resp +} + +func TestLocalResolver_WarmsLazyPeerOnResolve(t *testing.T) { + rec := nbdns.SimpleRecord{Name: "svc.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"} + resolver := NewResolver() + require.NoError(t, resolver.RegisterRecord(rec)) + + act := &recordingActivator{} + resolver.SetPeerActivator(act) + + resp := serveA(t, resolver, rec.Name) + require.NotNil(t, resp, "resolver must answer") + require.NotEmpty(t, resp.Answer, "answer must carry the A record") + + act.mu.Lock() + defer act.mu.Unlock() + assert.True(t, act.called, "activator must be invoked for an overlay A answer") + assert.Contains(t, act.ips, "100.64.0.7", "activator must receive the answer's peer IP") +} + +func TestLocalResolver_NoActivatorNoWarmup(t *testing.T) { + // With no activator wired the resolver behaves exactly as before. + rec := nbdns.SimpleRecord{Name: "svc.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"} + resolver := NewResolver() + require.NoError(t, resolver.RegisterRecord(rec)) + + resp := serveA(t, resolver, rec.Name) + require.NotNil(t, resp, "resolver must still answer without an activator") + require.NotEmpty(t, resp.Answer, "answer must carry the A record") +} + +func TestLocalResolver_NoWarmupForMissingRecord(t *testing.T) { + // A query that resolves to nothing must not invoke the activator (no IPs). + resolver := NewResolver() + require.NoError(t, resolver.RegisterRecord(nbdns.SimpleRecord{Name: "svc.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"})) + + act := &recordingActivator{} + resolver.SetPeerActivator(act) + + serveA(t, resolver, "absent.netbird.cloud.") + + act.mu.Lock() + defer act.mu.Unlock() + assert.False(t, act.called, "activator must not be invoked when there is no answer") +} diff --git a/client/internal/dns/server.go b/client/internal/dns/server.go index 7556c66cc..c4501c3e0 100644 --- a/client/internal/dns/server.go +++ b/client/internal/dns/server.go @@ -491,6 +491,13 @@ func (s *DefaultServer) SetFirewall(fw Firewall) { } } +// SetPeerActivator wires the DNS-time lazy-connection warm-up on the local +// resolver. Injected after the connection manager exists (it does not at +// DNS-server construction time). Pass nil to disable. +func (s *DefaultServer) SetPeerActivator(a local.PeerActivator) { + s.localResolver.SetPeerActivator(a) +} + // Stop stops the server func (s *DefaultServer) Stop() { s.ctxCancel() diff --git a/client/internal/dns_peer_activator.go b/client/internal/dns_peer_activator.go new file mode 100644 index 000000000..2e222f783 --- /dev/null +++ b/client/internal/dns_peer_activator.go @@ -0,0 +1,71 @@ +package internal + +import ( + "context" + "sync" + "time" + + "github.com/netbirdio/netbird/client/internal/peer" + "github.com/netbirdio/netbird/client/internal/peerstore" +) + +const dnsActivationPollInterval = 50 * time.Millisecond + +// dnsPeerActivator wakes lazy-connection peers from the DNS resolution path. It +// implements dns/local.PeerActivator. ConnMgr is not thread-safe (guarded by +// the engine's syncMsgMux) while DNS queries run on their own goroutines, so +// activation runs under that mutex; the connection wait runs without it. +type dnsPeerActivator struct { + connMgr *ConnMgr + peerStore *peerstore.Store + status *peer.Status + mu *sync.Mutex +} + +// ActivatePeersByIP triggers wake-up for the peer(s) owning ips and waits until +// one is connected or ctx expires. Unknown or already-connected IPs are skipped, +// so the steady-state (warm) path adds no latency. +func (a *dnsPeerActivator) ActivatePeersByIP(ctx context.Context, ips []string) { + if a == nil || a.connMgr == nil { + return + } + + var pending []string + a.mu.Lock() + for _, ip := range ips { + st, ok := a.status.PeerStateByIP(ip) + if !ok || st.ConnStatus == peer.StatusConnected { + continue + } + conn, ok := a.peerStore.PeerConn(st.PubKey) + if !ok { + continue + } + a.connMgr.ActivatePeer(ctx, conn) + pending = append(pending, ip) + } + a.mu.Unlock() + + if len(pending) == 0 { + return + } + a.waitConnected(ctx, pending) +} + +// waitConnected blocks until any of ips reports a connected peer or ctx expires. +func (a *dnsPeerActivator) waitConnected(ctx context.Context, ips []string) { + ticker := time.NewTicker(dnsActivationPollInterval) + defer ticker.Stop() + for { + for _, ip := range ips { + if st, ok := a.status.PeerStateByIP(ip); ok && st.ConnStatus == peer.StatusConnected { + return + } + } + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} diff --git a/client/internal/engine.go b/client/internal/engine.go index 1d00ed0d2..eee4e3e29 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -654,6 +654,18 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL) e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface) e.connMgr.Start(e.ctx) + // Wire DNS-time lazy-connection warm-up now that the connection manager + // exists (it does not at DNS-server construction time). A DNS answer that + // points at an idle peer then wakes it before the client's first request. + if ds, ok := e.dnsServer.(*dns.DefaultServer); ok { + ds.SetPeerActivator(&dnsPeerActivator{ + connMgr: e.connMgr, + peerStore: e.peerStore, + status: e.statusRecorder, + mu: e.syncMsgMux, + }) + } + e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg) e.srWatcher.Start(peer.IsForceRelayed())