Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit20fd687

Browse files
mafredriaslilac
authored andcommitted
refactor(agent/agentssh): move parsing of magic session and create type (#16630)
This change refactors the parsing of MagicSessionEnvs in the agentsshpackage and moves the logic to an earlier stage. Also intoduces enumsfor MagicSessionType.Refs#15139
1 parent794fd28 commit20fd687

File tree

3 files changed

+92
-56
lines changed

3 files changed

+92
-56
lines changed

‎agent/agent_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
138138
defersshClient.Close()
139139
session,err:=sshClient.NewSession()
140140
require.NoError(t,err)
141-
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable,agentssh.MagicSessionTypeVSCode)
141+
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable,string(agentssh.MagicSessionTypeVSCode))
142142
defersession.Close()
143143

144144
command:="sh -c 'echo $"+agentssh.MagicSessionTypeEnvironmentVariable+"'"
@@ -165,7 +165,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
165165
defersshClient.Close()
166166
session,err:=sshClient.NewSession()
167167
require.NoError(t,err)
168-
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable,agentssh.MagicSessionTypeVSCode)
168+
session.Setenv(agentssh.MagicSessionTypeEnvironmentVariable,string(agentssh.MagicSessionTypeVSCode))
169169
defersession.Close()
170170
stdin,err:=session.StdinPipe()
171171
require.NoError(t,err)

‎agent/agentssh/agentssh.go

Lines changed: 85 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"github.com/spf13/afero"
2727
"go.uber.org/atomic"
2828
gossh"golang.org/x/crypto/ssh"
29+
"golang.org/x/exp/slices"
2930
"golang.org/x/xerrors"
3031

3132
"cdr.dev/slog"
@@ -42,14 +43,6 @@ const (
4243
// unlikely to shadow other exit codes, which are typically 1, 2, 3, etc.
4344
MagicSessionErrorCode=229
4445

45-
// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
46-
// This is stripped from any commands being executed, and is counted towards connection stats.
47-
MagicSessionTypeEnvironmentVariable="CODER_SSH_SESSION_TYPE"
48-
// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
49-
MagicSessionTypeVSCode="vscode"
50-
// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
51-
// extension to identify itself.
52-
MagicSessionTypeJetBrains="jetbrains"
5346
// MagicProcessCmdlineJetBrains is a string in a process's command line that
5447
// uniquely identifies it as JetBrains software.
5548
MagicProcessCmdlineJetBrains="idea.vendor.name=JetBrains"
@@ -60,6 +53,29 @@ const (
6053
BlockedFileTransferErrorMessage="File transfer has been disabled."
6154
)
6255

56+
// MagicSessionType is a type that represents the type of session that is being
57+
// established.
58+
typeMagicSessionTypestring
59+
60+
const (
61+
// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
62+
// This is stripped from any commands being executed, and is counted towards connection stats.
63+
MagicSessionTypeEnvironmentVariable="CODER_SSH_SESSION_TYPE"
64+
)
65+
66+
// MagicSessionType enums.
67+
const (
68+
// MagicSessionTypeUnknown means the session type could not be determined.
69+
MagicSessionTypeUnknownMagicSessionType="unknown"
70+
// MagicSessionTypeSSH is the default session type.
71+
MagicSessionTypeSSHMagicSessionType="ssh"
72+
// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
73+
MagicSessionTypeVSCodeMagicSessionType="vscode"
74+
// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
75+
// extension to identify itself.
76+
MagicSessionTypeJetBrainsMagicSessionType="jetbrains"
77+
)
78+
6379
// BlockedFileTransferCommands contains a list of restricted file transfer commands.
6480
varBlockedFileTransferCommands= []string{"nc","rsync","scp","sftp"}
6581

@@ -255,14 +271,42 @@ func (s *Server) ConnStats() ConnStats {
255271
}
256272
}
257273

