opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-21 12:44:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent/agentssh): move parsing of magic session and create type (#16630)

Browse Source 
This change refactors the parsing of MagicSessionEnvs in the agentssh
package and moves the logic to an earlier stage. Also intoduces enums
for MagicSessionType.

Refs #15139
This commit is contained in:
Mathias Fredriksson
2025-02-19 22:18:31 +02:00
committed by GitHub
parent 570e42b7f7
commit 9f5ad23644
3 changed files with 92 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/agent_test.go
+2 -2
View File
@@ -138,7 +138,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable, agentssh.MagicSessionTypeVSCode)
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable, string(agentssh.MagicSessionTypeVSCode))
defer session.Close()
command := "sh -c 'echo $" + agentssh.MagicSessionTypeEnvironmentVariable + "'"
@@ -165,7 +165,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable, agentssh.MagicSessionTypeVSCode)
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable, string(agentssh.MagicSessionTypeVSCode))
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
agent/agentssh/agentssh.go
+85 -49
View File
@@ -26,6 +26,7 @@ import (
"github.com/spf13/afero"
"go.uber.org/atomic"
gossh "golang.org/x/crypto/ssh"
"golang.org/x/exp/slices"
"golang.org/x/xerrors"
"cdr.dev/slog"
@@ -42,14 +43,6 @@ const (
// unlikely to shadow other exit codes, which are typically 1, 2, 3, etc.
MagicSessionErrorCode = 229
// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
// This is stripped from any commands being executed, and is counted towards connection stats.
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 = "jetbrains"
// MagicProcessCmdlineJetBrains is a string in a process's command line that
// uniquely identifies it as JetBrains software.
MagicProcessCmdlineJetBrains = "idea.vendor.name=JetBrains"
@@ -60,6 +53,29 @@ const (
BlockedFileTransferErrorMessage = "File transfer has been disabled."
)
// MagicSessionType is a type that represents the type of session that is being
// established.
type MagicSessionType string
const (
// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
// This is stripped from any commands being executed, and is counted towards connection stats.
MagicSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
)
// MagicSessionType enums.
const (
// MagicSessionTypeUnknown means the session type could not be determined.
MagicSessionTypeUnknown MagicSessionType = "unknown"
// MagicSessionTypeSSH is the default session type.
MagicSessionTypeSSH MagicSessionType = "ssh"
// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
MagicSessionTypeVSCode MagicSessionType = "vscode"
// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
// extension to identify itself.
MagicSessionTypeJetBrains MagicSessionType = "jetbrains"
)
// BlockedFileTransferCommands contains a list of restricted file transfer commands.
var BlockedFileTransferCommands = []string{"nc", "rsync", "scp", "sftp"}
@@ -255,14 +271,42 @@ func (s *Server) ConnStats() ConnStats {
}
}
func extractMagicSessionType(env []string) (magicType MagicSessionType, rawType string, filteredEnv []string) {
for _, kv := range env {
if !strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable) {
continue
}
rawType = strings.TrimPrefix(kv, MagicSessionTypeEnvironmentVariable+"=")
// Keep going, we'll use the last instance of the env.
}
// Always force lowercase checking to be case-insensitive.
switch MagicSessionType(strings.ToLower(rawType)) {
case MagicSessionTypeVSCode:
magicType = MagicSessionTypeVSCode
case MagicSessionTypeJetBrains:
magicType = MagicSessionTypeJetBrains
case "", MagicSessionTypeSSH:
magicType = MagicSessionTypeSSH
default:
magicType = MagicSessionTypeUnknown
}
return magicType, rawType, slices.DeleteFunc(env, func(kv string) bool {
return strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable+"=")
})
}
func (s *Server) sessionHandler(session ssh.Session) {
ctx := session.Context()
id := uuid.New()
logger := s.logger.With(
slog.F("remote_addr", session.RemoteAddr()),
slog.F("local_addr", session.LocalAddr()),
// Assigning a random uuid for each session is useful for tracking
// logs for the same ssh session.
slog.F("id", uuid.NewString()),
slog.F("id", id.String()),
)
logger.Info(ctx, "handling ssh session")
@@ -274,16 +318,21 @@ func (s *Server) sessionHandler(session ssh.Session) {
}
defer s.trackSession(session, false)
extraEnv := make([]string, 0)
x11, hasX11 := session.X11()
if hasX11 {
display, handled := s.x11Handler(session.Context(), x11)
if !handled {
_ = session.Exit(1)
logger.Error(ctx, "x11 handler failed")
return
}
extraEnv = append(extraEnv, fmt.Sprintf("DISPLAY=localhost:%d.%d", display, x11.ScreenNumber))
env := session.Environ()
magicType, magicTypeRaw, env := extractMagicSessionType(env)
switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
// Do nothing here because JetBrains launches hundreds of ssh sessions.
// We instead track JetBrains in the single persistent tcp forwarding channel.
case MagicSessionTypeSSH:
s.connCountSSHSession.Add(1)
defer s.connCountSSHSession.Add(-1)
case MagicSessionTypeUnknown:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("raw_type", magicTypeRaw))
}
if s.fileTransferBlocked(session) {
@@ -309,7 +358,18 @@ func (s *Server) sessionHandler(session ssh.Session) {
return
}
err := s.sessionStart(logger, session, extraEnv)
x11, hasX11 := session.X11()
if hasX11 {
display, handled := s.x11Handler(session.Context(), x11)
if !handled {
_ = session.Exit(1)
logger.Error(ctx, "x11 handler failed")
return
}
env = append(env, fmt.Sprintf("DISPLAY=localhost:%d.%d", display, x11.ScreenNumber))
}
err := s.sessionStart(logger, session, env, magicType)
var exitError *exec.ExitError
if xerrors.As(err, &exitError) {
code := exitError.ExitCode()
@@ -379,32 +439,8 @@ func (s *Server) fileTransferBlocked(session ssh.Session) bool {
return false
}
func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv []string) (retErr error) {
func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, env []string, magicType MagicSessionType) (retErr error) {
ctx := session.Context()
env := append(session.Environ(), extraEnv...)
var magicType string
for index, kv := range env {
if !strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable) {
continue
}
magicType = strings.ToLower(strings.TrimPrefix(kv, MagicSessionTypeEnvironmentVariable+"="))
env = append(env[:index], env[index+1:]...)
}
// Always force lowercase checking to be case-insensitive.
switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
// 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)
default:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
}
magicTypeLabel := magicTypeMetricLabel(magicType)
sshPty, windowSize, isPty := session.Pty()
@@ -473,7 +509,7 @@ func (s *Server) startNonPTYSession(logger slog.Logger, session ssh.Session, mag
}()
go func() {
for sig := range sigs {
s.handleSignal(logger, sig, cmd.Process, magicTypeLabel)
handleSignal(logger, sig, cmd.Process, s.metrics, magicTypeLabel)
}
}()
return cmd.Wait()
@@ -558,7 +594,7 @@ func (s *Server) startPTYSession(logger slog.Logger, session ptySession, magicTy
sigs = nil
continue
}
s.handleSignal(logger, sig, process, magicTypeLabel)
handleSignal(logger, sig, process, s.metrics, magicTypeLabel)
case win, ok := <-windowSize:
if !ok {
windowSize = nil
@@ -612,7 +648,7 @@ func (s *Server) startPTYSession(logger slog.Logger, session ptySession, magicTy
return nil
}
func (s *Server) handleSignal(logger slog.Logger, ssig ssh.Signal, signaler interface{ Signal(os.Signal) error }, magicTypeLabel string) {
func handleSignal(logger slog.Logger, ssig ssh.Signal, signaler interface{ Signal(os.Signal) error }, metrics *sshServerMetrics, magicTypeLabel string) {
ctx := context.Background()
sig := osSignalFrom(ssig)
logger = logger.With(slog.F("ssh_signal", ssig), slog.F("signal", sig.String()))
@@ -620,7 +656,7 @@ func (s *Server) handleSignal(logger slog.Logger, ssig ssh.Signal, signaler inte
err := signaler.Signal(sig)
if err != nil {
logger.Warn(ctx, "signaling the process failed", slog.Error(err))
s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, "yes", "signal").Add(1)
metrics.sessionErrors.WithLabelValues(magicTypeLabel, "yes", "signal").Add(1)
}
}
agent/agentssh/metrics.go
+5 -5
View File
@@ -71,15 +71,15 @@ func newSSHServerMetrics(registerer prometheus.Registerer) *sshServerMetrics {
}
}
func magicTypeMetricLabel(magicType string) string {
func magicTypeMetricLabel(magicType MagicSessionType) string {
switch magicType {
case MagicSessionTypeVSCode:
case MagicSessionTypeJetBrains:
case "":
magicType = "ssh"
case MagicSessionTypeSSH:
case MagicSessionTypeUnknown:
default:
magicType = "unknown"
magicType = MagicSessionTypeUnknown
}
// Always be case insensitive
return strings.ToLower(magicType)
return strings.ToLower(string(magicType))
}