feat: implement jetbrains agentssh tracking

Emyrk · Emyrk · commit c9a226f96c5f · 2023-11-30T16:08:03.000-06:00
Based on tcp forwarding instead of ssh connections
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -111,7 +111,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 := NewChannelAcceptWatcher(s.logger, newChan, &s.connCountJetBrains)
+				ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
+			},
 			"direct-streamlocal@openssh.com": directStreamLocalHandler,
 			"session":                        ssh.DefaultSessionHandler,
 		},
diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
@@ -0,0 +1,63 @@
+package agentssh
+
+import (
+	"cdr.dev/slog"
+	"go.uber.org/atomic"
+	gossh "golang.org/x/crypto/ssh"
+)
+
+type localForwardChannelData struct {
+	DestAddr string
+	DestPort uint32
+
+	OriginAddr string
+	OriginPort uint32
+}
+
+type ChannelAcceptWatcher struct {
+	gossh.NewChannel
+	jetbrainsCounter *atomic.Int64
+}
+
+func NewChannelAcceptWatcher(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
+		return newChannel
+	}
+
+	//if !jetbrains {
+	// If this isn't jetbrains, then we don't need to do anything special.
+	//return newChannel
+	//}
+
+	return &ChannelAcceptWatcher{
+		NewChannel:       newChannel,
+		jetbrainsCounter: counter,
+	}
+}
+
+func (w *ChannelAcceptWatcher) 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
+	done func()
+}
+
+func (c *ChannelOnClose) Close() error {
+	c.done()
+	return c.Channel.Close()
+}
