diff --git a/snmp.go b/snmp.go index 1121919..a8b5b88 100644 --- a/snmp.go +++ b/snmp.go @@ -10,6 +10,8 @@ import ( "github.com/towerops-app/towerops-agent/pb" ) +const snmpMaxOIDsPerGet = 60 + // snmpQuerier abstracts SNMP operations for testability. type snmpQuerier interface { Get(oids []string) (*gosnmp.SnmpPacket, error) @@ -41,15 +43,24 @@ func executeSnmpJob(job *pb.AgentJob, resultCh chan<- *pb.SnmpResult) { } defer closeFn() - oidValues := make(map[string]string) + totalOIDs := 0 + for _, q := range job.Queries { + totalOIDs += len(q.Oids) + } + oidValues := make(map[string]string, totalOIDs) for _, q := range job.Queries { switch q.QueryType { case pb.QueryType_GET: - for _, oid := range q.Oids { - result, err := conn.Get([]string{oid}) + for i := 0; i < len(q.Oids); i += snmpMaxOIDsPerGet { + end := i + snmpMaxOIDsPerGet + if end > len(q.Oids) { + end = len(q.Oids) + } + batch := q.Oids[i:end] + result, err := conn.Get(batch) if err != nil { - slog.Warn("snmp get failed", "device", dev.Ip, "oid", oid, "error", err) + slog.Warn("snmp get failed", "device", dev.Ip, "oids", len(batch), "error", err) continue } for _, v := range result.Variables { diff --git a/snmp_test.go b/snmp_test.go index 7a40f75..ca846c9 100644 --- a/snmp_test.go +++ b/snmp_test.go @@ -606,6 +606,58 @@ func TestExecuteSnmpJob(t *testing.T) { }) } +func TestExecuteSnmpJobBatchesGets(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + var getCalls [][]string + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + getCalls = append(getCalls, oids) + var vars []gosnmp.SnmpPDU + for _, oid := range oids { + vars = append(vars, gosnmp.SnmpPDU{Name: oid, Type: gosnmp.Integer, Value: 42}) + } + return &gosnmp.SnmpPacket{Variables: vars}, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + // Create 150 OIDs — should be split into 3 batches of 60, 60, 30 + oids := make([]string, 150) + for i := range oids { + oids[i] = fmt.Sprintf(".1.3.6.1.2.1.1.%d.0", i) + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "batch-test", + DeviceId: "dev-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_GET, Oids: oids}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 150 { + t.Errorf("got %d oid values, want 150", len(result.OidValues)) + } + if len(getCalls) != 3 { + t.Errorf("got %d GET calls, want 3 (batches of 60)", len(getCalls)) + } + for i, call := range getCalls { + if i < 2 && len(call) != 60 { + t.Errorf("batch %d: got %d oids, want 60", i, len(call)) + } + if i == 2 && len(call) != 30 { + t.Errorf("batch %d: got %d oids, want 30", i, len(call)) + } + } +} + func TestExecuteCredentialTest(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.CredentialTestResult, 1)