From 378c570c06f678b7004520115df836c36f6fea35 Mon Sep 17 00:00:00 2001
From: Felix Maurer <felix@felix-maurer.de>
Date: Thu, 7 Nov 2019 20:56:51 +0200
Subject: [PATCH] Add support for SNMP over TCP (#5870)

---
 Gopkg.lock                       | 14 ++++++++++---
 Gopkg.toml                       |  2 +-
 plugins/inputs/snmp/README.md    |  2 +-
 plugins/inputs/snmp/snmp.go      |  4 ++++
 plugins/inputs/snmp/snmp_test.go | 34 ++++++++++++++++++++++++++++++++
 5 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/Gopkg.lock b/Gopkg.lock
index 410b9b284..78842d940 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -513,6 +513,14 @@
   revision = "636bf0302bc95575d69441b25a2603156ffdddf1"
   version = "v1.1.1"
 
+[[projects]]
+  digest = "1:530233672f656641b365f8efb38ed9fba80e420baff2ce87633813ab3755ed6d"
+  name = "github.com/golang/mock"
+  packages = ["gomock"]
+  pruneopts = ""
+  revision = "51421b967af1f557f93a59e0057aaf15ca02e29c"
+  version = "v1.2.0"
+
 [[projects]]
   digest = "1:f958a1c137db276e52f0b50efee41a1a389dcdded59a69711f3e872757dab34b"
   name = "github.com/golang/protobuf"
@@ -1104,12 +1112,12 @@
   version = "v1.0.5"
 
 [[projects]]
-  branch = "master"
-  digest = "1:4b0cabe65ca903a7b2a3e6272c5304eb788ce196d35ecb901c6563e5e7582443"
+  digest = "1:a1cb5e999ad98b9838147e11ed1bdb000e750ee8872e2e21c74d9464cc9110c0"
   name = "github.com/soniah/gosnmp"
   packages = ["."]
   pruneopts = ""
-  revision = "96b86229e9b3ffb4b954144cdc7f98fe3ee1003f"
+  revision = "40eae407a1f8cbbe3f3f14c57bde0b16db1cfe85"
+  version = "v1.22.0"
 
 [[projects]]
   branch = "master"
diff --git a/Gopkg.toml b/Gopkg.toml
index 3069cbf40..1bbeb962f 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -153,7 +153,7 @@
 
 [[constraint]]
   name = "github.com/soniah/gosnmp"
-  branch = "master"
+  version = "1.22.0"
 
 [[constraint]]
   name = "github.com/StackExchange/wmi"
diff --git a/plugins/inputs/snmp/README.md b/plugins/inputs/snmp/README.md
index dab28e9b0..a15e5ddb6 100644
--- a/plugins/inputs/snmp/README.md
+++ b/plugins/inputs/snmp/README.md
@@ -98,7 +98,7 @@ Resulting output:
 ### Config parameters
 
 * `agents`: Default: `[]`
-List of SNMP agents to connect to in the form of `IP[:PORT]`. If `:PORT` is unspecified, it defaults to `161`.
+List of SNMP agents to connect to in the form of `[tcp://]IP[:PORT]`. If `:PORT` is unspecified, it defaults to `161`. When using the optional prefix `tcp://`, SNMP over TCP will be used. Otherwise UDP is used as the transport protocol.
 
 * `version`: Default: `2`
 SNMP protocol version to use.
diff --git a/plugins/inputs/snmp/snmp.go b/plugins/inputs/snmp/snmp.go
index 18eed4e47..32968730e 100644
--- a/plugins/inputs/snmp/snmp.go
+++ b/plugins/inputs/snmp/snmp.go
@@ -623,6 +623,10 @@ 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"
+	}
 	host, portStr, err := net.SplitHostPort(agent)
 	if err != nil {
 		if err, ok := err.(*net.AddrError); !ok || err.Err != "missing port in address" {
diff --git a/plugins/inputs/snmp/snmp_test.go b/plugins/inputs/snmp/snmp_test.go
index db1a49605..efa426845 100644
--- a/plugins/inputs/snmp/snmp_test.go
+++ b/plugins/inputs/snmp/snmp_test.go
@@ -272,12 +272,46 @@ func TestGetSNMPConnection_v2(t *testing.T) {
 	assert.EqualValues(t, 567, gs.Port)
 	assert.Equal(t, gosnmp.Version2c, gs.Version)
 	assert.Equal(t, "foo", gs.Community)
+	assert.Equal(t, "udp", gs.Transport)
 
 	gsc, err = s.getConnection(1)
 	require.NoError(t, err)
 	gs = gsc.(gosnmpWrapper)
 	assert.Equal(t, "1.2.3.4", gs.Target)
 	assert.EqualValues(t, 161, gs.Port)
+	assert.Equal(t, "udp", gs.Transport)
+}
+
+func TestGetSNMPConnectionTCP(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go stubTCPServer(&wg)
+	wg.Wait()
+
+	s := &Snmp{
+		Agents: []string{"tcp://127.0.0.1:56789"},
+	}
+	err := s.init()
+	require.NoError(t, err)
+
+	wg.Add(1)
+	gsc, err := s.getConnection(0)
+	require.NoError(t, err)
+	gs := gsc.(gosnmpWrapper)
+	assert.Equal(t, "127.0.0.1", gs.Target)
+	assert.EqualValues(t, 56789, gs.Port)
+	assert.Equal(t, "tcp", gs.Transport)
+	wg.Wait()
+}
+
+func stubTCPServer(wg *sync.WaitGroup) {
+	defer wg.Done()
+	tcpAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:56789")
+	tcpServer, _ := net.ListenTCP("tcp", tcpAddr)
+	defer tcpServer.Close()
+	wg.Done()
+	conn, _ := tcpServer.AcceptTCP()
+	defer conn.Close()
 }
 
 func TestGetSNMPConnection_v3(t *testing.T) {