spb/internetpressure

1package main23import (4	"context"5	"encoding/json"6	"errors"7	"fmt"8	"log/slog"9	"os"10	"path/filepath"11	"sync"12	"sync/atomic"13	"time"1415	"internetpressure.io/probe-agent/internal/batcher"16	checkdns "internetpressure.io/probe-agent/internal/checks/dns"17	checkhttp "internetpressure.io/probe-agent/internal/checks/http"18	checkping "internetpressure.io/probe-agent/internal/checks/ping"19	checktrace "internetpressure.io/probe-agent/internal/checks/traceroute"20	"internetpressure.io/probe-agent/internal/client"21	"internetpressure.io/probe-agent/internal/config"22	"internetpressure.io/probe-agent/internal/health"23	"internetpressure.io/probe-agent/internal/identity"24	"internetpressure.io/probe-agent/internal/protocol"25	"internetpressure.io/probe-agent/internal/sched"26	"internetpressure.io/probe-agent/internal/signer"27	"internetpressure.io/probe-agent/internal/spool"28	"internetpressure.io/probe-agent/internal/update"29)3031const (32	cachedConfigFile   = "config.json"33	configBackoffMin   = 5 * time.Second34	configBackoffMax   = 5 * time.Minute35	defaultRefresh     = 300 * time.Second36	clockWarnThreshold = 120 * time.Second37	clockWarnEvery     = 10 * time.Minute38)3940// agent wires every component together.41type agent struct {42	cfg *config.Config43	log *slog.Logger4445	client   *client.Client46	spool    *spool.Spool47	health   *health.State48	batcher  *batcher.Batcher49	sched    *sched.Scheduler50	identity *identity.Service51	updater  *update.Updater5253	httpCheck  *checkhttp.Checker54	dnsCheck   *checkdns.Checker55	pingCheck  *checkping.Checker56	traceCheck *checktrace.Checker57	caps       []string5859	remote    atomic.Pointer[protocol.RemoteConfig]60	refreshCh chan struct{}61	lastClock time.Time62	clockMu   sync.Mutex63	updated   atomic.Bool64	cancel    context.CancelFunc65}6667func newAgent(cfg *config.Config, log *slog.Logger) (*agent, error) {68	if err := os.MkdirAll(cfg.DataDir, 0o700); err != nil {69		return nil, fmt.Errorf("data_dir: %w", err)70	}71	sg, err := signer.New(cfg.ProbeID, cfg.Key)72	if err != nil {73		return nil, err74	}75	cl, err := client.New(cfg.IngestURL, sg, version)76	if err != nil {77		return nil, err78	}79	sp, err := spool.Open(filepath.Join(cfg.DataDir, "spool"), spool.DefaultCap)80	if err != nil {81		return nil, err82	}8384	a := &agent{cfg: cfg, log: log, client: cl, spool: sp, refreshCh: make(chan struct{}, 1)}85	ua := client.UserAgent(version)86	a.httpCheck = &checkhttp.Checker{UserAgent: ua}87	a.dnsCheck = &checkdns.Checker{}88	a.pingCheck = &checkping.Checker{ICMP: checkping.Detect()}89	a.traceCheck = checktrace.Detect()90	a.caps = []string{"http", "dns"}91	if a.pingCheck.ICMP {92		a.caps = append(a.caps, "ping")93	} else {94		a.caps = append(a.caps, "tcp")95		log.Warn("unprivileged ICMP unavailable: ping checks fall back to tcp connects", "hint", checkping.LinuxHint())96	}97	if a.traceCheck.Available() {98		a.caps = append(a.caps, "traceroute")99	} else {100		log.Warn("traceroute binary not found: traceroute disabled")101	}102103	a.identity = identity.New(ua)104	a.sched = sched.New(a.runJob, sched.Options{105		MaxConcurrency: cfg.MaxConcurrency,106		Logger:         log,107		CanPing:        a.pingCheck.ICMP,108		CanTraceroute:  a.traceCheck.Available(),109	})110	a.health = health.New(cfg.ProbeID, version, a.caps, health.Gauges{111		Buffered:      func() int { return a.batcher.Buffered() },112		SpoolBytes:    sp.Bytes,113		ClockOffsetMs: cl.Clock.OffsetMs,114		Targets:       a.sched.Targets,115		Identity:      a.identity.Current,116	})117	a.batcher = batcher.New(cfg.ProbeID, version, cl, sp, batcher.Hooks{118		Health:         a.health.Snapshot,119		OnResponse:     a.onBatchResponse,120		OnFlushFailure: a.health.RecordFlushFailure,121		OnDelivered: func(t time.Time) {122			a.health.SetLastFlush(t)123			a.checkClock()124		},125	}, log)126	a.updater = &update.Updater{Fetcher: cl, Version: version, DataDir: cfg.DataDir, Log: log}127	return a, nil128}129130// Run starts every loop and blocks until ctx is cancelled (SIGINT/SIGTERM) or a self-update requires a restart.131func (a *agent) Run(parent context.Context) error {132	ctx, cancel := context.WithCancel(parent)133	a.cancel = cancel134	defer cancel()135	a.log.Info("ip-probe starting", "version", version, "probe_id", a.cfg.ProbeID, "ingest_url", a.cfg.IngestURL,136		"data_dir", a.cfg.DataDir, "listen", a.cfg.Listen, "capabilities", a.caps, "max_concurrency", a.cfg.MaxConcurrency)137138	// Offline start: reuse the last good config so checks begin before the server answers.139	if rc, err := a.loadCachedConfig(); err == nil {140		a.log.Info("starting with cached config", "config_version", rc.ConfigVersion, "targets", len(rc.Targets))141		a.applyRemote(rc, false)142	}143144	var wg sync.WaitGroup145	start := func(name string, f func(context.Context)) {146		wg.Add(1)147		go func() {148			defer wg.Done()149			f(ctx)150			a.log.Debug("loop stopped", "loop", name)151		}()152	}153	go func() {154		if err := a.health.Serve(ctx, a.cfg.Listen); err != nil {155			a.log.Error("local health endpoint failed", "listen", a.cfg.Listen, "err", err)156		}157	}()158	var schedDone sync.WaitGroup159	schedDone.Add(1)160	go func() { defer schedDone.Done(); a.sched.Run(ctx) }()161	start("batcher", a.batcher.Run)162	start("config", a.configLoop)163	start("identity", func(ctx context.Context) {164		a.identity.Run(ctx, func(id *protocol.Identity, err error) {165			if err != nil {166				a.log.Warn("identity lookup failed", "err", err)167				return168			}169			a.log.Info("identity", "public_ip", id.PublicIP, "asn", id.ASN, "org", id.Org, "country", id.Country,170				"city", id.City, "source", id.Source)171		})172	})173	if a.cfg.AllowSelfUpdate {174		start("update", func(ctx context.Context) {175			a.updater.Run(ctx, func() {176				a.updated.Store(true)177				a.log.Info("self-update installed; shutting down for the supervisor to restart")178				cancel()179			})180		})181	} else {182		a.log.Info("self-update disabled by configuration")183	}184185	<-ctx.Done()186	a.log.Info("shutting down", "reason", reason(parent, a.updated.Load()))187	// 1. stop scheduling and wait for in-flight checks; 2. final flush (one 5 s attempt, then spool).188	schedDone.Wait()189	wg.Wait()190	a.batcher.Stop()191	a.log.Info("stopped", "buffered", a.batcher.Buffered(), "spool_bytes", a.spool.Bytes())192	return nil193}194195func reason(parent context.Context, updated bool) string {196	switch {197	case updated:198		return "self-update"199	case parent.Err() != nil:200		return "signal"201	default:202		return "internal"203	}204}205206// runJob executes one scheduled (target, family) job. Measurements from cancelled checks are dropped.207func (a *agent) runJob(ctx context.Context, job sched.Job) {208	switch job.Family {209	case sched.FamilyHTTP:210		a.emit(ctx, a.httpCheck.Run(ctx, job.Target))211	case sched.FamilyDNS:212		a.emit(ctx, a.dnsCheck.Run(ctx, job.Target, a.resolvers())...)213	case sched.FamilyPing:214		a.emit(ctx, a.pingCheck.Run(ctx, job.Target))215	case sched.FamilyTCP:216		a.emit(ctx, a.pingCheck.RunTCP(ctx, job.Target))217	case sched.FamilyTraceroute:218		tr, err := a.traceCheck.Run(ctx, job.Target)219		if ctx.Err() != nil {220			return221		}222		if err != nil {223			a.log.Debug("traceroute failed", "target", job.Target.TargetID, "err", err)224			return225		}226		a.batcher.AddTraceroute(tr)227	}228}229230func (a *agent) emit(ctx context.Context, ms ...protocol.Measurement) {231	if ctx.Err() != nil {232		return233	}234	for _, m := range ms {235		a.health.Record(m)236	}237	a.batcher.Add(ms...)238}239240// resolvers returns the DNS resolver list: local override, else server config, else defaults.241func (a *agent) resolvers() []protocol.Resolver {242	if len(a.cfg.ResolversOverride) > 0 {243		out := make([]protocol.Resolver, 0, len(a.cfg.ResolversOverride))244		for _, r := range a.cfg.ResolversOverride {245			id, addr, err := config.ParseResolver(r)246			if err == nil {247				out = append(out, protocol.Resolver{ID: id, Address: addr})248			}249		}250		return out251	}252	if rc := a.remote.Load(); rc != nil && len(rc.Resolvers) > 0 {253		return rc.Resolvers254	}255	return checkdns.DefaultResolvers256}257258// configLoop fetches /config with backoff until it succeeds, then every config_refresh_seconds or on demand.259func (a *agent) configLoop(ctx context.Context) {260	backoff := configBackoffMin261	for {262		rc, err := a.client.GetConfig(ctx)263		if ctx.Err() != nil {264			return265		}266		var wait time.Duration267		if err != nil {268			a.log.Warn("config fetch failed", "err", err, "retry_in", backoff)269			wait = backoff270			backoff *= 2271			if backoff > configBackoffMax {272				backoff = configBackoffMax273			}274		} else {275			backoff = configBackoffMin276			a.applyRemote(rc, true)277			wait = defaultRefresh278			if rc.Schedule.ConfigRefreshSeconds > 0 {279				wait = time.Duration(rc.Schedule.ConfigRefreshSeconds) * time.Second280			}281			if wait < 30*time.Second {282				wait = 30 * time.Second283			}284		}285		select {286		case <-ctx.Done():287			return288		case <-time.After(wait):289		case <-a.refreshCh:290		}291	}292}293294// applyRemote installs a remote config everywhere and (optionally) persists it for offline starts.295func (a *agent) applyRemote(rc *protocol.RemoteConfig, persist bool) {296	if rc.Probe.ProbeID != "" && rc.Probe.ProbeID != a.cfg.ProbeID {297		a.log.Warn("server config is for another probe id", "got", rc.Probe.ProbeID, "want", a.cfg.ProbeID)298	}299	prev := a.remote.Load()300	a.remote.Store(rc)301	a.sched.Apply(rc)302	a.sched.SetBoost(rc.Schedule.Boost)303	a.batcher.Configure(rc.Schedule.BatchFlushSeconds, rc.Schedule.MaxBatch)304	a.health.SetLastConfig(time.Now(), rc.ConfigVersion)305	if !rc.Probe.Enabled {306		a.log.Warn("probe disabled by server: idling (health only)")307	}308	if prev == nil || prev.ConfigVersion != rc.ConfigVersion {309		a.log.Info("config applied", "config_version", rc.ConfigVersion, "targets", len(rc.Targets),310			"resolvers", len(rc.Resolvers), "enabled", rc.Probe.Enabled, "tiers", rc.Schedule.Tiers,311			"dns_every", rc.Schedule.DNSEvery, "ping_every", rc.Schedule.PingEvery, "traceroute_every", rc.Schedule.TracerouteEvery)312	}313	if persist {314		if err := a.saveCachedConfig(rc); err != nil {315			a.log.Warn("could not persist config", "err", err)316		}317	}318	a.checkClock()319}320321// onBatchResponse reacts to a batch reply: new config_version → refresh now; boost → apply.322func (a *agent) onBatchResponse(resp *protocol.BatchResponse) {323	if resp.Boost != nil {324		a.sched.SetBoost(resp.Boost)325	}326	cur := a.remote.Load()327	if resp.ConfigVersion != "" && (cur == nil || cur.ConfigVersion != resp.ConfigVersion) {328		a.log.Info("server announced a new config version", "config_version", resp.ConfigVersion)329		select {330		case a.refreshCh <- struct{}{}:331		default:332		}333	}334}335336// checkClock warns (rate-limited) when the estimated clock offset exceeds 120 s.337func (a *agent) checkClock() {338	off := a.client.Clock.Offset()339	if off < 0 {340		off = -off341	}342	if off <= clockWarnThreshold {343		return344	}345	a.clockMu.Lock()346	defer a.clockMu.Unlock()347	if time.Since(a.lastClock) < clockWarnEvery {348		return349	}350	a.lastClock = time.Now()351	a.log.Warn("local clock differs from server by more than 120 s; fix NTP", "clock_offset_ms", a.client.Clock.OffsetMs())352}353354func (a *agent) cachedConfigPath() string { return filepath.Join(a.cfg.DataDir, cachedConfigFile) }355356func (a *agent) loadCachedConfig() (*protocol.RemoteConfig, error) {357	data, err := os.ReadFile(a.cachedConfigPath())358	if err != nil {359		return nil, err360	}361	var rc protocol.RemoteConfig362	if err := json.Unmarshal(data, &rc); err != nil {363		return nil, err364	}365	if rc.ConfigVersion == "" {366		return nil, errors.New("cached config has no version")367	}368	return &rc, nil369}370371func (a *agent) saveCachedConfig(rc *protocol.RemoteConfig) error {372	data, err := json.Marshal(rc)373	if err != nil {374		return err375	}376	tmp := a.cachedConfigPath() + ".tmp"377	if err := os.WriteFile(tmp, data, 0o600); err != nil {378		return err379	}380	return os.Rename(tmp, a.cachedConfigPath())381}382