274+
funcextractMagicSessionType(env []string) (magicTypeMagicSessionType,rawTypestring,filteredEnv []string) {
275+
for_,kv:=rangeenv {
276+
if!strings.HasPrefix(kv,MagicSessionTypeEnvironmentVariable) {
277+
continue
278+
}
279+
280+
rawType=strings.TrimPrefix(kv,MagicSessionTypeEnvironmentVariable+"=")
281+
// Keep going, we'll use the last instance of the env.
282+
}
283+
284+
// Always force lowercase checking to be case-insensitive.
285+
switchMagicSessionType(strings.ToLower(rawType)) {
286+
caseMagicSessionTypeVSCode:
287+
magicType=MagicSessionTypeVSCode
288+
caseMagicSessionTypeJetBrains:
289+
magicType=MagicSessionTypeJetBrains
290+
case"",MagicSessionTypeSSH:
291+
magicType=MagicSessionTypeSSH
292+
default:
293+
magicType=MagicSessionTypeUnknown
294+
}
295+
296+
returnmagicType,rawType,slices.DeleteFunc(env,func(kvstring)bool {
297+
returnstrings.HasPrefix(kv,MagicSessionTypeEnvironmentVariable+"=")
298+
})
299+
}
300+
258301
func (s*Server)sessionHandler(session ssh.Session) {
259302
ctx:=session.Context()
303+
id:=uuid.New()
260304
logger:=s.logger.With(
261305
slog.F("remote_addr",session.RemoteAddr()),
262306
slog.F("local_addr",session.LocalAddr()),
263307
// Assigning a random uuid for each session is useful for tracking
264308
// logs for the same ssh session.
265-
slog.F("id",uuid.NewString()),
309+
slog.F("id",id.String()),
266310
)
267311
logger.Info(ctx,"handling ssh session")
268312

@@ -274,16 +318,21 @@ func (s *Server) sessionHandler(session ssh.Session) {
274318
}
275319
defers.trackSession(session,false)
276320

277-
extraEnv:=make([]string,0)
278-
x11,hasX11:=session.X11()
279-
ifhasX11 {
280-
display,handled:=s.x11Handler(session.Context(),x11)
281-
if!handled {
282-
_=session.Exit(1)
283-
logger.Error(ctx,"x11 handler failed")
284-
return
285-
}
286-
extraEnv=append(extraEnv,fmt.Sprintf("DISPLAY=localhost:%d.%d",display,x11.ScreenNumber))
321+
env:=session.Environ()
322+
magicType,magicTypeRaw,env:=extractMagicSessionType(env)
323+
324+
switchmagicType {
325+
caseMagicSessionTypeVSCode:
326+
s.connCountVSCode.Add(1)
327+
defers.connCountVSCode.Add(-1)
328+
caseMagicSessionTypeJetBrains:
329+
// Do nothing here because JetBrains launches hundreds of ssh sessions.
330+
// We instead track JetBrains in the single persistent tcp forwarding channel.
331+
caseMagicSessionTypeSSH:
332+
s.connCountSSHSession.Add(1)
333+
defers.connCountSSHSession.Add(-1)
334+
caseMagicSessionTypeUnknown:
335+
logger.Warn(ctx,"invalid magic ssh session type specified",slog.F("raw_type",magicTypeRaw))
287336
}
288337

289338
ifs.fileTransferBlocked(session) {
@@ -309,7 +358,18 @@ func (s *Server) sessionHandler(session ssh.Session) {
309358
return
310359
}
311360

312-
err:=s.sessionStart(logger,session,extraEnv)
361+
x11,hasX11:=session.X11()
362+
ifhasX11 {
363+
display,handled:=s.x11Handler(session.Context(),x11)
364+
if!handled {
365+
_=session.Exit(1)
366+
logger.Error(ctx,"x11 handler failed")
367+
return
368+
}
369+
env=append(env,fmt.Sprintf("DISPLAY=localhost:%d.%d",display,x11.ScreenNumber))
370+
}
371+
372+
err:=s.sessionStart(logger,session,env,magicType)
313373
varexitError*exec.ExitError
314374
ifxerrors.As(err,&exitError) {
315375
code:=exitError.ExitCode()
@@ -379,32 +439,8 @@ func (s *Server) fileTransferBlocked(session ssh.Session) bool {
379439
returnfalse
380440
}
381441

382-
func (s*Server)sessionStart(logger slog.Logger,session ssh.Session,extraEnv []string) (retErrerror) {
442+
func (s*Server)sessionStart(logger slog.Logger,session ssh.Session,env []string,magicTypeMagicSessionType) (retErrerror) {
383443
ctx:=session.Context()
384-
env:=append(session.Environ(),extraEnv...)
385-
varmagicTypestring
386-
forindex,kv:=rangeenv {
387-
if!strings.HasPrefix(kv,MagicSessionTypeEnvironmentVariable) {
388-
continue
389-
}
390-
magicType=strings.ToLower(strings.TrimPrefix(kv,MagicSessionTypeEnvironmentVariable+"="))
391-
env=append(env[:index],env[index+1:]...)
392-
}
393-
394-
// Always force lowercase checking to be case-insensitive.
395-
switchmagicType {
396-
caseMagicSessionTypeVSCode:
397-
s.connCountVSCode.Add(1)
398-
defers.connCountVSCode.Add(-1)
399-
caseMagicSessionTypeJetBrains:
400-
// Do nothing here because JetBrains launches hundreds of ssh sessions.
401-
// We instead track JetBrains in the single persistent tcp forwarding channel.
402-
case"":
403-
s.connCountSSHSession.Add(1)
404-
defers.connCountSSHSession.Add(-1)
405-
default:
406-
logger.Warn(ctx,"invalid magic ssh session type specified",slog.F("type",magicType))
407-
}
408444

409445
magicTypeLabel:=magicTypeMetricLabel(magicType)
410446
sshPty,windowSize,isPty:=session.Pty()
@@ -473,7 +509,7 @@ func (s *Server) startNonPTYSession(logger slog.Logger, session ssh.Session, mag
473509
}()
474510
gofunc() {
475511
forsig:=rangesigs {
476-
s.handleSignal(logger,sig,cmd.Process,magicTypeLabel)
512+
handleSignal(logger,sig,cmd.Process,s.metrics,magicTypeLabel)
477513
}
478514
}()
479515
returncmd.Wait()
@@ -558,7 +594,7 @@ func (s *Server) startPTYSession(logger slog.Logger, session ptySession, magicTy
558594
sigs=nil
559595
continue
560596
}
561-
s.handleSignal(logger,sig,process,magicTypeLabel)
597+
handleSignal(logger,sig,process,s.metrics,magicTypeLabel)
562598
casewin,ok:=<-windowSize:
563599
if!ok {
564600
windowSize=nil
@@ -612,15 +648,15 @@ func (s *Server) startPTYSession(logger slog.Logger, session ptySession, magicTy
612648
returnnil
613649
}
614650

615-
func(s*Server)handleSignal(logger slog.Logger,ssig ssh.Signal,signalerinterface{Signal(os.Signal)error },magicTypeLabelstring) {
651+
funchandleSignal(logger slog.Logger,ssig ssh.Signal,signalerinterface{Signal(os.Signal)error },metrics*sshServerMetrics,magicTypeLabelstring) {
616652
ctx:=context.Background()
617653
sig:=osSignalFrom(ssig)
618654
logger=logger.With(slog.F("ssh_signal",ssig),slog.F("signal",sig.String()))
619655
logger.Info(ctx,"received signal from client")
620656
err:=signaler.Signal(sig)
621657
iferr!=nil {
622658
logger.Warn(ctx,"signaling the process failed",slog.Error(err))
623-
s.metrics.sessionErrors.WithLabelValues(magicTypeLabel,"yes","signal").Add(1)
659+
metrics.sessionErrors.WithLabelValues(magicTypeLabel,"yes","signal").Add(1)
624660
}
625661
}
626662

‎agent/agentssh/metrics.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -71,15 +71,15 @@ func newSSHServerMetrics(registerer prometheus.Registerer) *sshServerMetrics {
7171
}
7272
}
7373

74-
funcmagicTypeMetricLabel(magicTypestring)string {
74+
funcmagicTypeMetricLabel(magicTypeMagicSessionType)string {
7575
switchmagicType {
7676
caseMagicSessionTypeVSCode:
7777
caseMagicSessionTypeJetBrains:
78-
case"":
79-
magicType="ssh"
78+
caseMagicSessionTypeSSH:
79+
caseMagicSessionTypeUnknown:
8080
default:
81-
magicType="unknown"
81+
magicType=MagicSessionTypeUnknown
8282
}
8383
// Always be case insensitive
84-
returnstrings.ToLower(magicType)
84+
returnstrings.ToLower(string(magicType))
8585
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp