Implement port process inspection

coder · code-asher · Dec 7, 2023 · Nov 30, 2023 · Nov 30, 2023 · Nov 30, 2023
commit adf2fb37a942cb7b1885ae692262103e79355517
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,35 @@ 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 invocation.
+	cmdline, err := getListeningPortProcessCmdline(d.DestPort)
+	if err != nil {
+		logger.Warn(ctx, "port inspection failed",
+			slog.F("destination_port", d.DestPort),
+			slog.Error(err))
+		return newChannel
+	}
+	logger.Debug(ctx, "checking forwarded process",
+		slog.F("cmdline", cmdline),
+		slog.F("destination_port", d.DestPort))
+
+	// 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(cmdline), "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,31 @@
+//go:build linux
+
+package agentssh
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/cakturk/go-netstat/netstat"
+	"golang.org/x/xerrors"
+)
+
+func getListeningPortProcessCmdline(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 "", xerrors.Errorf("inspect port %d: %w", port, err)
+	}
+	if len(tabs) == 0 {
+		return "", nil
+	}
+	// The process name provided by go-netstat does not include the full command
+	// line so grab that instead.
+	pid := tabs[0].Process.Pid
+	data, err := os.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
+	if err != nil {
+		return "", xerrors.Errorf("read /proc/%d/cmdline: %w", pid, err)
+	}
+	return string(data), nil
+}
diff --git a/agent/agentssh/portinspection_unsupported.go b/agent/agentssh/portinspection_unsupported.go
@@ -0,0 +1,9 @@
+//go:build !linux
+
+package agentssh
+
+func getListeningPortProcessCmdline(port uint32) (string, error) {
+	// We are not worrying about other platforms at the moment because Gateway
+	// only supports Linux anyway.
+	return "", nil
+}