mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-22 00:11:29 +02:00
Compare commits
7 Commits
feature/dn
...
feature/na
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8fd098dcc7 | ||
|
|
820303c71d | ||
|
|
ad31406494 | ||
|
|
4d8d0e30db | ||
|
|
5aa2a748d7 | ||
|
|
e1a24376ab | ||
|
|
8f901f8899 |
@@ -6,7 +6,6 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -40,30 +39,6 @@ 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
|
||||
@@ -76,9 +51,6 @@ 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
|
||||
@@ -104,14 +76,6 @@ 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
|
||||
}
|
||||
@@ -158,9 +122,6 @@ 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
|
||||
@@ -557,33 +518,6 @@ 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 {
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
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")
|
||||
}
|
||||
@@ -491,13 +491,6 @@ 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()
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
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
|
||||
// ctx is the engine's long-lived context. The connection dial is tied to it
|
||||
// (not the per-query DNS wait budget) so a handshake that outlasts the wait
|
||||
// still completes in the background rather than being cancelled at the deadline.
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
// ActivatePeersByIP triggers wake-up for the peer(s) owning ips and waits until
|
||||
// one is connected or ctx (the per-query DNS wait budget) expires. Activation
|
||||
// itself is tied to the engine's long-lived context so the dial survives a wait
|
||||
// that times out. 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(a.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:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -654,19 +654,6 @@ 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,
|
||||
ctx: e.ctx,
|
||||
})
|
||||
}
|
||||
|
||||
e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg)
|
||||
e.srWatcher.Start(peer.IsForceRelayed())
|
||||
|
||||
|
||||
@@ -91,14 +91,7 @@ func availableProviders() []providerCase {
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
}
|
||||
// A valid Bedrock inference-profile id (region prefix + date + version),
|
||||
// overridable per account. `global.` profiles can be invoked from any
|
||||
// region; set AWS_BEDROCK_MODEL to match the enabled profile for the token.
|
||||
model := os.Getenv("AWS_BEDROCK_MODEL")
|
||||
if model == "" {
|
||||
model = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
}
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: model, kind: harness.WireBedrock})
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock})
|
||||
}
|
||||
return ps
|
||||
}
|
||||
@@ -115,16 +108,8 @@ func providerRequest(pc providerCase) api.AgentNetworkProviderRequest {
|
||||
Enabled: ptr(true),
|
||||
}
|
||||
if pc.kind != harness.WireVertex {
|
||||
// The router matches the normalized catalog id. Bedrock's request model
|
||||
// travels as a region-prefixed inference-profile id in the URL path
|
||||
// (us.anthropic...), which the router strips before matching, so register
|
||||
// the normalized form here or routing fails as model_not_routable.
|
||||
modelID := pc.model
|
||||
if pc.kind == harness.WireBedrock {
|
||||
modelID = catalogModel(pc)
|
||||
}
|
||||
req.Models = &[]api.AgentNetworkProviderModel{
|
||||
{Id: modelID, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
{Id: pc.model, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
}
|
||||
}
|
||||
return req
|
||||
@@ -216,13 +201,11 @@ func TestProvidersMatrix(t *testing.T) {
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
|
||||
// the proxy peer so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
|
||||
for _, pc := range matrix {
|
||||
pc := pc
|
||||
|
||||
@@ -4,7 +4,6 @@ package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -16,29 +15,13 @@ import (
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// bedrockRegionPrefixes and bedrockVersionSuffix mirror the proxy's Bedrock
|
||||
// model normalization (region/inference-profile prefix + version suffix) so the
|
||||
// provider is registered under the same catalog key the router matches against.
|
||||
var (
|
||||
bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
|
||||
bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
|
||||
)
|
||||
|
||||
// catalogModel returns the normalized catalog id the proxy stamps for a
|
||||
// path-routed provider's configured model — the form the router and guardrail
|
||||
// allowlist compare against (Bedrock region prefix + version stripped, Vertex
|
||||
// @version stripped).
|
||||
// path-routed provider's configured model — the form the guardrail allowlist is
|
||||
// compared against (region prefix / @version stripped).
|
||||
func catalogModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
m := pc.model
|
||||
for _, p := range bedrockRegionPrefixes {
|
||||
if strings.HasPrefix(m, p) {
|
||||
m = m[len(p):]
|
||||
break
|
||||
}
|
||||
}
|
||||
return bedrockVersionSuffix.ReplaceAllString(m, "")
|
||||
return strings.TrimPrefix(pc.model, "us.")
|
||||
case harness.WireVertex:
|
||||
return strings.SplitN(pc.model, "@", 2)[0]
|
||||
default:
|
||||
@@ -164,13 +147,11 @@ func TestModelAllowlistEnforced(t *testing.T) {
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
|
||||
// the proxy peer so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
|
||||
for _, pc := range providers {
|
||||
pc := pc
|
||||
|
||||
@@ -104,13 +104,11 @@ func TestProviderSkipTLSVerification(t *testing.T) {
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
|
||||
// the proxy peer so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
|
||||
// Positive: skip=true reaches the self-signed upstream. Retry to absorb
|
||||
// tunnel/DNS jitter on the first call; success also proves the path works.
|
||||
|
||||
@@ -106,13 +106,11 @@ func TestVLLMProvider(t *testing.T) {
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
|
||||
// the proxy peer so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
|
||||
before, _ := srv.ListAccessLogs(ctx)
|
||||
sessionID := "e2e-session-vllm"
|
||||
|
||||
1
go.mod
1
go.mod
@@ -103,6 +103,7 @@ require (
|
||||
github.com/rs/xid v1.3.0
|
||||
github.com/shirou/gopsutil/v4 v4.25.8
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
|
||||
github.com/soheilhy/cmux v0.1.5
|
||||
github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8
|
||||
github.com/stretchr/testify v1.11.1
|
||||
github.com/testcontainers/testcontainers-go v0.37.0
|
||||
|
||||
3
go.sum
3
go.sum
@@ -603,6 +603,8 @@ github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w
|
||||
github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g=
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
|
||||
github.com/soheilhy/cmux v0.1.5 h1:jjzc5WVemNEDTLwv9tlmemhC73tI08BNOIGwBOo10Js=
|
||||
github.com/soheilhy/cmux v0.1.5/go.mod h1:T7TcVDs9LWfQgPlPsdngu6I6QIoyIFZDDC6sNE1GqG0=
|
||||
github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8 h1:TG/diQgUe0pntT/2D9tmUCz4VNwm9MfrtPr0SU2qSX8=
|
||||
github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8/go.mod h1:P5HUIBuIWKbyjl083/loAegFkfbFNx5i2qEP4CNbm7E=
|
||||
github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY=
|
||||
@@ -757,6 +759,7 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
|
||||
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
|
||||
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
|
||||
golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
|
||||
@@ -24,13 +24,13 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/encryption"
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
|
||||
accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
|
||||
"github.com/netbirdio/netbird/management/server/activity"
|
||||
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
nbcache "github.com/netbirdio/netbird/management/server/cache"
|
||||
nbContext "github.com/netbirdio/netbird/management/server/context"
|
||||
nbhttp "github.com/netbirdio/netbird/management/server/http"
|
||||
@@ -184,7 +184,12 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
|
||||
grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
|
||||
}
|
||||
|
||||
if s.Config.HttpConfig.LetsEncryptDomain != "" {
|
||||
// With the native transport enabled, TLS is terminated at the listeners
|
||||
// (cmux-split shared listener and legacy port), so transport credentials
|
||||
// must not be set or the server would attempt a second handshake.
|
||||
if nativeGRPCEnabled() { //nolint:gocritic
|
||||
log.Info("native gRPC transport enabled, TLS is terminated at the listeners")
|
||||
} else if s.Config.HttpConfig.LetsEncryptDomain != "" {
|
||||
certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to create certificate service: %v", err)
|
||||
|
||||
@@ -6,12 +6,16 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/soheilhy/cmux"
|
||||
"go.opentelemetry.io/otel/metric"
|
||||
"golang.org/x/crypto/acme/autocert"
|
||||
"golang.org/x/net/http2"
|
||||
@@ -36,8 +40,18 @@ const (
|
||||
DefaultSelfHostedDomain = "netbird.selfhosted"
|
||||
|
||||
ContainerKeyBaseServer = "baseServer"
|
||||
|
||||
// NativeGRPCEnvVar enables serving gRPC on the native gRPC transport,
|
||||
// multiplexed with HTTP on the shared listener, instead of through the
|
||||
// net/http ServeHTTP path which costs two extra goroutines per stream.
|
||||
NativeGRPCEnvVar = "NB_MGMT_NATIVE_GRPC"
|
||||
)
|
||||
|
||||
func nativeGRPCEnabled() bool {
|
||||
enabled, _ := strconv.ParseBool(os.Getenv(NativeGRPCEnvVar))
|
||||
return enabled
|
||||
}
|
||||
|
||||
type Server interface {
|
||||
Start(ctx context.Context) error
|
||||
Stop() error
|
||||
@@ -182,11 +196,22 @@ func (s *BaseServer) Start(ctx context.Context) error {
|
||||
}
|
||||
}
|
||||
|
||||
// With the native transport enabled the gRPC server carries no transport
|
||||
// credentials, so TLS must be terminated at each of its listeners.
|
||||
var grpcTLSConfig *tls.Config
|
||||
if nativeGRPCEnabled() {
|
||||
if s.certManager != nil {
|
||||
grpcTLSConfig = s.certManager.TLSConfig()
|
||||
} else {
|
||||
grpcTLSConfig = tlsConfig
|
||||
}
|
||||
}
|
||||
|
||||
var compatListener net.Listener
|
||||
if s.mgmtPort != ManagementLegacyPort && !s.disableLegacyManagementPort {
|
||||
// The Management gRPC server was running on port 33073 previously. Old agents that are already connected to it
|
||||
// are using port 33073. For compatibility purposes we keep running a 2nd gRPC server on port 33073.
|
||||
compatListener, err = s.serveGRPC(srvCtx, s.GRPCServer(), ManagementLegacyPort)
|
||||
compatListener, err = s.serveGRPC(srvCtx, s.GRPCServer(), ManagementLegacyPort, grpcTLSConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -196,22 +221,38 @@ func (s *BaseServer) Start(ctx context.Context) error {
|
||||
rootHandler := s.handlerFunc(srvCtx, s.GRPCServer(), s.APIHandler(), s.IDPHandler(), s.Metrics().GetMeter())
|
||||
switch {
|
||||
case s.certManager != nil:
|
||||
// a call to certManager.Listener() always creates a new listener so we do it once
|
||||
cml := s.certManager.Listener()
|
||||
if s.mgmtPort == 443 {
|
||||
// CertManager, HTTP and gRPC API all on the same port
|
||||
rootHandler = s.certManager.HTTPHandler(rootHandler)
|
||||
s.listener = cml
|
||||
if nativeGRPCEnabled() {
|
||||
var tcpListener net.Listener
|
||||
tcpListener, err = net.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed creating TCP listener on port %d: %v", s.mgmtPort, err)
|
||||
}
|
||||
s.listener = tls.NewListener(tcpListener, preferHTTP1ForDualProtoClients(s.certManager.TLSConfig()))
|
||||
} else {
|
||||
s.listener = s.certManager.Listener()
|
||||
}
|
||||
} else {
|
||||
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), s.certManager.TLSConfig())
|
||||
mgmtTLSConfig := s.certManager.TLSConfig()
|
||||
if nativeGRPCEnabled() {
|
||||
mgmtTLSConfig = preferHTTP1ForDualProtoClients(mgmtTLSConfig)
|
||||
}
|
||||
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), mgmtTLSConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
|
||||
}
|
||||
cml := s.certManager.Listener()
|
||||
log.WithContext(ctx).Infof("running HTTP server (LetsEncrypt challenge handler): %s", cml.Addr().String())
|
||||
s.serveHTTP(ctx, cml, s.certManager.HTTPHandler(nil))
|
||||
}
|
||||
case tlsConfig != nil:
|
||||
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), tlsConfig)
|
||||
mgmtTLSConfig := tlsConfig
|
||||
if nativeGRPCEnabled() {
|
||||
mgmtTLSConfig = preferHTTP1ForDualProtoClients(mgmtTLSConfig)
|
||||
}
|
||||
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), mgmtTLSConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
|
||||
}
|
||||
@@ -224,7 +265,12 @@ func (s *BaseServer) Start(ctx context.Context) error {
|
||||
|
||||
log.WithContext(ctx).Infof("management server version %s", version.NetbirdVersion())
|
||||
log.WithContext(ctx).Infof("running HTTP server and gRPC server on the same port: %s", s.listener.Addr().String())
|
||||
s.serveGRPCWithHTTP(ctx, s.listener, rootHandler, tlsEnabled)
|
||||
if nativeGRPCEnabled() {
|
||||
log.WithContext(ctx).Infof("serving gRPC on the native transport (multiplexed with HTTP)")
|
||||
s.serveMultiplexed(ctx, s.listener, s.GRPCServer(), rootHandler, tlsEnabled)
|
||||
} else {
|
||||
s.serveGRPCWithHTTP(ctx, s.listener, rootHandler, tlsEnabled)
|
||||
}
|
||||
|
||||
s.update = version.NewUpdateAndStart("nb/management")
|
||||
s.update.SetDaemonVersion(version.NetbirdVersion())
|
||||
@@ -331,11 +377,14 @@ func (s *BaseServer) handlerFunc(_ context.Context, gRPCHandler *grpc.Server, ht
|
||||
})
|
||||
}
|
||||
|
||||
func (s *BaseServer) serveGRPC(ctx context.Context, grpcServer *grpc.Server, port int) (net.Listener, error) {
|
||||
func (s *BaseServer) serveGRPC(ctx context.Context, grpcServer *grpc.Server, port int, tlsConf *tls.Config) (net.Listener, error) {
|
||||
listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if tlsConf != nil {
|
||||
listener = tls.NewListener(listener, tlsConf)
|
||||
}
|
||||
|
||||
s.wg.Add(1)
|
||||
go func() {
|
||||
@@ -399,6 +448,69 @@ func (s *BaseServer) serveGRPCWithHTTP(ctx context.Context, listener net.Listene
|
||||
}()
|
||||
}
|
||||
|
||||
// preferHTTP1ForDualProtoClients steers TLS clients that offer both "h2" and
|
||||
// "http/1.1" in ALPN (browsers, REST clients) to HTTP/1.1. gRPC clients offer
|
||||
// only "h2", so with this steering every HTTP/2 connection on the shared
|
||||
// listener carries gRPC and can be routed to the native transport without
|
||||
// inspecting frames. ACME "acme-tls/1" and single-protocol clients keep the
|
||||
// base configuration.
|
||||
func preferHTTP1ForDualProtoClients(base *tls.Config) *tls.Config {
|
||||
h1Config := base.Clone()
|
||||
h1Config.NextProtos = []string{"http/1.1"}
|
||||
steered := base.Clone()
|
||||
steered.GetConfigForClient = func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
|
||||
if slices.Contains(hello.SupportedProtos, "http/1.1") && slices.Contains(hello.SupportedProtos, "h2") {
|
||||
return h1Config, nil
|
||||
}
|
||||
return base, nil
|
||||
}
|
||||
return steered
|
||||
}
|
||||
|
||||
// serveMultiplexed splits the shared listener by protocol: HTTP/2 connections
|
||||
// go to the native gRPC transport (see preferHTTP1ForDualProtoClients for why
|
||||
// they are all gRPC), everything else is served by net/http.
|
||||
//
|
||||
// Content-type based classification cannot be used here: cmux's SendSettings
|
||||
// matchers greet non-matching HTTP/2 connections and corrupt them for any
|
||||
// subsequent handler, while read-only matchers deadlock grpc-go clients,
|
||||
// which do not send HEADERS until they receive the server SETTINGS frame.
|
||||
func (s *BaseServer) serveMultiplexed(ctx context.Context, listener net.Listener, grpcServer *grpc.Server, handler http.Handler, tlsEnabled bool) {
|
||||
mux := cmux.New(listener)
|
||||
grpcListener := mux.Match(cmux.HTTP2())
|
||||
httpListener := mux.Match(cmux.Any())
|
||||
|
||||
httpHandler := handler
|
||||
if !tlsEnabled {
|
||||
//nolint:staticcheck // h2c also handles the HTTP/1 Upgrade mechanism, which http.Server's UnencryptedHTTP2 does not
|
||||
httpHandler = h2c.NewHandler(handler, &http2.Server{})
|
||||
}
|
||||
|
||||
s.wg.Add(3)
|
||||
go func() {
|
||||
defer s.wg.Done()
|
||||
s.reportServeError(ctx, grpcServer.Serve(grpcListener))
|
||||
}()
|
||||
go func() {
|
||||
defer s.wg.Done()
|
||||
s.reportServeError(ctx, http.Serve(httpListener, httpHandler))
|
||||
}()
|
||||
go func() {
|
||||
defer s.wg.Done()
|
||||
s.reportServeError(ctx, mux.Serve())
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *BaseServer) reportServeError(ctx context.Context, err error) {
|
||||
if ctx.Err() != nil || err == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case s.errCh <- err:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveDomains determines dnsDomain and mgmtSingleAccModeDomain based on store state.
|
||||
// Fresh installs use the default self-hosted domain, while existing installs reuse the
|
||||
// persisted account domain to keep addressing stable across config changes.
|
||||
|
||||
109
management/internals/server/server_multiplex_test.go
Normal file
109
management/internals/server/server_multiplex_test.go
Normal file
@@ -0,0 +1,109 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
"net"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/health"
|
||||
healthpb "google.golang.org/grpc/health/grpc_health_v1"
|
||||
)
|
||||
|
||||
func newSelfSignedCert(t *testing.T) tls.Certificate {
|
||||
t.Helper()
|
||||
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
|
||||
template := x509.Certificate{
|
||||
SerialNumber: big.NewInt(1),
|
||||
Subject: pkix.Name{CommonName: "127.0.0.1"},
|
||||
NotBefore: time.Now().Add(-time.Hour),
|
||||
NotAfter: time.Now().Add(time.Hour),
|
||||
IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
|
||||
}
|
||||
der, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
|
||||
require.NoError(t, err)
|
||||
|
||||
return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}
|
||||
}
|
||||
|
||||
func TestServeMultiplexedRoutesProtocols(t *testing.T) {
|
||||
tcpListener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = tcpListener.Close() })
|
||||
|
||||
baseTLSConfig := &tls.Config{
|
||||
Certificates: []tls.Certificate{newSelfSignedCert(t)},
|
||||
NextProtos: []string{"h2", "http/1.1"},
|
||||
}
|
||||
tlsListener := tls.NewListener(tcpListener, preferHTTP1ForDualProtoClients(baseTLSConfig))
|
||||
|
||||
grpcServer := grpc.NewServer()
|
||||
healthpb.RegisterHealthServer(grpcServer, health.NewServer())
|
||||
t.Cleanup(grpcServer.Stop)
|
||||
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = fmt.Fprintf(w, "proto=%d", r.ProtoMajor)
|
||||
})
|
||||
|
||||
s := &BaseServer{errCh: make(chan error, 4)}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancel)
|
||||
s.serveMultiplexed(ctx, tlsListener, grpcServer, handler, true)
|
||||
|
||||
addr := tcpListener.Addr().String()
|
||||
url := "https://" + addr + "/"
|
||||
|
||||
grpcConn, err := grpc.NewClient(addr,
|
||||
grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})))
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = grpcConn.Close() })
|
||||
|
||||
checkCtx, checkCancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer checkCancel()
|
||||
resp, err := healthpb.NewHealthClient(grpcConn).Check(checkCtx, &healthpb.HealthCheckRequest{})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, healthpb.HealthCheckResponse_SERVING, resp.Status)
|
||||
|
||||
get := func(client *http.Client) string {
|
||||
t.Helper()
|
||||
res, err := client.Get(url)
|
||||
require.NoError(t, err)
|
||||
body, err := io.ReadAll(res.Body)
|
||||
require.NoError(t, err)
|
||||
_ = res.Body.Close()
|
||||
return string(body)
|
||||
}
|
||||
|
||||
dualProtoClient := &http.Client{
|
||||
Timeout: 5 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
ForceAttemptHTTP2: true,
|
||||
},
|
||||
}
|
||||
require.Equal(t, "proto=1", get(dualProtoClient), "dual-ALPN client should be steered to HTTP/1.1")
|
||||
|
||||
h1OnlyClient := &http.Client{
|
||||
Timeout: 5 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true, NextProtos: []string{"http/1.1"}},
|
||||
},
|
||||
}
|
||||
require.Equal(t, "proto=1", get(h1OnlyClient))
|
||||
}
|
||||
@@ -215,12 +215,13 @@ func (s *Server) Job(srv proto.ManagementService_JobServer) error {
|
||||
return status.Errorf(codes.Unauthenticated, "peer is not registered")
|
||||
}
|
||||
|
||||
s.startResponseReceiver(ctx, srv)
|
||||
|
||||
updates := s.jobManager.CreateJobChannel(ctx, accountID, peer.ID)
|
||||
stream := s.jobManager.RegisterStream(ctx, accountID, peer.ID, func(event *job.Event) error {
|
||||
return s.sendJob(ctx, peerKey, event, srv)
|
||||
})
|
||||
defer s.jobManager.UnregisterStream(ctx, accountID, peer.ID, stream)
|
||||
log.WithContext(ctx).Debugf("Job: took %v", time.Since(reqStart))
|
||||
|
||||
return s.sendJobsLoop(ctx, accountID, peerKey, peer, updates, srv)
|
||||
return s.receiveJobResponses(ctx, peerKey, srv)
|
||||
}
|
||||
|
||||
// Sync validates the existence of a connecting peer, sends an initial state (all available for the connecting peers) and
|
||||
@@ -362,51 +363,26 @@ func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementServic
|
||||
return peerKey, nil
|
||||
}
|
||||
|
||||
func (s *Server) startResponseReceiver(ctx context.Context, srv proto.ManagementService_JobServer) {
|
||||
go func() {
|
||||
for {
|
||||
msg, err := srv.Recv()
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
|
||||
return
|
||||
}
|
||||
log.WithContext(ctx).Warnf("recv job response error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
jobResp := &proto.JobResponse{}
|
||||
if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
|
||||
log.WithContext(ctx).Warnf("invalid job response: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
|
||||
log.WithContext(ctx).Errorf("handle job response failed: %v", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates *job.Channel, srv proto.ManagementService_JobServer) error {
|
||||
// todo figure out better error handling strategy
|
||||
defer s.jobManager.CloseChannel(ctx, accountID, peer.ID)
|
||||
|
||||
func (s *Server) receiveJobResponses(ctx context.Context, peerKey wgtypes.Key, srv proto.ManagementService_JobServer) error {
|
||||
for {
|
||||
event, err := updates.Event(ctx)
|
||||
msg, err := srv.Recv()
|
||||
if err != nil {
|
||||
if errors.Is(err, job.ErrJobChannelClosed) {
|
||||
log.WithContext(ctx).Debugf("jobs channel for peer %s was closed", peerKey.String())
|
||||
return nil
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) || ctx.Err() != nil {
|
||||
log.WithContext(ctx).Debugf("job stream of peer %s has been closed", peerKey.String())
|
||||
return nil //nolint:nilerr
|
||||
}
|
||||
|
||||
// happens when connection drops, e.g. client disconnects
|
||||
log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
|
||||
return ctx.Err()
|
||||
log.WithContext(ctx).Warnf("recv job response error: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
if err := s.sendJob(ctx, peerKey, event, srv); err != nil {
|
||||
log.WithContext(ctx).Warnf("send job failed: %v", err)
|
||||
return nil
|
||||
jobResp := &proto.JobResponse{}
|
||||
if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
|
||||
log.WithContext(ctx).Warnf("invalid job response: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
|
||||
log.WithContext(ctx).Errorf("handle job response failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,8 +43,9 @@ type TimeBasedAuthSecretsManager struct {
|
||||
updateManager network_map.PeersUpdateManager
|
||||
settingsManager settings.Manager
|
||||
groupsManager groups.Manager
|
||||
turnCancelMap map[string]chan struct{}
|
||||
relayCancelMap map[string]chan struct{}
|
||||
scheduler *refreshScheduler
|
||||
turnJobs map[string]*refreshJob
|
||||
relayJobs map[string]*refreshJob
|
||||
wgKey wgtypes.Key
|
||||
}
|
||||
|
||||
@@ -60,8 +61,6 @@ func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager
|
||||
updateManager: updateManager,
|
||||
turnCfg: turnCfg,
|
||||
relayCfg: relayCfg,
|
||||
turnCancelMap: make(map[string]chan struct{}),
|
||||
relayCancelMap: make(map[string]chan struct{}),
|
||||
settingsManager: settingsManager,
|
||||
groupsManager: groupsManager,
|
||||
wgKey: key,
|
||||
@@ -127,16 +126,16 @@ func (m *TimeBasedAuthSecretsManager) GenerateRelayToken() (*Token, error) {
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) cancelTURN(peerID string) {
|
||||
if channel, ok := m.turnCancelMap[peerID]; ok {
|
||||
close(channel)
|
||||
delete(m.turnCancelMap, peerID)
|
||||
if job, ok := m.turnJobs[peerID]; ok {
|
||||
m.scheduler.cancel(job)
|
||||
delete(m.turnJobs, peerID)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) cancelRelay(peerID string) {
|
||||
if channel, ok := m.relayCancelMap[peerID]; ok {
|
||||
close(channel)
|
||||
delete(m.relayCancelMap, peerID)
|
||||
if job, ok := m.relayJobs[peerID]; ok {
|
||||
m.scheduler.cancel(job)
|
||||
delete(m.relayJobs, peerID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,6 +147,31 @@ func (m *TimeBasedAuthSecretsManager) CancelRefresh(peerID string) {
|
||||
m.cancelRelay(peerID)
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) ensureScheduler() {
|
||||
if m.scheduler == nil {
|
||||
m.scheduler = newRefreshScheduler(m.runRefreshJob)
|
||||
m.turnJobs = make(map[string]*refreshJob)
|
||||
m.relayJobs = make(map[string]*refreshJob)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) runRefreshJob(job *refreshJob) {
|
||||
switch job.kind {
|
||||
case refreshKindTURN:
|
||||
m.pushNewTURNAndRelayTokens(job.ctx, job.accountID, job.peerID)
|
||||
case refreshKindRelay:
|
||||
m.pushNewRelayTokens(job.ctx, job.accountID, job.peerID)
|
||||
}
|
||||
}
|
||||
|
||||
func refreshInterval(ttl time.Duration) time.Duration {
|
||||
interval := ttl / 4 * 3
|
||||
if interval <= 0 {
|
||||
interval = defaultDuration / 4 * 3
|
||||
}
|
||||
return interval
|
||||
}
|
||||
|
||||
// SetupRefresh starts peer credentials refresh
|
||||
func (m *TimeBasedAuthSecretsManager) SetupRefresh(ctx context.Context, accountID, peerID string) {
|
||||
m.mux.Lock()
|
||||
@@ -157,54 +181,38 @@ func (m *TimeBasedAuthSecretsManager) SetupRefresh(ctx context.Context, accountI
|
||||
m.cancelRelay(peerID)
|
||||
|
||||
if m.turnCfg != nil && m.turnCfg.TimeBasedCredentials {
|
||||
turnCancel := make(chan struct{}, 1)
|
||||
m.turnCancelMap[peerID] = turnCancel
|
||||
go m.refreshTURNTokens(ctx, accountID, peerID, turnCancel)
|
||||
m.ensureScheduler()
|
||||
job := &refreshJob{
|
||||
ctx: ctx,
|
||||
accountID: accountID,
|
||||
peerID: peerID,
|
||||
kind: refreshKindTURN,
|
||||
interval: refreshInterval(m.turnCfg.CredentialsTTL.Duration),
|
||||
}
|
||||
m.turnJobs[peerID] = job
|
||||
m.scheduler.schedule(job)
|
||||
log.WithContext(ctx).Debugf("starting TURN refresh for %s", peerID)
|
||||
} else {
|
||||
log.WithContext(ctx).Debugf("no TURN configuration, skipping TURN refresh for %s", peerID)
|
||||
}
|
||||
|
||||
if m.relayCfg != nil {
|
||||
relayCancel := make(chan struct{}, 1)
|
||||
m.relayCancelMap[peerID] = relayCancel
|
||||
go m.refreshRelayTokens(ctx, accountID, peerID, relayCancel)
|
||||
m.ensureScheduler()
|
||||
job := &refreshJob{
|
||||
ctx: ctx,
|
||||
accountID: accountID,
|
||||
peerID: peerID,
|
||||
kind: refreshKindRelay,
|
||||
interval: refreshInterval(m.relayCfg.CredentialsTTL.Duration),
|
||||
}
|
||||
m.relayJobs[peerID] = job
|
||||
m.scheduler.schedule(job)
|
||||
log.WithContext(ctx).Tracef("starting relay refresh for %s", peerID)
|
||||
} else {
|
||||
log.WithContext(ctx).Tracef("no relay configuration, skipping relay refresh for %s", peerID)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) refreshTURNTokens(ctx context.Context, accountID, peerID string, cancel chan struct{}) {
|
||||
ticker := time.NewTicker(m.turnCfg.CredentialsTTL.Duration / 4 * 3)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-cancel:
|
||||
log.WithContext(ctx).Tracef("stopping TURN refresh for %s", peerID)
|
||||
return
|
||||
case <-ticker.C:
|
||||
m.pushNewTURNAndRelayTokens(ctx, accountID, peerID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) refreshRelayTokens(ctx context.Context, accountID, peerID string, cancel chan struct{}) {
|
||||
ticker := time.NewTicker(m.relayCfg.CredentialsTTL.Duration / 4 * 3)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-cancel:
|
||||
log.WithContext(ctx).Tracef("stopping relay refresh for %s", peerID)
|
||||
return
|
||||
case <-ticker.C:
|
||||
m.pushNewRelayTokens(ctx, accountID, peerID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *TimeBasedAuthSecretsManager) pushNewTURNAndRelayTokens(ctx context.Context, accountID, peerID string) {
|
||||
turnToken, err := m.turnHmacToken.GenerateToken(sha1.New)
|
||||
if err != nil {
|
||||
|
||||
@@ -112,12 +112,12 @@ func TestTimeBasedAuthSecretsManager_SetupRefresh(t *testing.T) {
|
||||
|
||||
tested.SetupRefresh(ctx, "someAccountID", peer)
|
||||
|
||||
if _, ok := tested.turnCancelMap[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in the turn cancel map, got not present")
|
||||
if _, ok := tested.turnJobs[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in the turn jobs map, got not present")
|
||||
}
|
||||
|
||||
if _, ok := tested.relayCancelMap[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in the relay cancel map, got not present")
|
||||
if _, ok := tested.relayJobs[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in the relay jobs map, got not present")
|
||||
}
|
||||
|
||||
var updates []*network_map.UpdateMessage
|
||||
@@ -212,19 +212,26 @@ func TestTimeBasedAuthSecretsManager_CancelRefresh(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
tested.SetupRefresh(context.Background(), "someAccountID", peer)
|
||||
if _, ok := tested.turnCancelMap[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in turn cancel map, got not present")
|
||||
if _, ok := tested.turnJobs[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in turn jobs map, got not present")
|
||||
}
|
||||
if _, ok := tested.relayCancelMap[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in relay cancel map, got not present")
|
||||
if _, ok := tested.relayJobs[peer]; !ok {
|
||||
t.Errorf("expecting peer to be present in relay jobs map, got not present")
|
||||
}
|
||||
|
||||
tested.CancelRefresh(peer)
|
||||
if _, ok := tested.turnCancelMap[peer]; ok {
|
||||
t.Errorf("expecting peer to be not present in turn cancel map, got present")
|
||||
if _, ok := tested.turnJobs[peer]; ok {
|
||||
t.Errorf("expecting peer to be not present in turn jobs map, got present")
|
||||
}
|
||||
if _, ok := tested.relayCancelMap[peer]; ok {
|
||||
t.Errorf("expecting peer to be not present in relay cancel map, got present")
|
||||
if _, ok := tested.relayJobs[peer]; ok {
|
||||
t.Errorf("expecting peer to be not present in relay jobs map, got present")
|
||||
}
|
||||
|
||||
tested.scheduler.mu.Lock()
|
||||
heapLen := len(tested.scheduler.jobs)
|
||||
tested.scheduler.mu.Unlock()
|
||||
if heapLen != 0 {
|
||||
t.Errorf("expecting scheduler heap to be empty after cancel, got %d entries", heapLen)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
158
management/internals/shared/grpc/token_refresh_scheduler.go
Normal file
158
management/internals/shared/grpc/token_refresh_scheduler.go
Normal file
@@ -0,0 +1,158 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
refreshWorkerCount = 4
|
||||
refreshWorkQueueSize = 1024
|
||||
)
|
||||
|
||||
type refreshKind int
|
||||
|
||||
const (
|
||||
refreshKindTURN refreshKind = iota
|
||||
refreshKindRelay
|
||||
)
|
||||
|
||||
type refreshJob struct {
|
||||
ctx context.Context
|
||||
accountID string
|
||||
peerID string
|
||||
kind refreshKind
|
||||
interval time.Duration
|
||||
nextRun time.Time
|
||||
index int
|
||||
cancelled atomic.Bool
|
||||
}
|
||||
|
||||
type refreshJobHeap []*refreshJob
|
||||
|
||||
func (h refreshJobHeap) Len() int { return len(h) }
|
||||
func (h refreshJobHeap) Less(i, j int) bool { return h[i].nextRun.Before(h[j].nextRun) }
|
||||
|
||||
func (h refreshJobHeap) Swap(i, j int) {
|
||||
h[i], h[j] = h[j], h[i]
|
||||
h[i].index = i
|
||||
h[j].index = j
|
||||
}
|
||||
|
||||
func (h *refreshJobHeap) Push(x any) {
|
||||
job := x.(*refreshJob)
|
||||
job.index = len(*h)
|
||||
*h = append(*h, job)
|
||||
}
|
||||
|
||||
func (h *refreshJobHeap) Pop() any {
|
||||
old := *h
|
||||
n := len(old)
|
||||
job := old[n-1]
|
||||
old[n-1] = nil
|
||||
job.index = -1
|
||||
*h = old[:n-1]
|
||||
return job
|
||||
}
|
||||
|
||||
// refreshScheduler executes periodic credential refresh jobs for all peers
|
||||
// from one timer goroutine and a fixed worker pool, instead of two parked
|
||||
// goroutines per connected peer.
|
||||
type refreshScheduler struct {
|
||||
mu sync.Mutex
|
||||
jobs refreshJobHeap
|
||||
wake chan struct{}
|
||||
work chan *refreshJob
|
||||
run func(job *refreshJob)
|
||||
}
|
||||
|
||||
func newRefreshScheduler(run func(job *refreshJob)) *refreshScheduler {
|
||||
s := &refreshScheduler{
|
||||
wake: make(chan struct{}, 1),
|
||||
work: make(chan *refreshJob, refreshWorkQueueSize),
|
||||
run: run,
|
||||
}
|
||||
go s.loop()
|
||||
for range refreshWorkerCount {
|
||||
go s.worker()
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *refreshScheduler) schedule(job *refreshJob) {
|
||||
s.mu.Lock()
|
||||
job.nextRun = time.Now().Add(job.interval)
|
||||
heap.Push(&s.jobs, job)
|
||||
s.mu.Unlock()
|
||||
|
||||
select {
|
||||
case s.wake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *refreshScheduler) cancel(job *refreshJob) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
job.cancelled.Store(true)
|
||||
if job.index >= 0 {
|
||||
heap.Remove(&s.jobs, job.index)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *refreshScheduler) loop() {
|
||||
timer := time.NewTimer(time.Hour)
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
|
||||
for {
|
||||
s.mu.Lock()
|
||||
now := time.Now()
|
||||
var due []*refreshJob
|
||||
for len(s.jobs) > 0 && !s.jobs[0].nextRun.After(now) {
|
||||
job := s.jobs[0]
|
||||
job.nextRun = job.nextRun.Add(job.interval)
|
||||
if !job.nextRun.After(now) {
|
||||
job.nextRun = now.Add(job.interval)
|
||||
}
|
||||
heap.Fix(&s.jobs, 0)
|
||||
due = append(due, job)
|
||||
}
|
||||
wait := time.Duration(-1)
|
||||
if len(s.jobs) > 0 {
|
||||
wait = time.Until(s.jobs[0].nextRun)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
for _, job := range due {
|
||||
s.work <- job
|
||||
}
|
||||
|
||||
if wait < 0 {
|
||||
<-s.wake
|
||||
continue
|
||||
}
|
||||
|
||||
timer.Reset(wait)
|
||||
select {
|
||||
case <-s.wake:
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *refreshScheduler) worker() {
|
||||
for job := range s.work {
|
||||
if job.cancelled.Load() {
|
||||
continue
|
||||
}
|
||||
s.run(job)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRefreshInterval(t *testing.T) {
|
||||
defaultInterval := defaultDuration / 4 * 3
|
||||
|
||||
require.Equal(t, 9*time.Hour, refreshInterval(12*time.Hour))
|
||||
require.Equal(t, defaultInterval, refreshInterval(0))
|
||||
require.Equal(t, defaultInterval, refreshInterval(-time.Second))
|
||||
require.Equal(t, defaultInterval, refreshInterval(3*time.Nanosecond))
|
||||
require.Positive(t, refreshInterval(4*time.Nanosecond))
|
||||
}
|
||||
|
||||
func TestWorkerSkipsCancelledJob(t *testing.T) {
|
||||
var ran atomic.Int32
|
||||
scheduler := newRefreshScheduler(func(*refreshJob) {
|
||||
ran.Add(1)
|
||||
})
|
||||
|
||||
cancelledJob := &refreshJob{interval: time.Hour}
|
||||
cancelledJob.cancelled.Store(true)
|
||||
liveJob := &refreshJob{interval: time.Hour}
|
||||
|
||||
scheduler.work <- cancelledJob
|
||||
scheduler.work <- liveJob
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
return ran.Load() == 1
|
||||
}, 2*time.Second, 10*time.Millisecond, "live job should run exactly once, cancelled job never")
|
||||
}
|
||||
|
||||
func TestCancelBeforeFirePreventsRun(t *testing.T) {
|
||||
var ran atomic.Int32
|
||||
scheduler := newRefreshScheduler(func(*refreshJob) {
|
||||
ran.Add(1)
|
||||
})
|
||||
|
||||
job := &refreshJob{interval: 50 * time.Millisecond}
|
||||
scheduler.schedule(job)
|
||||
scheduler.cancel(job)
|
||||
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
require.Zero(t, ran.Load(), "cancelled job must never fire")
|
||||
|
||||
scheduler.mu.Lock()
|
||||
heapLen := len(scheduler.jobs)
|
||||
scheduler.mu.Unlock()
|
||||
require.Zero(t, heapLen)
|
||||
}
|
||||
@@ -1,19 +1,24 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
// nolint:gosec
|
||||
_ "net/http/pprof"
|
||||
"os"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/cmd"
|
||||
)
|
||||
|
||||
func main() {
|
||||
go func() {
|
||||
log.Println(http.ListenAndServe("localhost:6060", nil))
|
||||
}()
|
||||
if pprofAddr := os.Getenv("NB_PPROF_ADDR"); pprofAddr != "" {
|
||||
log.Infof("pprof enabled, listening on: %s", pprofAddr)
|
||||
go func() {
|
||||
log.Println(http.ListenAndServe(pprofAddr, nil))
|
||||
}()
|
||||
}
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
@@ -4,7 +4,6 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
@@ -21,11 +20,17 @@ type Event struct {
|
||||
Response *proto.JobResponse
|
||||
}
|
||||
|
||||
// PeerStream is the send side of a peer's Job stream. Sends are serialized by
|
||||
// its mutex; the receive side runs on the stream's gRPC handler goroutine.
|
||||
type PeerStream struct {
|
||||
send func(*Event) error
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
type Manager struct {
|
||||
mu *sync.RWMutex
|
||||
jobChannels map[string]*Channel // per-peer job streams
|
||||
pending map[string]*Event // jobID → event
|
||||
responseWait time.Duration
|
||||
streams map[string]*PeerStream // per-peer job streams
|
||||
pending map[string]*Event // jobID → event
|
||||
metrics telemetry.AppMetrics
|
||||
Store store.Store
|
||||
peersManager peers.Manager
|
||||
@@ -34,9 +39,8 @@ type Manager struct {
|
||||
func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager peers.Manager) *Manager {
|
||||
|
||||
return &Manager{
|
||||
jobChannels: make(map[string]*Channel),
|
||||
streams: make(map[string]*PeerStream),
|
||||
pending: make(map[string]*Event),
|
||||
responseWait: 5 * time.Minute,
|
||||
metrics: metrics,
|
||||
mu: &sync.RWMutex{},
|
||||
Store: store,
|
||||
@@ -44,8 +48,9 @@ func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager
|
||||
}
|
||||
}
|
||||
|
||||
// CreateJobChannel creates or replaces a channel for a peer
|
||||
func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID string) *Channel {
|
||||
// RegisterStream registers the send side of a peer's Job stream, replacing any
|
||||
// previous registration for the peer.
|
||||
func (jm *Manager) RegisterStream(ctx context.Context, accountID, peerID string, send func(*Event) error) *PeerStream {
|
||||
// all pending jobs stored in db for this peer should be failed
|
||||
if err := jm.Store.MarkAllPendingJobsAsFailed(ctx, accountID, peerID, "Pending job cleanup: marked as failed automatically due to being stuck too long"); err != nil {
|
||||
log.WithContext(ctx).Error(err.Error())
|
||||
@@ -54,23 +59,41 @@ func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID strin
|
||||
jm.mu.Lock()
|
||||
defer jm.mu.Unlock()
|
||||
|
||||
if ch, ok := jm.jobChannels[peerID]; ok {
|
||||
ch.Close()
|
||||
delete(jm.jobChannels, peerID)
|
||||
}
|
||||
stream := &PeerStream{send: send}
|
||||
jm.streams[peerID] = stream
|
||||
return stream
|
||||
}
|
||||
|
||||
ch := NewChannel()
|
||||
jm.jobChannels[peerID] = ch
|
||||
return ch
|
||||
// UnregisterStream removes a peer's stream registration and fails its pending
|
||||
// jobs. It is a no-op if the registration was already replaced by a newer
|
||||
// stream of the same peer.
|
||||
func (jm *Manager) UnregisterStream(ctx context.Context, accountID, peerID string, stream *PeerStream) {
|
||||
jm.mu.Lock()
|
||||
defer jm.mu.Unlock()
|
||||
|
||||
if jm.streams[peerID] != stream {
|
||||
return
|
||||
}
|
||||
delete(jm.streams, peerID)
|
||||
|
||||
for jobID, ev := range jm.pending {
|
||||
if ev.PeerID == peerID {
|
||||
// if the client disconnect and there is pending job then mark it as failed
|
||||
if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
|
||||
}
|
||||
delete(jm.pending, jobID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SendJob sends a job to a peer and tracks it as pending
|
||||
func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *proto.JobRequest) error {
|
||||
jm.mu.RLock()
|
||||
ch, ok := jm.jobChannels[peerID]
|
||||
stream, ok := jm.streams[peerID]
|
||||
jm.mu.RUnlock()
|
||||
if !ok {
|
||||
return fmt.Errorf("peer %s has no channel", peerID)
|
||||
return fmt.Errorf("peer %s has no stream", peerID)
|
||||
}
|
||||
|
||||
event := &Event{
|
||||
@@ -82,7 +105,10 @@ func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *p
|
||||
jm.pending[string(req.ID)] = event
|
||||
jm.mu.Unlock()
|
||||
|
||||
if err := ch.AddEvent(ctx, jm.responseWait, event); err != nil {
|
||||
stream.mu.Lock()
|
||||
err := stream.send(event)
|
||||
stream.mu.Unlock()
|
||||
if err != nil {
|
||||
jm.cleanup(ctx, accountID, string(req.ID), err.Error())
|
||||
return err
|
||||
}
|
||||
@@ -127,27 +153,6 @@ func (jm *Manager) HandleResponse(ctx context.Context, resp *proto.JobResponse,
|
||||
return nil
|
||||
}
|
||||
|
||||
// CloseChannel closes a peer’s channel and cleans up its jobs
|
||||
func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string) {
|
||||
jm.mu.Lock()
|
||||
defer jm.mu.Unlock()
|
||||
|
||||
if ch, ok := jm.jobChannels[peerID]; ok {
|
||||
ch.Close()
|
||||
delete(jm.jobChannels, peerID)
|
||||
}
|
||||
|
||||
for jobID, ev := range jm.pending {
|
||||
if ev.PeerID == peerID {
|
||||
// if the client disconnect and there is pending job then mark it as failed
|
||||
if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
|
||||
}
|
||||
delete(jm.pending, jobID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanup removes a pending job safely
|
||||
func (jm *Manager) cleanup(ctx context.Context, accountID, jobID string, reason string) {
|
||||
jm.mu.Lock()
|
||||
@@ -165,7 +170,7 @@ func (jm *Manager) IsPeerConnected(peerID string) bool {
|
||||
jm.mu.RLock()
|
||||
defer jm.mu.RUnlock()
|
||||
|
||||
_, ok := jm.jobChannels[peerID]
|
||||
_, ok := jm.streams[peerID]
|
||||
return ok
|
||||
}
|
||||
|
||||
|
||||
90
management/server/job/manager_test.go
Normal file
90
management/server/job/manager_test.go
Normal file
@@ -0,0 +1,90 @@
|
||||
package job
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
func newTestManager(t *testing.T) (*Manager, *store.MockStore) {
|
||||
t.Helper()
|
||||
ctrl := gomock.NewController(t)
|
||||
t.Cleanup(ctrl.Finish)
|
||||
mockStore := store.NewMockStore(ctrl)
|
||||
return NewJobManager(nil, mockStore, nil), mockStore
|
||||
}
|
||||
|
||||
func TestSendJobDeliversThroughRegisteredStream(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
manager, mockStore := newTestManager(t)
|
||||
mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
|
||||
|
||||
var sent []*Event
|
||||
manager.RegisterStream(ctx, "acc", "peer1", func(event *Event) error {
|
||||
sent = append(sent, event)
|
||||
return nil
|
||||
})
|
||||
require.True(t, manager.IsPeerConnected("peer1"))
|
||||
|
||||
err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, sent, 1)
|
||||
require.Equal(t, "peer1", sent[0].PeerID)
|
||||
require.True(t, manager.IsPeerHasPendingJobs("peer1"))
|
||||
}
|
||||
|
||||
func TestSendJobWithoutStream(t *testing.T) {
|
||||
manager, _ := newTestManager(t)
|
||||
err := manager.SendJob(context.Background(), "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSendJobFailureCleansPending(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
manager, mockStore := newTestManager(t)
|
||||
mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
|
||||
mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
|
||||
|
||||
manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error {
|
||||
return errors.New("stream broken")
|
||||
})
|
||||
|
||||
err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
|
||||
require.Error(t, err)
|
||||
require.False(t, manager.IsPeerHasPendingJobs("peer1"))
|
||||
}
|
||||
|
||||
func TestUnregisterStreamIgnoresSupersededRegistration(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
manager, mockStore := newTestManager(t)
|
||||
mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil).Times(2)
|
||||
|
||||
first := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
|
||||
second := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
|
||||
|
||||
manager.UnregisterStream(ctx, "acc", "peer1", first)
|
||||
require.True(t, manager.IsPeerConnected("peer1"), "stale unregister must not remove the replacement stream")
|
||||
|
||||
manager.UnregisterStream(ctx, "acc", "peer1", second)
|
||||
require.False(t, manager.IsPeerConnected("peer1"))
|
||||
}
|
||||
|
||||
func TestUnregisterStreamFailsPendingJobs(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
manager, mockStore := newTestManager(t)
|
||||
mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
|
||||
mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
|
||||
|
||||
stream := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
|
||||
require.NoError(t, manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")}))
|
||||
require.True(t, manager.IsPeerHasPendingJobs("peer1"))
|
||||
|
||||
manager.UnregisterStream(ctx, "acc", "peer1", stream)
|
||||
require.False(t, manager.IsPeerHasPendingJobs("peer1"))
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ssh/auth"
|
||||
@@ -42,6 +43,14 @@ type NetworkMapComponents struct {
|
||||
PostureFailedPeers map[string]map[string]struct{}
|
||||
|
||||
RouterPeers map[string]*nbpeer.Peer
|
||||
|
||||
routesByPeerOnce sync.Once
|
||||
routesByPeerIdx map[string][]routeIndexEntry
|
||||
}
|
||||
|
||||
type routeIndexEntry struct {
|
||||
route *route.Route
|
||||
viaGroup bool
|
||||
}
|
||||
|
||||
type AccountSettingsInfo struct {
|
||||
@@ -530,33 +539,43 @@ func (c *NetworkMapComponents) getRoutingPeerRoutes(peerID string) (enabledRoute
|
||||
disabledRoutes = append(disabledRoutes, r)
|
||||
}
|
||||
|
||||
for _, r := range c.Routes {
|
||||
for _, groupID := range r.PeerGroups {
|
||||
group := c.GetGroupInfo(groupID)
|
||||
if group == nil {
|
||||
continue
|
||||
}
|
||||
for _, id := range group.Peers {
|
||||
if id != peerID {
|
||||
continue
|
||||
}
|
||||
|
||||
newPeerRoute := r.Copy()
|
||||
newPeerRoute.Peer = id
|
||||
newPeerRoute.PeerGroups = nil
|
||||
newPeerRoute.ID = route.ID(string(r.ID) + ":" + id)
|
||||
takeRoute(newPeerRoute)
|
||||
break
|
||||
}
|
||||
}
|
||||
if r.Peer == peerID {
|
||||
takeRoute(r.Copy())
|
||||
for _, entry := range c.routesByPeer()[peerID] {
|
||||
if entry.viaGroup {
|
||||
newPeerRoute := entry.route.Copy()
|
||||
newPeerRoute.PeerGroups = nil
|
||||
newPeerRoute.ID = route.ID(string(entry.route.ID) + ":" + peerID)
|
||||
takeRoute(newPeerRoute)
|
||||
continue
|
||||
}
|
||||
takeRoute(entry.route.Copy())
|
||||
}
|
||||
|
||||
return enabledRoutes, disabledRoutes
|
||||
}
|
||||
|
||||
func (c *NetworkMapComponents) routesByPeer() map[string][]routeIndexEntry {
|
||||
c.routesByPeerOnce.Do(func() {
|
||||
idx := make(map[string][]routeIndexEntry)
|
||||
for _, r := range c.Routes {
|
||||
for _, groupID := range r.PeerGroups {
|
||||
group := c.GetGroupInfo(groupID)
|
||||
if group == nil {
|
||||
continue
|
||||
}
|
||||
for _, id := range group.Peers {
|
||||
idx[id] = append(idx[id], routeIndexEntry{route: r, viaGroup: true})
|
||||
}
|
||||
}
|
||||
if r.Peer != "" {
|
||||
idx[r.Peer] = append(idx[r.Peer], routeIndexEntry{route: r})
|
||||
}
|
||||
}
|
||||
c.routesByPeerIdx = idx
|
||||
})
|
||||
|
||||
return c.routesByPeerIdx
|
||||
}
|
||||
|
||||
func (c *NetworkMapComponents) filterRoutesByGroups(routes []*route.Route, groupListMap LookupMap) []*route.Route {
|
||||
var filteredRoutes []*route.Route
|
||||
for _, r := range routes {
|
||||
|
||||
Reference in New Issue
Block a user