Implement port process inspection

code-asher · code-asher · commit e7f9516a92f3 · 2023-11-30T15:23:59.000-09:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -192,10 +192,13 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	// This test name being "Jetbrains" is required to be a certain string.
-	// It must match the regex check in the agent for Jetbrains.
-	t.Run("Jetbrains", func(t *testing.T) {
+	// This test name must contain the string checked for by the agent, since it
+	// looks for this string in the process name.
+	t.Run("TracksIdea.vendor.name=JetBrains", func(t *testing.T) {
 		t.Parallel()
+		if runtime.GOOS != "linux" {
+			t.Skip("JetBrains tracking is only supported on Linux")
+		}
 		ctx := testutil.Context(t, testutil.WaitLong)
 
 		rl, err := net.Listen("tcp", "127.0.0.1:0")
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -113,7 +113,7 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
 		ChannelHandlers: map[string]ssh.ChannelHandler{
 			"direct-tcpip": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
 				// wrapper is designed to find and track jetbrains gateway connections.
-				wrapped := NewChannelAcceptWatcher(s.logger, newChan, &s.connCountJetBrains)
+				wrapped := NewChannelAcceptWatcher(ctx, s.logger, newChan, &s.connCountJetBrains)
 				ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
 			},
 			"direct-streamlocal@openssh.com": directStreamLocalHandler,
diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
@@ -1,9 +1,11 @@
 package agentssh
 
 import (
+	"strings"
 	"sync"
 
 	"cdr.dev/slog"
+	"github.com/gliderlabs/ssh"
 	"go.uber.org/atomic"
 	gossh "golang.org/x/crypto/ssh"
 )
@@ -21,17 +23,32 @@ type ChannelAcceptWatcher struct {
 	jetbrainsCounter *atomic.Int64
 }
 
-func NewChannelAcceptWatcher(logger slog.Logger, newChannel gossh.NewChannel, counter *atomic.Int64) gossh.NewChannel {
+func NewChannelAcceptWatcher(ctx ssh.Context, logger slog.Logger, newChannel gossh.NewChannel, counter *atomic.Int64) gossh.NewChannel {
 	d := localForwardChannelData{}
 	if err := gossh.Unmarshal(newChannel.ExtraData(), &d); err != nil {
-		// If the data fails to unmarshal, do nothing
+		// If the data fails to unmarshal, do nothing.
 		return newChannel
 	}
 
-	//if !jetbrains {
-	// If this isn't jetbrains, then we don't need to do anything special.
-	//return newChannel
-	//}
+	// If we do get a port, we should be able to get the matching PID and from
+	// there look up the name.
+	name, err := getListeningPortProcessName(d.DestPort)
+	if err != nil {
+		logger.Warn(ctx, "port inspection failed",
+			slog.F("destination_port", d.DestPort),
+			slog.Error(err))
+		return newChannel
+	}
+
+	// If this is not JetBrains, then we do not need to do anything special.  We
+	// attempt to match on something that appears unique to JetBrains software and
+	// the vendor name flag seems like it might be a reasonable choice.
+	if strings.Contains(strings.ToLower(name), "idea.vendor.name=jetbrains") {
+		return newChannel
+	}
+
+	logger.Debug(ctx, "discovered forwarded JetBrains process",
+		slog.F("destination_port", d.DestPort))
 
 	return &ChannelAcceptWatcher{
 		NewChannel:       newChannel,
diff --git a/agent/agentssh/portinspection_supported.go b/agent/agentssh/portinspection_supported.go
@@ -0,0 +1,20 @@
+//go:build linux || (windows && amd64)
+
+package agentssh
+
+import (
+	"github.com/cakturk/go-netstat/netstat"
+)
+
+func getListeningPortProcessName(port uint32) (string, error) {
+	tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool {
+		return s.LocalAddr != nil && uint32(s.LocalAddr.Port) == port
+	})
+	if err != nil {
+		return "", err
+	}
+	if len(tabs) == 0 {
+		return "", nil
+	}
+	return tabs[0].Process.Name, nil
+}
diff --git a/agent/agentssh/portinspection_unsupported.go b/agent/agentssh/portinspection_unsupported.go
@@ -0,0 +1,9 @@
+//go:build !linux && !(windows && amd64)
+
+package agentssh
+
+func getListeningPortProcessName(port uint32) (string, error) {
+	// We are not worrying about other platforms at the moment because Gateway
+	// only supports Linux anyway.
+	return "", nil
+}
