Skip to content

fix: track JetBrains connections #10968

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Dec 7, 2023
Merged

fix: track JetBrains connections #10968

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
c9a226f
feat: implement jetbrains agentssh tracking
Emyrk Nov 30, 2023
2447970
Add unit test to confirm tracking
Emyrk Nov 30, 2023
76d3a24
implement unit test to verify jetbrains functionality
Emyrk Nov 30, 2023
adf2fb3
Implement port process inspection
code-asher Nov 30, 2023
ad034f2
Add JetBrains tracking to bottom bar
code-asher Dec 1, 2023
34b7c5e
Elaborate on process name check comment
code-asher Dec 1, 2023
dce56fd
Comment that localForwardChannelData is copied
code-asher Dec 1, 2023
7139448
Comment ChannelAccepterWatcher
code-asher Dec 1, 2023
6e8f235
Rename channel watcher to be specific to Jetbrains
code-asher Dec 1, 2023
4d65478
Log unmarshal failure
code-asher Dec 1, 2023
254a5b6
Add constant for JetBrains magic string
code-asher Dec 1, 2023
a75ed6c
Fix JetBrains tracking test
code-asher Dec 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
implement unit test to verify jetbrains functionality
  • Loading branch information
@Emyrk
Emyrk committed Nov 30, 2023
commit 76d3a24094f8146594d6cda68be1c72a43b42d8c
37 changes: 33 additions & 4 deletions agent/agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,12 +206,41 @@ func TestAgent_Stats_Magic(t *testing.T) {
remotePort := tcpAddr.Port
go echoOnce(t, rl)

sshClient := setupAgentSSHClient(ctx, t)
//nolint:dogsled
conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)

conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
require.NoError(t, err)
defer conn.Close()
requireEcho(t, conn)
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,
)

// Manually closing the connection
requireEcho(t, tunneledConn)
_ = rl.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,
)
})
}

Expand Down
7 changes: 6 additions & 1 deletion agent/agentssh/jetbrainstrack.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package agentssh

import (
"sync"

"cdr.dev/slog"
"go.uber.org/atomic"
gossh "golang.org/x/crypto/ssh"
Expand Down Expand Up @@ -54,10 +56,13 @@ func (w *ChannelAcceptWatcher) Accept() (gossh.Channel, <-chan *gossh.Request, e

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.done()
c.once.Do(c.done)
return c.Channel.Close()
}