fix: track JetBrains connections (#10968)

* feat: implement jetbrains agentssh tracking

Based on tcp forwarding instead of ssh connections

* Add JetBrains tracking to bottom bar

agent/agent_test.go

@@ -1,6 +1,7 @@
 package agent_test
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
@@ -152,7 +153,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, expected, strings.TrimSpace(string(output)))
 	})
-	t.Run("Tracks", func(t *testing.T) {
+	t.Run("TracksVSCode", func(t *testing.T) {
 		t.Parallel()
 		if runtime.GOOS == "window" {
 			t.Skip("Sleeping for infinity doesn't work on Windows")
@@ -191,6 +192,77 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		err = session.Wait()
 		require.NoError(t, err)
 	})
+
+	t.Run("TracksJetBrains", 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)
+
+		// JetBrains tracking works by looking at the process name listening on the
+		// forwarded port.  If the process's command line includes the magic string
+		// we are looking for, then we assume it is a JetBrains editor.  So when we
+		// connect to the port we must ensure the process includes that magic string
+		// to fool the agent into thinking this is JetBrains.  To do this we need to
+		// spawn an external process (in this case a simple echo server) so we can
+		// control the process name.  The -D here is just to mimic how Java options
+		// are set but is not necessary as the agent looks only for the magic
+		// string itself anywhere in the command.
+		_, b, _, ok := runtime.Caller(0)
+		require.True(t, ok)
+		dir := filepath.Join(filepath.Dir(b), "../scripts/echoserver/main.go")
+		echoServerCmd := exec.Command("go", "run", dir,
+			"-D", agentssh.MagicProcessCmdlineJetBrains)
+		stdout, err := echoServerCmd.StdoutPipe()
+		require.NoError(t, err)
+		err = echoServerCmd.Start()
+		require.NoError(t, err)
+		defer echoServerCmd.Process.Kill()
+
+		// The echo server prints its port as the first line.
+		sc := bufio.NewScanner(stdout)
+		sc.Scan()
+		remotePort := sc.Text()
+
+		//nolint:dogsled
+		conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
+		sshClient, err := conn.SSHClient(ctx)
+		require.NoError(t, err)
+
+		tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", remotePort))
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			// always close on failure of test
+			_ = conn.Close()
+			_ = tunneledConn.Close()
+		})
+
+		var s *agentsdk.Stats
+		require.Eventuallyf(t, func() bool {
+			var ok bool
+			s, ok = <-stats
+			return ok && s.ConnectionCount > 0 &&
+				s.SessionCountJetBrains == 1
+		}, testutil.WaitLong, testutil.IntervalFast,
+			"never saw stats with conn open: %+v", s,
+		)
+
+		// Kill the server and connection after checking for the echo.
+		requireEcho(t, tunneledConn)
+		_ = echoServerCmd.Process.Kill()
+		_ = tunneledConn.Close()
+
+		require.Eventuallyf(t, func() bool {
+			var ok bool
+			s, ok = <-stats
+			return ok && s.ConnectionCount == 0 &&
+				s.SessionCountJetBrains == 0
+		}, testutil.WaitLong, testutil.IntervalFast,
+			"never saw stats after conn closes: %+v", s,
+		)
+	})
 }
 
 func TestAgent_SessionExec(t *testing.T) {

agent/agentssh/agentssh.go

@@ -47,8 +47,12 @@ const (
 	MagicSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
 	// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
 	MagicSessionTypeVSCode = "vscode"
-	// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains extension to identify itself.
+	// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
+	// extension to identify itself.
 	MagicSessionTypeJetBrains = "jetbrains"
+	// MagicProcessCmdlineJetBrains is a string in a process's command line that
+	// uniquely identifies it as JetBrains software.
+	MagicProcessCmdlineJetBrains = "idea.vendor.name=JetBrains"
 )
 
 type Server struct {
@@ -111,7 +115,11 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
 
 	srv := &ssh.Server{
 		ChannelHandlers: map[string]ssh.ChannelHandler{
-			"direct-tcpip":                   ssh.DirectTCPIPHandler,
+			"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 := NewJetbrainsChannelWatcher(ctx, s.logger, newChan, &s.connCountJetBrains)
+				ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
+			},
 			"direct-streamlocal@openssh.com": directStreamLocalHandler,
 			"session":                        ssh.DefaultSessionHandler,
 		},
@@ -291,8 +299,8 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv
 		s.connCountVSCode.Add(1)
 		defer s.connCountVSCode.Add(-1)
 	case MagicSessionTypeJetBrains:
-		s.connCountJetBrains.Add(1)
-		defer s.connCountJetBrains.Add(-1)
+		// Do nothing here because JetBrains launches hundreds of ssh sessions.
+		// We instead track JetBrains in the single persistent tcp forwarding channel.
 	case "":
 		s.connCountSSHSession.Add(1)
 		defer s.connCountSSHSession.Add(-1)

agent/agentssh/jetbrainstrack.go

@@ -0,0 +1,89 @@
+package agentssh
+
+import (
+	"strings"
+	"sync"
+
+	"cdr.dev/slog"
+	"github.com/gliderlabs/ssh"
+	"go.uber.org/atomic"
+	gossh "golang.org/x/crypto/ssh"
+)
+
+// localForwardChannelData is copied from the ssh package.
+type localForwardChannelData struct {
+	DestAddr string
+	DestPort uint32
+
+	OriginAddr string
+	OriginPort uint32
+}
+
+// JetbrainsChannelWatcher is used to track JetBrains port forwarded (Gateway)
+// channels. If the port forward is something other than JetBrains, this struct
+// is a noop.
+type JetbrainsChannelWatcher struct {
+	gossh.NewChannel
+	jetbrainsCounter *atomic.Int64
+}
+
+func NewJetbrainsChannelWatcher(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.
+		logger.Warn(ctx, "failed to unmarshal port forward data", slog.Error(err))
+		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, "failed to inspect port",
+			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.
+	if !strings.Contains(strings.ToLower(cmdline), strings.ToLower(MagicProcessCmdlineJetBrains)) {
+		return newChannel
+	}
+
+	logger.Debug(ctx, "discovered forwarded JetBrains process",
+		slog.F("destination_port", d.DestPort))
+
+	return &JetbrainsChannelWatcher{
+		NewChannel:       newChannel,
+		jetbrainsCounter: counter,
+	}
+}
+
+func (w *JetbrainsChannelWatcher) Accept() (gossh.Channel, <-chan *gossh.Request, error) {
+	c, r, err := w.NewChannel.Accept()
+	if err != nil {
+		return c, r, err
+	}
+	w.jetbrainsCounter.Add(1)
+
+	return &ChannelOnClose{
+		Channel: c,
+		done: func() {
+			w.jetbrainsCounter.Add(-1)
+		},
+	}, r, err
+}
+
+type ChannelOnClose struct {
+	gossh.Channel
+	// once ensures close only decrements the counter once.
+	// Because close can be called multiple times.
+	once sync.Once
+	done func()
+}
+
+func (c *ChannelOnClose) Close() error {
+	c.once.Do(c.done)
+	return c.Channel.Close()
+}

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
+}

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
+}