nmon/internal/simagent/agent.go

// Package simagent is an SNMP agent that serves snmpsim-format .snmprec
// recordings over UDP (v1, v2c, v3 USM) with dynamic counters and fault
// injection. Tests and the loadtest use it in place of Python snmpsim.
package simagent

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gosnmp/gosnmp"

	"github.com/v10networks/nmon/internal/clock"
	"github.com/v10networks/nmon/internal/snmprec"
)

// Behaviour injects faults.
type Behaviour struct {
	Latency         time.Duration // delay before every response, slept on the agent's clock
	DropEvery       int           // silently drop every Nth request (0 = never)
	TooBigAbove     int           // GET/GETNEXT with more varbinds than this returns tooBig (0 = no limit)
	MaxResponseSize int           // bytes; larger responses are tooBig (GET) or truncated (GETBULK). Default 65000
}

// V3User is a USM user. Auth and priv are required; the agent only serves authPriv.
type V3User struct {
	Name      string
	AuthProto gosnmp.SnmpV3AuthProtocol
	AuthPass  string
	PrivProto gosnmp.SnmpV3PrivProtocol
	PrivPass  string
}

// Config describes one agent.
type Config struct {
	Addr      string // UDP listen address, default "[::1]:0"
	Recording *snmprec.Recording
	Community string // v1/v2c community, default "public"
	V3        []V3User
	// Contexts serve other recordings under the same agent, the way an IOS switch serves one bridge table
	// per VLAN: a v1/v2c request with community "<community>@<name>" or a v3 request with context name
	// "<name>" is answered from Contexts[name]; a name not listed gets the default recording.
	Contexts  map[string]*snmprec.Recording
	EngineID  string // v3 authoritative engine ID, default derived from the bound address
	Behaviour Behaviour
	Clock     clock.Clock // default clock.Real()
}

// Stats counts requests that passed community or USM checks, by PDU type.
type Stats struct {
	Get, GetNext, GetBulk uint64
	Set                   uint64 // rejected, counted for completeness
	Discovery             uint64 // v3 engine-ID discovery probes
	Dropped               uint64 // dropped by Behaviour.DropEvery
	Invalid               uint64 // undecodable, wrong community, unknown user or bad key
	TooBig                uint64 // tooBig responses sent
	NotInTimeWindow       uint64 // v3 notInTimeWindow reports sent
}

// Agent is a running simulated device.
type Agent struct {
	cfg       Config
	rec       *snmprec.Recording
	clk       clock.Clock
	start     time.Time
	conn      *net.UDPConn
	addr      netip.AddrPort
	community string
	engineID  string
	dec       *gosnmp.GoSNMP // decodes requests; holds the USM user table
	hasV3     bool
	maxSize   int

	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
	once   sync.Once

	reqN atomic.Uint64
	salt atomic.Uint64
	st   struct{ get, next, bulk, set, disc, dropped, invalid, tooBig, notInTime atomic.Uint64 }
}

const (
	engineBoots           = 1
	usmStatsNotInTime     = ".1.3.6.1.6.3.15.1.1.2.0"
	usmStatsUnknownEngine = ".1.3.6.1.6.3.15.1.1.4.0"
)

// Start binds cfg.Addr and serves until ctx is cancelled or Close is called.
func Start(ctx context.Context, cfg Config) (*Agent, error) {
	if cfg.Recording == nil {
		return nil, errors.New("simagent: no recording")
	}
	if cfg.Addr == "" {
		cfg.Addr = "[::1]:0"
	}
	ap, err := netip.ParseAddrPort(cfg.Addr)
	if err != nil {
		return nil, fmt.Errorf("simagent: %w", err)
	}
	conn, err := net.ListenUDP("udp", net.UDPAddrFromAddrPort(ap))
	if err != nil {
		return nil, fmt.Errorf("simagent: %w", err)
	}
	a := &Agent{
		cfg: cfg, rec: cfg.Recording, clk: cfg.Clock, conn: conn,
		addr:      conn.LocalAddr().(*net.UDPAddr).AddrPort(),
		community: cfg.Community, maxSize: cfg.Behaviour.MaxResponseSize,
		hasV3: len(cfg.V3) > 0,
	}
	if a.clk == nil {
		a.clk = clock.Real()
	}
	if a.community == "" {
		a.community = "public"
	}
	if a.maxSize <= 0 {
		a.maxSize = 65000
	}
	a.start = a.clk.Now()
	id := "nmon-" + a.addr.String()
	if len(id) > 27 {
		id = id[:27]
	}
	a.engineID = "\x80\x00\x1f\x88\x04" + id
	if cfg.EngineID != "" {
		a.engineID = cfg.EngineID
	}
	a.dec = &gosnmp.GoSNMP{Version: gosnmp.Version3, SecurityModel: gosnmp.UserSecurityModel}
	if a.hasV3 {
		if err := a.initUsers(); err != nil {
			_ = conn.Close() // best-effort: the init error is the one to report
			return nil, err
		}
	}
	a.ctx, a.cancel = context.WithCancel(ctx)
	a.wg.Add(1)
	go a.serve()
	context.AfterFunc(a.ctx, func() { _ = a.Close() }) // idempotent; the caller sees errors from its own Close
	return a, nil
}

// Addr returns the bound address.
func (a *Agent) Addr() netip.AddrPort { return a.addr }

// Close stops the agent and waits for in-flight requests.
func (a *Agent) Close() error {
	var err error
	a.once.Do(func() {
		a.cancel()
		err = a.conn.Close()
		a.wg.Wait()
	})
	return err
}

// Stats returns a snapshot of the request counters.
func (a *Agent) Stats() Stats {
	return Stats{
		Get: a.st.get.Load(), GetNext: a.st.next.Load(), GetBulk: a.st.bulk.Load(), Set: a.st.set.Load(),
		Discovery: a.st.disc.Load(), Dropped: a.st.dropped.Load(), Invalid: a.st.invalid.Load(),
		TooBig: a.st.tooBig.Load(), NotInTimeWindow: a.st.notInTime.Load(),
	}
}

// ResetStats zeroes the counters (and the DropEvery sequence).
func (a *Agent) ResetStats() {
	for _, c := range []*atomic.Uint64{&a.st.get, &a.st.next, &a.st.bulk, &a.st.set, &a.st.disc,
		&a.st.dropped, &a.st.invalid, &a.st.tooBig, &a.st.notInTime, &a.reqN} {
		c.Store(0)
	}
}

func (a *Agent) serve() {
	defer a.wg.Done()
	buf := make([]byte, 8192)
	for {
		n, from, err := a.conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			if a.ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
				return
			}
			continue
		}
		if a.cfg.Behaviour.Latency > 0 {
			raw := append([]byte(nil), buf[:n]...)
			a.wg.Add(1)
			go func() {
				defer a.wg.Done()
				a.handle(raw, from)
			}()
			continue
		}
		a.handle(buf[:n], from)
	}
}

func (a *Agent) handle(raw []byte, from netip.AddrPort) {
	req, err := a.dec.UnmarshalTrap(raw, false)
	if err != nil || !a.authorised(req) {
		a.st.invalid.Add(1)
		return
	}
	disc := req.Version == gosnmp.Version3 && req.MsgFlags&gosnmp.AuthNoPriv == 0
	if !a.count(req, disc) || a.injectFault() {
		return
	}
	out := a.respond(req, disc, a.recordingFor(req))
	if out == nil {
		return
	}
	// UDP is best-effort: a failed write looks like a dropped reply to the client, which is what it is.
	_, _ = a.conn.WriteToUDPAddrPort(out, from)
}

// authorised applies the community (v1/v2c) or USM-presence (v3) check.
func (a *Agent) authorised(req *gosnmp.SnmpPacket) bool {
	switch req.Version {
	case gosnmp.Version1, gosnmp.Version2c:
		base, _, _ := strings.Cut(req.Community, "@")
		return base == a.community
	case gosnmp.Version3:
		return a.hasV3 && req.SecurityModel == gosnmp.UserSecurityModel
	default:
		return false
	}
}

// recordingFor picks the recording a request addresses: the context's, or the default one.
func (a *Agent) recordingFor(req *gosnmp.SnmpPacket) *snmprec.Recording {
	name := ""
	switch req.Version {
	case gosnmp.Version3:
		name = strings.TrimPrefix(req.ContextName, "vlan-")
	default:
		_, name, _ = strings.Cut(req.Community, "@")
	}
	if rec, ok := a.cfg.Contexts[name]; ok && name != "" {
		return rec
	}
	return a.rec
}

// count bumps the per-PDU counter; false means the PDU type is not served.
func (a *Agent) count(req *gosnmp.SnmpPacket, disc bool) bool {
	switch {
	case disc:
		a.st.disc.Add(1)
	case req.PDUType == gosnmp.GetRequest:
		a.st.get.Add(1)
	case req.PDUType == gosnmp.GetNextRequest:
		a.st.next.Add(1)
	case req.PDUType == gosnmp.GetBulkRequest:
		a.st.bulk.Add(1)
	case req.PDUType == gosnmp.SetRequest:
		a.st.set.Add(1)
	default:
		return false
	}
	return true
}

// injectFault applies the configured drop and latency; true means no reply is sent.
func (a *Agent) injectFault() bool {
	if n := a.cfg.Behaviour.DropEvery; n > 0 && a.reqN.Add(1)%uint64(n) == 0 {
		a.st.dropped.Add(1)
		return true
	}
	if d := a.cfg.Behaviour.Latency; d > 0 && a.clk.Sleep(a.ctx, d) != nil {
		return true
	}
	return false
}

