From f7d7ae9b817aec5f27365dbc74b1c93c7c3d82e6 Mon Sep 17 00:00:00 2001
From: Sebastian Widmer <widmer.sebastian@gmail.com>
Date: Sun, 12 Jan 2020 08:59:56 +0100
Subject: [PATCH] Add option to skip process lookup

---
 go.mod          |  2 ++
 netstat.go      | 23 ++++++++++++++++++++++-
 netstat_test.go | 26 ++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 7e036a4..c30c9f5 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/bastjan/netstat
 
 require github.com/google/go-cmp v0.2.0
+
+go 1.13
diff --git a/netstat.go b/netstat.go
index 0963bfa..4248c84 100644
--- a/netstat.go
+++ b/netstat.go
@@ -45,6 +45,14 @@ var (
 	UDP6 = &Protocol{"udp6", "net/udp6"}
 )
 
+type option int
+
+const (
+	// SkipProcessLookup skips the rather slow PID lookup.
+	// Exe, Cmdline, and Pid will be empty if called with this options.
+	SkipProcessLookup option = iota + 1
+)
+
 var (
 	procFdLinkParseType1 = regexp.MustCompile(`^socket:\[(\d+)\]$`)
 	procFdLinkParseType2 = regexp.MustCompile(`^\[0000\]:(\d+)$`)
@@ -52,10 +60,14 @@ var (
 
 // Connections queries the given /proc/net file and returns the found connections.
 // Returns an error if the /proc/net file can't be read.
-func (p *Protocol) Connections() ([]*Connection, error) {
+func (p *Protocol) Connections(options ...option) ([]*Connection, error) {
 	inodeToPid := make(chan map[uint64]int)
 
 	go func() {
+		if hasOption(options, SkipProcessLookup) {
+			inodeToPid <- make(map[uint64]int)
+			return
+		}
 		inodeToPid <- procFdInodeToPid()
 	}()
 
@@ -227,3 +239,12 @@ func procFdExtractInode(fdLinkTarget string) (inode uint64, found bool) {
 	inode, _ = strconv.ParseUint(match[1], 10, 64)
 	return inode, true
 }
+
+func hasOption(os []option, o option) bool {
+	for _, ot := range os {
+		if ot == o {
+			return true
+		}
+	}
+	return false
+}
diff --git a/netstat_test.go b/netstat_test.go
index 8cbc40d..31a244b 100644
--- a/netstat_test.go
+++ b/netstat_test.go
@@ -67,6 +67,20 @@ func TestConnections(t *testing.T) {
 	compareResult(netstat.TCP6, []*netstat.Connection{tcp6Connection})
 }
 
+func TestConnectionsOptionSkipPidLookup(t *testing.T) {
+	compareResult := func(p *netstat.Protocol, expected []*netstat.Connection) {
+		connections, err := p.Connections(netstat.SkipProcessLookup)
+		if err != nil {
+			t.Error("Connections() returned unexpected errors:", err)
+		}
+		if diff := cmp.Diff(connections, expected); diff != "" {
+			t.Error("Connections() returned connections differ from expected connections:\n", diff)
+		}
+	}
+	compareResult(netstat.TCP, []*netstat.Connection{stripProcessInfo(tcpConnection)})
+	compareResult(netstat.TCP6, []*netstat.Connection{stripProcessInfo(tcp6Connection)})
+}
+
 func TestConnectionsProcNetNotFound(t *testing.T) {
 	_, err := (&netstat.Protocol{RelPath: "./nothere"}).Connections()
 	expectError(t, err, "test/proc/nothere: no such file or directory", "Connections() should return an error if the proc file can't be found")
@@ -87,3 +101,15 @@ func expectError(t *testing.T, err error, expectedErr, nilMessage string) {
 	}
 	t.Error("Error message should contain filename and error.", "Expected:", expectedErr, "Got:", err.Error())
 }
+
+func stripProcessInfo(c *netstat.Connection) *netstat.Connection {
+	r := new(netstat.Connection)
+
+	*r = *c
+
+	r.Exe = ""
+	r.Cmdline = []string{}
+	r.Pid = 0
+
+	return r
+}