From 8d647c4ebf1f62efc96fa62ee1774981c4a50585 Mon Sep 17 00:00:00 2001
From: Daniel Nelson <daniel.nelson@influxdb.com>
Date: Tue, 14 Jan 2020 16:26:44 -0800
Subject: [PATCH] Fix URL agent address form with udp in snmp input (#6899)

---
 plugins/inputs/snmp/snmp.go      | 31 +++++++++++++++++++++----------
 plugins/inputs/snmp/snmp_test.go |  9 ++++++++-
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/plugins/inputs/snmp/snmp.go b/plugins/inputs/snmp/snmp.go
index 75c9b7836..2fc56ff97 100644
--- a/plugins/inputs/snmp/snmp.go
+++ b/plugins/inputs/snmp/snmp.go
@@ -7,6 +7,7 @@ import (
 	"log"
 	"math"
 	"net"
+	"net/url"
 	"os/exec"
 	"strconv"
 	"strings"
@@ -609,20 +610,30 @@ func (s *Snmp) getConnection(idx int) (snmpConnection, error) {
 	gs := gosnmpWrapper{&gosnmp.GoSNMP{}}
 	s.connectionCache[idx] = gs
 
-	if strings.HasPrefix(agent, "tcp://") {
-		agent = strings.TrimPrefix(agent, "tcp://")
-		gs.Transport = "tcp"
+	if !strings.Contains(agent, "://") {
+		agent = "udp://" + agent
 	}
-	host, portStr, err := net.SplitHostPort(agent)
+
+	u, err := url.Parse(agent)
 	if err != nil {
-		if err, ok := err.(*net.AddrError); !ok || err.Err != "missing port in address" {
-			return nil, Errorf(err, "parsing host")
-		}
-		host = agent
+		return nil, err
+	}
+
+	switch u.Scheme {
+	case "tcp":
+		gs.Transport = "tcp"
+	case "", "udp":
+		gs.Transport = "udp"
+	default:
+		return nil, fmt.Errorf("unsupported scheme: %v", u.Scheme)
+	}
+
+	gs.Target = u.Hostname()
+
+	portStr := u.Port()
+	if portStr == "" {
 		portStr = "161"
 	}
-	gs.Target = host
-
 	port, err := strconv.ParseUint(portStr, 10, 16)
 	if err != nil {
 		return nil, Errorf(err, "parsing port")
diff --git a/plugins/inputs/snmp/snmp_test.go b/plugins/inputs/snmp/snmp_test.go
index 3e174e224..25382bd7d 100644
--- a/plugins/inputs/snmp/snmp_test.go
+++ b/plugins/inputs/snmp/snmp_test.go
@@ -232,7 +232,7 @@ func TestSnmpInit_noTranslate(t *testing.T) {
 
 func TestGetSNMPConnection_v2(t *testing.T) {
 	s := &Snmp{
-		Agents:    []string{"1.2.3.4:567", "1.2.3.4"},
+		Agents:    []string{"1.2.3.4:567", "1.2.3.4", "udp://127.0.0.1"},
 		Timeout:   internal.Duration{Duration: 3 * time.Second},
 		Retries:   4,
 		Version:   2,
@@ -256,6 +256,13 @@ func TestGetSNMPConnection_v2(t *testing.T) {
 	assert.Equal(t, "1.2.3.4", gs.Target)
 	assert.EqualValues(t, 161, gs.Port)
 	assert.Equal(t, "udp", gs.Transport)
+
+	gsc, err = s.getConnection(2)
+	require.NoError(t, err)
+	gs = gsc.(gosnmpWrapper)
+	assert.Equal(t, "127.0.0.1", gs.Target)
+	assert.EqualValues(t, 161, gs.Port)
+	assert.Equal(t, "udp", gs.Transport)
 }
 
 func TestGetSNMPConnectionTCP(t *testing.T) {