// respond builds the marshalled reply, or nil for none.
func (a *Agent) respond(req *gosnmp.SnmpPacket, disc bool, rec *snmprec.Recording) []byte {
	elapsed := a.clk.Now().Sub(a.start)
	limit := a.maxSize
	if req.Version == gosnmp.Version3 {
		if report, ok := a.v3Report(req, disc, elapsed); ok {
			return report
		}
		if m := int(req.MsgMaxSize); m > 0 && m < limit {
			limit = m
		}
	}
	resp := a.newResponse(req, elapsed)
	if out, done := a.refuse(req, resp); done {
		return out
	}
	a.fill(req, resp, elapsed, rec)
	return a.fit(req, resp, limit)
}

// v3Report answers discovery, unknown-engine and out-of-window v3 requests with
// a report PDU; ok is false when the request should be served normally.
func (a *Agent) v3Report(req *gosnmp.SnmpPacket, disc bool, elapsed time.Duration) (out []byte, ok bool) {
	if disc {
		return a.marshal(a.discoveryReport(req, elapsed)), true
	}
	if req.SecurityParameters.(*gosnmp.UsmSecurityParameters).AuthoritativeEngineID != a.engineID {
		return a.marshal(a.discoveryReport(req, elapsed)), true // RFC 3414 3.2.3: unknown engine ID
	}
	if !a.inTimeWindow(req, elapsed) {
		a.st.notInTime.Add(1)
		return a.marshal(a.timeReport(req, elapsed)), true
	}
	return nil, false
}

// refuse covers requests answered without looking at the data: SET is always
// notWritable, v1 has no GetBulk, and oversized GET/GETNEXT is tooBig.
func (a *Agent) refuse(req, resp *gosnmp.SnmpPacket) (out []byte, done bool) {
	switch req.PDUType {
	case gosnmp.SetRequest:
		resp.Variables = req.Variables
		resp.Error, resp.ErrorIndex = gosnmp.NotWritable, 1
		if req.Version == gosnmp.Version1 {
			resp.Error = gosnmp.NoSuchName
		}
		return a.marshal(resp), true
	case gosnmp.GetRequest, gosnmp.GetNextRequest:
		if n := a.cfg.Behaviour.TooBigAbove; n > 0 && len(req.Variables) > n {
			return a.tooBig(resp), true
		}
	case gosnmp.GetBulkRequest:
		if req.Version == gosnmp.Version1 {
			return nil, true
		}
	}
	return nil, false
}

// fit marshals resp, trimming GetBulk varbinds (or answering tooBig) until it
// fits the size limit.
func (a *Agent) fit(req, resp *gosnmp.SnmpPacket, limit int) []byte {
	out, err := resp.MarshalMsg()
	if err != nil {
		return nil
	}
	for len(out) > limit {
		if req.PDUType != gosnmp.GetBulkRequest || len(resp.Variables) < 2 {
			return a.tooBig(resp)
		}
		keep := len(resp.Variables) * limit / len(out)
		if keep >= len(resp.Variables) {
			keep = len(resp.Variables) - 1
		}
		resp.Variables = resp.Variables[:keep]
		if out, err = resp.MarshalMsg(); err != nil {
			return nil
		}
	}
	return out
}

func (a *Agent) marshal(p *gosnmp.SnmpPacket) []byte {
	out, err := p.MarshalMsg()
	if err != nil {
		return nil
	}
	return out
}

func (a *Agent) tooBig(resp *gosnmp.SnmpPacket) []byte {
	a.st.tooBig.Add(1)
	resp.Variables, resp.Error, resp.ErrorIndex = nil, gosnmp.TooBig, 0
	return a.marshal(resp)
}

// fill computes the response varbinds and error status for req.
func (a *Agent) fill(req, resp *gosnmp.SnmpPacket, elapsed time.Duration, rec *snmprec.Recording) {
	v1 := req.Version == gosnmp.Version1
	switch req.PDUType {
	case gosnmp.GetRequest:
		out := make([]gosnmp.SnmpPDU, 0, len(req.Variables))
		for i, vb := range req.Variables {
			oid := trimDot(vb.Name)
			r := exact(rec, oid, v1)
			if r == nil {
				if v1 {
					resp.Variables, resp.Error, resp.ErrorIndex = req.Variables, gosnmp.NoSuchName, uint8(i+1)
					return
				}
				out = append(out, gosnmp.SnmpPDU{Name: oid, Type: missing(rec, oid)})
				continue
			}
			out = append(out, pdu(r, elapsed))
		}
		resp.Variables = out
	case gosnmp.GetNextRequest:
		out := make([]gosnmp.SnmpPDU, 0, len(req.Variables))
		for i, vb := range req.Variables {
			oid := trimDot(vb.Name)
			r := next(rec, oid, v1)
			if r == nil {
				if v1 {
					resp.Variables, resp.Error, resp.ErrorIndex = req.Variables, gosnmp.NoSuchName, uint8(i+1)
					return
				}
				out = append(out, gosnmp.SnmpPDU{Name: oid, Type: gosnmp.EndOfMibView})
				continue
			}
			out = append(out, pdu(r, elapsed))
		}
		resp.Variables = out
	case gosnmp.GetBulkRequest:
		resp.Variables = a.bulk(req, elapsed, rec)
	}
}

func (a *Agent) bulk(req *gosnmp.SnmpPacket, elapsed time.Duration, rec *snmprec.Recording) []gosnmp.SnmpPDU {
	vbs := req.Variables
	nr := min(int(req.NonRepeaters), len(vbs))
	var out []gosnmp.SnmpPDU
	for _, vb := range vbs[:nr] {
		oid := trimDot(vb.Name)
		if r := next(rec, oid, false); r != nil {
			out = append(out, pdu(r, elapsed))
		} else {
			out = append(out, gosnmp.SnmpPDU{Name: oid, Type: gosnmp.EndOfMibView})
		}
	}
	cur := make([]string, 0, len(vbs)-nr)
	for _, vb := range vbs[nr:] {
		cur = append(cur, trimDot(vb.Name))
	}
	for rep := uint32(0); rep < req.MaxRepetitions && len(cur) > 0; rep++ {
		if len(out) > a.maxSize/8 {
			break
		}
		ended := 0
		for i, oid := range cur {
			r := next(rec, oid, false)
			if r == nil {
				out = append(out, gosnmp.SnmpPDU{Name: oid, Type: gosnmp.EndOfMibView})
				ended++
				continue
			}
			cur[i] = r.OID
			out = append(out, pdu(r, elapsed))
		}
		if ended == len(cur) {
			break
		}
	}
	return out
}

func trimDot(s string) string {
	if len(s) > 0 && s[0] == '.' {
		return s[1:]
	}
	return s
}

func pdu(r *snmprec.Record, elapsed time.Duration) gosnmp.SnmpPDU {
	return gosnmp.SnmpPDU{Name: r.OID, Type: r.Type, Value: r.At(elapsed)}
}

// visible reports whether r can be sent on the wire version (v1 has no Counter64).
func visible(r *snmprec.Record, v1 bool) bool { return !v1 || r.Type != gosnmp.Counter64 }

func exact(rec *snmprec.Recording, oid string, v1 bool) *snmprec.Record {
	recs := rec.Records()
	i := rec.Search(oid)
	if i < len(recs) && recs[i].OID == oid && visible(&recs[i], v1) {
		return &recs[i]
	}
	return nil
}

func next(rec *snmprec.Recording, oid string, v1 bool) *snmprec.Record {
	recs := rec.Records()
	i := rec.Search(oid)
	if i < len(recs) && recs[i].OID == oid {
		i++
	}
	for ; i < len(recs); i++ {
		if visible(&recs[i], v1) {
			return &recs[i]
		}
	}
	return nil
}

// missing classifies an absent OID the way net-snmp does: noSuchInstance when
// the object exists but not this instance (a sibling instance or a direct
// child exists), noSuchObject otherwise.
func missing(rec *snmprec.Recording, oid string) gosnmp.Asn1BER {
	recs := rec.Records()
	i := rec.Search(oid)
	arcs := strings.Count(oid, ".") + 1
	parent := oid
	if j := strings.LastIndexByte(oid, '.'); j >= 0 {
		parent = oid[:j]
	}
	for _, k := range []int{i - 1, i} {
		if k < 0 || k >= len(recs) {
			continue
		}
		r := recs[k].OID
		n := strings.Count(r, ".") + 1
		if snmprec.HasOIDPrefix(r, parent) && n == arcs || snmprec.HasOIDPrefix(r, oid) && n == arcs+1 {
			return gosnmp.NoSuchInstance
		}
	}
	return gosnmp.NoSuchObject
}

func (a *Agent) newResponse(req *gosnmp.SnmpPacket, elapsed time.Duration) *gosnmp.SnmpPacket {
	p := &gosnmp.SnmpPacket{
		Version: req.Version, Community: req.Community,
		PDUType: gosnmp.GetResponse, RequestID: req.RequestID,
	}
	if req.Version == gosnmp.Version3 {
		p.MsgID = req.MsgID
		p.MsgFlags = req.MsgFlags & gosnmp.AuthPriv
		p.SecurityModel = gosnmp.UserSecurityModel
		p.ContextEngineID = a.engineID
		p.ContextName = req.ContextName
		p.SecurityParameters = a.usmFor(req, elapsed, true)
	}
	return p
}