opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-10-16 14:48:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add client_session_id to ssh connections (#29701)

Browse Source 
Add a new /tcp endpoint on the agent that is used to pass along
the client session ID (with the existing baggage mechanisms). The agent
then can upgrade that into the desired connection type, based on the
port. For now that is only the standard SSH port and reconnecting pty.

The `ssh`, `ssh --stdio`, `ssh --stdio` over Coder Connect, `exp rpty`,
and web terminal paths are all updated to use this new endpoint, and
fall back to the old method when the endpoint is not available.

That ID is added to the log context, and is also used to populate the
session ID column in the connection logs table.

There are some logs related to a connection that will not be tagged with
the ID (like the "accepted connection" line) because they happen before
the http handler gets to run.
This commit is contained in:
Asher
2026-10-02 16:40:24 -08:00
committed by GitHub
parent c7da54cf03
commit cc0c868c55
24 changed files with 1403 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/agent.go
+38 -5
View File
@@ -128,10 +128,10 @@ type Options struct {
}
type Client interface {
// ConnectRPC211WithRole connects to the Agent API v2.11. The workspace
// ConnectRPC212WithRole connects to the Agent API v2.12. The workspace
// agent should use role "agent" to enable connection monitoring.
ConnectRPC211WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC212WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient212, tailnetproto.DRPCTailnetClient28, error,
)
tailnet.DERPMapRewriter
agentsdk.RefreshableSessionTokenProvider
@@ -1008,7 +1008,7 @@ func (a *agent) run() (retErr error) {
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs.
// We pass role "agent" to enable connection monitoring on the server, which tracks
// the agent's connectivity state (first_connected_at, last_connected_at, disconnected_at).
aAPI, tAPI, err := a.client.ConnectRPC211WithRole(a.hardCtx, "agent")
aAPI, tAPI, err := a.client.ConnectRPC212WithRole(a.hardCtx, "agent")
if err != nil {
return err
}
@@ -1838,9 +1838,42 @@ func (a *agent) createTailnet(
_ = apiListener.Close()
}
}()
upgrader := &httpUpgrader{
logger: a.logger,
addr: apiListener.Addr(),
}
sshUpgradeListener := upgrader.listen(workspacesdk.AgentStandardSSHPort)
ptyUpgradeListener := upgrader.listen(workspacesdk.AgentReconnectingPTYPort)
defer func() {
if err != nil {
_ = sshUpgradeListener.Close()
_ = ptyUpgradeListener.Close()
}
}()
if err = a.trackGoroutine(func() {
_ = a.sshServer.Serve(sshUpgradeListener)
}); err != nil {
return nil, err
}
if err = a.trackGoroutine(func() {
go func() {
select {
case <-ctx.Done():
case <-a.hardCtx.Done():
}
_ = ptyUpgradeListener.Close()
}()
_ = a.reconnectingPTYServer.Serve(a.gracefulCtx, a.hardCtx, ptyUpgradeListener)
}); err != nil {
return nil, err
}
if err = a.trackGoroutine(func() {
defer apiListener.Close()
apiHandler := a.apiHandler()
apiHandler := a.apiHandler(upgrader)
server := &http.Server{
BaseContext: func(net.Listener) context.Context { return ctx },
Handler: apiHandler,
agent/agent_test.go
+2 -2
View File
@@ -131,8 +131,8 @@ type eofClient struct {
*agenttest.Client
}
func (*eofClient) ConnectRPC211WithRole(context.Context, string) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
func (*eofClient) ConnectRPC212WithRole(context.Context, string) (
proto.DRPCAgentClient212, tailnetproto.DRPCTailnetClient28, error,
) {
return nil, nil, io.EOF
}
agent/agentssh/agentssh.go
+79 -20
View File
@@ -11,6 +11,7 @@ import (
"os"
"os/exec"
"path/filepath"
"reflect"
"runtime"
"slices"
"strings"
@@ -222,12 +223,20 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
ChannelHandlers: map[string]ssh.ChannelHandler{
"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, s.config.ConnectionReporter, newChan, s.startSession)
wrapped := NewJetbrainsChannelWatcher(ctx, s.logger,
s.config.ConnectionReporter,
newChan,
s.startSession,
clientSessionIDFromContext(ctx),
)
ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
},
"direct-streamlocal@openssh.com": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
if s.config.BlockLocalPortForwarding {
s.logger.Warn(ctx, "unix local port forward blocked")
s.logger.Warn(ctx, "unix local port forward blocked",
slog.F("remote_addr", conn.RemoteAddr()),
slog.F("local_addr", conn.LocalAddr()),
slog.F("client_session_id", clientSessionIDFromContext(ctx)))
_ = newChan.Reject(gossh.Prohibited, "local port forwarding is disabled")
return
}
@@ -235,19 +244,33 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
},
"session": ssh.DefaultSessionHandler,
},
ConnCallback: func(ctx ssh.Context, conn net.Conn) net.Conn {
ctx.SetValue(clientSessionIDContextKey{}, ClientSessionIDFromConn(conn))
return conn
},
ConnectionFailedCallback: func(conn net.Conn, err error) {
// The conn here is our conn wrapped in a ssh.serverConn, so unwrap to get
// the client session ID from the inner conn.
clientSessionID := ""
val := reflect.ValueOf(conn)
if val.Kind() == reflect.Ptr {
val = val.Elem()
}
if val.Kind() == reflect.Struct {
f := val.FieldByName("Conn")
if f.IsValid() && f.CanInterface() {
if inner, ok := f.Interface().(net.Conn); ok {
clientSessionID = ClientSessionIDFromConn(inner)
}
}
}
s.logger.Warn(ctx, "ssh connection failed",
slog.F("remote_addr", conn.RemoteAddr()),
slog.F("local_addr", conn.LocalAddr()),
slog.F("client_session_id", clientSessionID),
slog.Error(err))
metrics.failedConnectionsTotal.Add(1)
},
ConnectionCompleteCallback: func(conn *gossh.ServerConn, err error) {
s.logger.Info(ctx, "ssh connection complete",
slog.F("remote_addr", conn.RemoteAddr()),
slog.F("local_addr", conn.LocalAddr()),
slog.Error(err))
},
Handler: s.sessionHandler,
// HostSigners are intentionally empty, as the host key will
// be set before we start listening.
@@ -409,12 +432,16 @@ func extractContainerInfo(env []string) (container, containerUser string, filter
func (s *Server) sessionHandler(session ssh.Session) {
ctx := session.Context()
id := uuid.New()
clientSessionID := clientSessionIDFromContext(ctx)
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", id.String()),
// The client session ID tracks multiple SSH sessions over the lifetime of a
// single client (IDE) session, for debugging reconnects/disconnects.
slog.F("client_session_id", clientSessionID),
)
logger.Info(ctx, "handling ssh session")
@@ -455,10 +482,11 @@ func (s *Server) sessionHandler(session ssh.Session) {
reason := "unable to accept new session, server is closing"
// Report connection attempt even if we couldn't accept it.
connReporter := s.config.ConnectionReporter.Connect(proto.ConnectEvent{
ID: id,
Type: connectionType,
AppName: appName,
IP: remoteAddrString,
ID: id,
Type: connectionType,
AppName: appName,
IP: remoteAddrString,
ClientSessionID: clientSessionID,
})
defer connReporter.Disconnect(proto.DisconnectEvent{
Code: 1,
@@ -492,10 +520,11 @@ func (s *Server) sessionHandler(session ssh.Session) {
session = scr
connReporter := s.config.ConnectionReporter.Connect(proto.ConnectEvent{
ID: id,
Type: connectionType,
AppName: appName,
IP: remoteAddrString,
ID: id,
Type: connectionType,
AppName: appName,
IP: remoteAddrString,
ClientSessionID: clientSessionID,
})
defer func() {
logger.Info(ctx, "ssh session closed",
@@ -1066,8 +1095,10 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string,
return cmd, nil
}
// Serve starts the server to handle incoming connections on the provided listener.
// It returns an error if no host keys are set or if there is an issue accepting connections.
// Serve starts the server to handle incoming connections on the provided
// listener. It returns an error if no host keys are set or if there is an
// issue accepting connections. If the listener returns UpgradedConn, then the
// client session ID will be extracted from those connections.
func (s *Server) Serve(l net.Listener) (retErr error) {
// Ensure we're not mutating HostSigners as we're reading it.
s.mu.RLock()
@@ -1096,11 +1127,19 @@ func (s *Server) Serve(l net.Listener) (retErr error) {
}
}
// handleConn serves a single SSH connection. l is the listener from which the
// connection was accepted. If the conn is an UpgradedConn, then the client
// session ID will be extracted from it.
func (s *Server) handleConn(l net.Listener, c net.Conn) {
clientSessionID := ClientSessionIDFromConn(c)
logger := s.logger.With(
slog.F("remote_addr", c.RemoteAddr()),
slog.F("local_addr", c.LocalAddr()),
slog.F("listen_addr", l.Addr()))
slog.F("listen_addr", l.Addr()),
slog.F("client_session_id", clientSessionID))
defer logger.Info(context.Background(), "ssh connection complete")
defer c.Close()
if !s.trackConn(l, c, true) {
@@ -1111,7 +1150,6 @@ func (s *Server) handleConn(l net.Listener, c net.Conn) {
}
defer s.trackConn(l, c, false)
logger.Info(context.Background(), "started serving ssh connection")
// note: srv.ConnectionCompleteCallback logs completion of the connection
s.srv.HandleConn(c)
}
@@ -1402,3 +1440,24 @@ func CoderSigner(seed int64) (gossh.Signer, error) {
coderSigner, err := gossh.NewSignerFromKey(coderHostKey)
return coderSigner, err
}
type clientSessionIDContextKey struct{}
func clientSessionIDFromContext(ctx ssh.Context) string {
id, _ := ctx.Value(clientSessionIDContextKey{}).(string)
return id
}
// UpgradedConn is a net.Conn that has a client session ID attached. Listeners
// can return these to augment logs for a connection with a client session ID.
type UpgradedConn interface {
net.Conn
ClientSessionID() string
}
func ClientSessionIDFromConn(conn net.Conn) string {
if uc, ok := conn.(UpgradedConn); ok {
return uc.ClientSessionID()
}
return ""
}
agent/agentssh/agentssh_test.go
+125 -29
View File
@@ -1,4 +1,4 @@
// Package agentssh_test provides tests for basic functinoality of the agentssh
// Package agentssh_test provides tests for basic functionality of the agentssh
// package, more test coverage can be found in the `agent` and `cli` package(s).
package agentssh_test
@@ -28,6 +28,7 @@ import (
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/agent/agentexec"
"github.com/coder/coder/v2/agent/agentssh"
"github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/testutil"
"github.com/coder/coder/v2/testutil/expecter"
)
@@ -36,44 +37,139 @@ func TestMain(m *testing.M) {
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
}
type testConnectionReporter struct {
mu sync.Mutex
connects []proto.ConnectEvent
disconnects []proto.DisconnectEvent
}
func (r *testConnectionReporter) Connect(connectEvent proto.ConnectEvent) proto.DisconnectionReporter {
r.mu.Lock()
defer r.mu.Unlock()
r.connects = append(r.connects, connectEvent)
return &testDisconnectionReporter{
callback: func(disconnectEvent proto.DisconnectEvent) {
r.mu.Lock()
defer r.mu.Unlock()
r.disconnects = append(r.disconnects, disconnectEvent)
},
}
}
type testDisconnectionReporter struct {
callback func(proto.DisconnectEvent)
}
func (r *testDisconnectionReporter) Disconnect(disconnectEvent proto.DisconnectEvent) {
r.callback(disconnectEvent)
}
func TestNewServer_ServeClient(t *testing.T) {
t.Parallel()
ctx := context.Background()
logger := testutil.Logger(t)
s, err := agentssh.NewServer(ctx, logger, prometheus.NewRegistry(), afero.NewMemMapFs(), agentexec.DefaultExecer, nil)
require.NoError(t, err)
defer s.Close()
err = s.UpdateHostSigner(42)
assert.NoError(t, err)
tests := []struct {
name string
clientSessionID string
}{
{
name: "NoClientSessionID",
},
{
name: "WithClientSessionID",
clientSessionID: "0123456789abcdef0123456789abcdef",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
var (
ctx = context.Background()
logger = testutil.Logger(t)
)
done := make(chan struct{})
go func() {
defer close(done)
err := s.Serve(ln)
assert.Error(t, err) // Server is closed.
}()
reporter := &testConnectionReporter{}
c := sshClient(t, ln.Addr().String())
s, err := agentssh.NewServer(ctx, logger,
prometheus.NewRegistry(),
afero.NewMemMapFs(),
agentexec.DefaultExecer,
&agentssh.Config{
ConnectionReporter: reporter,
},
)
require.NoError(t, err)
defer s.Close()
err = s.UpdateHostSigner(42)
assert.NoError(t, err)
var b bytes.Buffer
sess, err := c.NewSession()
require.NoError(t, err)
sess.Stdout = &b
err = sess.Start("echo hello")
require.NoError(t, err)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
err = sess.Wait()
require.NoError(t, err)
if tc.clientSessionID != "" {
ln = &wrappedListener{
ln,
tc.clientSessionID,
}
}
require.Equal(t, "hello", strings.TrimSpace(b.String()))
done := make(chan struct{})
go func() {
defer close(done)
err := s.Serve(ln)
assert.Error(t, err) // Server is closed.
}()
err = s.Close()
require.NoError(t, err)
<-done
c := sshClient(t, ln.Addr().String())
var b bytes.Buffer
sess, err := c.NewSession()
require.NoError(t, err)
sess.Stdout = &b
err = sess.Start("echo hello")
require.NoError(t, err)
err = sess.Wait()
require.NoError(t, err)
require.Equal(t, "hello", strings.TrimSpace(b.String()))
err = s.Close()
require.NoError(t, err)
<-done
// The server will not close until the disconnect has been reported, so
// these reports should be populated and safe to access now.
require.Len(t, reporter.connects, 1)
require.Len(t, reporter.disconnects, 1)
// We only need to check the connect report's client session ID;
// disconnect inherit the connect report's ID via the closure, so they
// will always be the same.
require.Equal(t, tc.clientSessionID, reporter.connects[0].ClientSessionID, "connect client session id")
})
}
}
// wrappedListener wraps a net.Listener to augment all connections with the
// provided client session ID.
type wrappedListener struct {
net.Listener
clientSessionID string
}
type wrappedConn struct {
net.Conn
clientSessionID string
}
func (w wrappedConn) ClientSessionID() string { return w.clientSessionID }
func (ln *wrappedListener) Accept() (net.Conn, error) {
conn, err := ln.Listener.Accept()
return wrappedConn{
Conn: conn,
clientSessionID: ln.clientSessionID,
}, err
}
func TestNewServer_ExecuteShebang(t *testing.T) {
agent/agentssh/jetbrainstrack.go
+4
View File
@@ -41,7 +41,10 @@ type JetbrainsChannelWatcher struct {
func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger,
connectionReporter proto.ConnectionReporter,
newChannel gossh.NewChannel, startSession startSessionFunc,
clientSessionID string,
) gossh.NewChannel {
logger = logger.With(slog.F("client_session_id", clientSessionID))
d := localForwardChannelData{}
if err := gossh.Unmarshal(newChannel.ExtraData(), &d); err != nil {
// If the data fails to unmarshal, do nothing.
@@ -74,6 +77,7 @@ func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger,
logger: logger.With(slog.F("destination_port", d.DestPort)),
originAddr: d.OriginAddr,
connectionReporter: connectionReporter,
clientSessionID: clientSessionID,
}
}
agent/agenttest/client.go
+7
View File
@@ -185,6 +185,13 @@ func (c *Client) ConnectRPC211WithRole(ctx context.Context, _ string) (
return c.ConnectRPC210(ctx)
}
// v2.12 adds no RPCs, so the v2.11 client satisfies it as-is.
func (c *Client) ConnectRPC212WithRole(ctx context.Context, role string) (
agentproto.DRPCAgentClient212, proto.DRPCTailnetClient28, error,
) {
return c.ConnectRPC211WithRole(ctx, role)
}
func (c *Client) ConnectRPC29(ctx context.Context) (
agentproto.DRPCAgentClient29, proto.DRPCTailnetClient28, error,
) {
agent/api.go
+2 -1
View File
@@ -15,7 +15,7 @@ import (
"github.com/coder/coder/v2/httpmw"
)
func (a *agent) apiHandler() http.Handler {
func (a *agent) apiHandler(upgradeListener *httpUpgrader) http.Handler {
r := chi.NewRouter()
r.Use(
httpmw.Recover(a.logger),
@@ -68,6 +68,7 @@ func (a *agent) apiHandler() http.Handler {
r.Get("/api/v0/listening-ports", a.listeningPortsHandler.handler)
r.Get("/api/v0/netcheck", a.HandleNetcheck)
r.Get("/api/v0/tcp/{port}", upgradeListener.handler)
r.Get("/debug/logs", a.HandleHTTPDebugLogs)
r.Get("/debug/magicsock", a.HandleHTTPDebugMagicsock)
r.Get("/debug/magicsock/debug-logging/{state}", a.HandleHTTPMagicsockDebugLoggingState)
agent/proto/agent_drpc_old.go
+6
View File
@@ -105,3 +105,9 @@ type DRPCAgentClient210 interface {
type DRPCAgentClient211 interface {
DRPCAgentClient210
}
// DRPCAgentClient212 is the Agent API at v2.12. It adds the ability to track a
// client session id with connections. No new RPCs.
type DRPCAgentClient212 interface {
DRPCAgentClient211
}
agent/reconnectingpty/server.go
+7 -4
View File
@@ -69,9 +69,11 @@ func (s *Server) Serve(ctx, hardCtx context.Context, l net.Listener) (retErr err
retErr = err
break
}
clientSessionID := agentssh.ClientSessionIDFromConn(conn)
clog := s.logger.With(
slog.F("remote", conn.RemoteAddr()),
slog.F("local", conn.LocalAddr()))
slog.F("local", conn.LocalAddr()),
slog.F("client_session_id", clientSessionID))
clog.Info(ctx, "accepted conn")
// It's not safe to assume RemoteAddr() returns a non-nil value. slog.F usage is fine because it correctly
@@ -85,9 +87,10 @@ func (s *Server) Serve(ctx, hardCtx context.Context, l net.Listener) (retErr err
wg.Add(1)
connReporter := s.connectionReporter.Connect(proto.ConnectEvent{
ID: uuid.New(),
Type: proto.Connection_RECONNECTING_PTY,
IP: remoteAddrString,
ID: uuid.New(),
Type: proto.Connection_RECONNECTING_PTY,
IP: remoteAddrString,
ClientSessionID: clientSessionID,
})
closed := make(chan struct{})
go func() {
agent/upgrade.go
+160
View File
@@ -0,0 +1,160 @@
package agent
import (
"bufio"
"net"
"net/http"
"strconv"
"strings"
"sync"
"time"
"github.com/go-chi/chi/v5"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/coderd/util/syncmap"
"github.com/coder/coder/v2/codersdk"
)
type upgradedConn struct {
net.Conn
clientSessionID string
r *bufio.Reader
}
func (c *upgradedConn) Read(p []byte) (int, error) { return c.r.Read(p) }
func (c *upgradedConn) ClientSessionID() string { return c.clientSessionID }
type httpUpgradeListener struct {
mu sync.Mutex
addr net.Addr
conn chan *upgradedConn
closed chan struct{}
}
func (ln *httpUpgradeListener) Accept() (net.Conn, error) {
var c *upgradedConn
select {
case c = <-ln.conn:
case <-ln.closed:
return nil, xerrors.Errorf("tcp upgrade listener: %w", net.ErrClosed)
}
return c, nil
}
func (ln *httpUpgradeListener) Addr() net.Addr { return ln.addr }
func (ln *httpUpgradeListener) Close() error {
ln.mu.Lock()
defer ln.mu.Unlock()
select {
case <-ln.closed:
default:
close(ln.closed)
}
return nil
}
type httpUpgrader struct {
listeners syncmap.Map[uint16, *httpUpgradeListener]
logger slog.Logger
addr net.Addr
}
// listen listens for upgrades on the specified port.
func (up *httpUpgrader) listen(port uint16) net.Listener {
listener, _ := up.listeners.LoadOrStore(port, &httpUpgradeListener{
addr: up.addr,
closed: make(chan struct{}),
conn: make(chan *upgradedConn),
})
return listener
}
// handler upgrades an HTTP connection based on the port number, passing along
// any client session ID that exists on the request context.
func (up *httpUpgrader) handler(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
port, err := strconv.ParseUint(chi.URLParam(r, "port"), 10, 16)
if err != nil {
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
Message: "Invalid port",
Detail: err.Error(),
})
return
}
listener, ok := up.listeners.Load(uint16(port))
if !ok {
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
Message: "Unsupported port for upgrading",
})
return
}
if r.Header.Get("upgrade") == "" {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Missing upgrade header",
})
return
}
if !strings.EqualFold(r.Header.Get("upgrade"), "tcp") {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid upgrade header",
})
return
}
clientSessionID := tracing.ClientSessionID(r)
if clientSessionID == "" {
up.logger.Warn(ctx, "client session ID is missing")
}
hijacker, ok := rw.(http.Hijacker)
if !ok {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Hijacking not supported",
})
return
}
conn, brw, err := hijacker.Hijack()
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to upgrade connection",
Detail: err.Error(),
})
return
}
_, err = conn.Write([]byte(strings.Join([]string{
"HTTP/1.1 101 Switching Protocols",
"Connection: Upgrade",
"Upgrade: tcp",
}, "\r\n") + "\r\n\r\n"))
if err != nil {
_ = conn.Close()
return
}
uconn := &upgradedConn{
Conn: conn,
clientSessionID: clientSessionID,
r: brw.Reader,
}
t := time.NewTimer(30 * time.Second)
defer t.Stop()
select {
case listener.conn <- uconn:
return
case <-listener.closed:
up.logger.Info(ctx, "listener closed; closing connection")
case <-t.C:
up.logger.Info(ctx, "listener timed out accepting; closing connection")
}
_ = conn.Close()
}
agent/upgrade_test.go
+243
View File
@@ -0,0 +1,243 @@
package agent_test
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/testutil"
)
func TestAgent_TCPUpgrade_Error(t *testing.T) {
t.Parallel()
tests := []struct {
name string
port int
portStr string
upgrade string
error string
status int
}{
{
name: "MissingPort",
status: http.StatusNotFound,
},
{
name: "InvalidPort",
portStr: "invalid",
status: http.StatusNotFound,
error: "Invalid port",
},
{
name: "UnsupportedPort",
port: 9999,
status: http.StatusNotFound,
error: "Unsupported port",
},
{
name: "MissingUpgrade",
port: workspacesdk.AgentStandardSSHPort,
status: http.StatusBadRequest,
error: "Missing upgrade header",
},
{
name: "InvalidUpgrade",
upgrade: "invalid",
port: workspacesdk.AgentStandardSSHPort,
status: http.StatusBadRequest,
error: "Invalid upgrade header",
},
}
//nolint:dogsled
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
client := &http.Client{
Transport: &http.Transport{
DisableKeepAlives: true,
DialContext: agentConn.DialContext,
},
}
t.Cleanup(client.CloseIdleConnections)
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
port := tc.portStr
if port == "" && tc.port != 0 {
port = strconv.Itoa(tc.port)
}
url := fmt.Sprintf("http://127.0.0.1:%d/api/v0/tcp/%s",
workspacesdk.AgentHTTPAPIServerPort, port)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
if tc.upgrade != "" {
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Upgrade", tc.upgrade)
}
require.True(t, agentConn.AwaitReachable(ctx))
res, err := client.Do(req)
require.NoError(t, err)
defer res.Body.Close()
require.Equal(t, tc.status, res.StatusCode)
if tc.error != "" {
var apiErr codersdk.Response
require.NoError(t, json.NewDecoder(res.Body).Decode(&apiErr))
require.Contains(t, apiErr.Message, tc.error)
}
})
}
}
func TestAgent_SSHUpgrade(t *testing.T) {
t.Parallel()
tests := []struct {
name string
id string
expect string
}{
{
name: "EmptyClientSessionID",
},
{
name: "InvalidClientSessionID",
id: "invalid",
},
{
name: "ValidClientSessionID",
id: "0123456789abcdef0123456789abcdef",
expect: "0123456789abcdef0123456789abcdef",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, client, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
require.True(t, conn.AwaitReachable(ctx))
if tc.id != "" {
conn.SetExtraHeaders(http.Header{
"baggage": []string{tracing.SessionIDBaggageKey + "=" + tc.id},
})
}
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
defer sshClient.Close()
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
stdin, err := session.StdinPipe()
require.NoError(t, err)
err = session.Shell()
require.NoError(t, err)
// Generate SSH traffic so the connstats window sees the session.
_, err = stdin.Write([]byte("echo test\n"))
require.NoError(t, err)
assertSSHStats(t, stats)
_, err = stdin.Write([]byte("exit 0\n"))
require.NoError(t, err, "writing exit to stdin")
_ = stdin.Close()
err = session.Wait()
require.NoError(t, err, "waiting for session to exit")
assertConnectionReport(t, client,
proto.ConnectEvent{
Type: proto.Connection_SSH,
ClientSessionID: tc.expect,
},
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
)
})
}
}
func TestAgent_ReconnectingPTYUpgrade(t *testing.T) {
t.Parallel()
tests := []struct {
name string
id string
expect string
}{
{
name: "EmptyClientSessionID",
},
{
name: "InvalidClientSessionID",
id: "invalid",
},
{
name: "ValidClientSessionID",
id: "0123456789abcdef0123456789abcdef",
expect: "0123456789abcdef0123456789abcdef",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
require.True(t, conn.AwaitReachable(ctx))
if tc.id != "" {
conn.SetExtraHeaders(http.Header{
"baggage": []string{tracing.SessionIDBaggageKey + "=" + tc.id},
})
}
ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "bash")
require.NoError(t, err)
defer ptyConn.Close()
data, err := json.Marshal(workspacesdk.ReconnectingPTYRequest{
Data: "echo test\r\n",
})
require.NoError(t, err)
_, err = ptyConn.Write(data)
require.NoError(t, err)
err = ptyConn.Close()
require.NoError(t, err)
assertConnectionReport(t, client,
proto.ConnectEvent{
Type: proto.Connection_RECONNECTING_PTY,
ClientSessionID: tc.expect,
},
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
)
})
}
}
cli/configssh_test.go
+1 -1
View File
@@ -108,7 +108,7 @@ func TestConfigSSH(t *testing.T) {
break
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
ssh, err := agentConn.SSH(ctx)
ssh, err := agentConn.SSHTCPConn(ctx)
cancel()
assert.NoError(t, err)
wg.Add(2)
cli/exp_rpty.go
+1
View File
@@ -25,6 +25,7 @@ func (r *RootCmd) rptyCommand() *serpent.Command {
var args handleRPTYArgs
cmd := &serpent.Command{
Annotations: serpent.Annotations(workspaceCommand).Mark(annotationClientSessionID, ""),
Handler: func(inv *serpent.Invocation) error {
if r.disableDirect {
return xerrors.New("direct connections are disabled, but you can try websocat ;-)")
cli/exp_rpty_test.go
+61 -39
View File
@@ -15,6 +15,10 @@ import (
"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/connectionlog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
"github.com/coder/coder/v2/testutil/expecter"
)
@@ -22,50 +26,68 @@ import (
func TestExpRpty(t *testing.T) {
t.Parallel()
t.Run("DefaultCommand", func(t *testing.T) {
t.Parallel()
randstr := uuid.NewString()
client, workspace, agentToken := setupWorkspaceForAgent(t)
inv, root := clitest.New(t, "exp", "rpty", workspace.Name)
clitest.SetupConfig(t, client, root)
stdin := testutil.NewWriterAttachedToInvocation(t, testutil.Logger(t), inv)
tests := []struct {
name string // test name
args []string // args to append
stdin string // stdin to write
stdout string // expected stdout
id string // expected client session id
}{
{name: "DefaultCommand", stdin: "exit"},
{name: "Command", args: []string{"echo", randstr}, stdout: randstr},
{name: "ClientSessionID", stdin: "exit", id: "0123456789abcdef0123456789abcdef"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
connLogger := connectionlog.NewFake()
client, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
ConnectionLogger: connLogger,
})
client.SetLogger(testutil.Logger(t).Named("client"))
first := coderdtest.CreateFirstUser(t, client)
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
r.Username = "myuser"
})
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
Name: "myworkspace",
OrganizationID: first.OrganizationID,
OwnerID: user.ID,
}).WithAgent().Do()
_ = agenttest.New(t, client.URL, agentToken)
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
args := []string{"exp", "rpty", r.Workspace.Name}
args = append(args, tc.args...)
inv, root := clitest.New(t, args...)
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
clitest.SetupConfig(t, userClient, root)
if tc.id != "" {
inv.Environ.Set("CODER_TRACE_SESSION_ID", tc.id)
}
stdin := testutil.NewWriterAttachedToInvocation(t, testutil.Logger(t), inv)
stdout := expecter.NewAttachedToInvocation(t, inv)
_ = agenttest.New(t, client.URL, r.AgentToken)
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).Wait()
ctx := testutil.Context(t, testutil.WaitLong)
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
if tc.stdin != "" {
stdin.WriteLine(tc.stdin)
}
if tc.stdout != "" {
stdout.ExpectMatch(ctx, tc.stdout)
}
<-cmdDone
assertConnLog(t, connLogger, r.Workspace, database.ConnectionTypeReconnectingPty, tc.id)
})
stdin.WriteLine("exit")
<-cmdDone
})
t.Run("Command", func(t *testing.T) {
t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t)
randStr := uuid.NewString()
inv, root := clitest.New(t, "exp", "rpty", workspace.Name, "echo", randStr)
clitest.SetupConfig(t, client, root)
stdout := expecter.NewAttachedToInvocation(t, inv)
ctx := testutil.Context(t, testutil.WaitLong)
_ = agenttest.New(t, client.URL, agentToken)
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
stdout.ExpectMatch(ctx, randStr)
<-cmdDone
})
}
t.Run("NotFound", func(t *testing.T) {
t.Parallel()
cli/ssh.go
+70 -16
View File
@@ -25,11 +25,12 @@ import (
"github.com/google/uuid"
"github.com/mattn/go-isatty"
"github.com/spf13/afero"
"go.opentelemetry.io/otel/baggage"
"go.opentelemetry.io/otel/propagation"
gossh "golang.org/x/crypto/ssh"
gosshagent "golang.org/x/crypto/ssh/agent"
"golang.org/x/term"
"golang.org/x/xerrors"
"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
"tailscale.com/types/netlogtype"
"cdr.dev/slog/v3"
@@ -38,6 +39,7 @@ import (
"github.com/coder/coder/v2/cli/cliui"
"github.com/coder/coder/v2/cli/cliutil"
"github.com/coder/coder/v2/coderd/autobuild/notify"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/coderd/util/maps"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/codersdk"
@@ -481,7 +483,16 @@ func (r *RootCmd) ssh() *serpent.Command {
})
defer closeUsage()
}
return runCoderConnectStdio(ctx, fmt.Sprintf("%s:22", coderConnectHost), stdioReader, stdioWriter, stack, logger)
return runCoderConnectStdio(ctx, coderConnectOpts{
host: coderConnectHost,
httpPort: workspacesdk.AgentHTTPAPIServerPort,
tcpPort: workspacesdk.AgentStandardSSHPort,
stdin: stdioReader,
stdout: stdioWriter,
stack: stack,
logger: logger,
clientSessionID: sessionID,
})
}
}
@@ -546,7 +557,7 @@ func (r *RootCmd) ssh() *serpent.Command {
}
if stdio {
rawSSH, err := conn.SSH(ctx)
rawSSH, err := conn.SSHTCPConn(ctx)
if err != nil {
return xerrors.Errorf("connect SSH: %w", err)
}
@@ -1476,7 +1487,7 @@ func (c *closerStack) push(name string, closer io.Closer) error {
// rawSSHCopier handles copying raw SSH data between the conn and the pair (r, w).
type rawSSHCopier struct {
conn *gonet.TCPConn
conn workspacesdk.TCPConn
logger slog.Logger
r io.Reader
w io.Writer
@@ -1484,7 +1495,7 @@ type rawSSHCopier struct {
done chan struct{}
}
func newRawSSHCopier(logger slog.Logger, conn *gonet.TCPConn, r io.Reader, w io.Writer) *rawSSHCopier {
func newRawSSHCopier(logger slog.Logger, conn workspacesdk.TCPConn, r io.Reader, w io.Writer) *rawSSHCopier {
return &rawSSHCopier{conn: conn, logger: logger, r: r, w: w, done: make(chan struct{})}
}
@@ -1721,26 +1732,69 @@ func testOrDefaultDialer(ctx context.Context) coderConnectDialer {
return dialer
}
func runCoderConnectStdio(ctx context.Context, addr string, stdin io.Reader, stdout io.Writer, stack *closerStack, logger slog.Logger) error {
dialer := testOrDefaultDialer(ctx)
type coderConnectOpts struct {
host string
httpPort uint16
tcpPort uint16
clientSessionID string
stdin io.Reader
stdout io.Writer
stack *closerStack
logger slog.Logger
}
func runCoderConnectStdio(ctx context.Context, opts coderConnectOpts) error {
var conn net.Conn
if err := retryWithInterval(ctx, logger, sshRetryInterval, sshMaxAttempts, func() error {
var err error
conn, err = dialer.DialContext(ctx, "tcp", addr)
if err != nil {
return xerrors.Errorf("dial coder connect host %q over tcp: %w", addr, err)
if err := retryWithInterval(ctx, opts.logger, sshRetryInterval, sshMaxAttempts, func() error {
headers := http.Header{}
// Propagate any baggage and add the session ID.
if opts.clientSessionID != "" {
member, err := baggage.NewMemberRaw(tracing.SessionIDBaggageKey, opts.clientSessionID)
if err != nil {
return err
}
bctx := propagation.Baggage{}.Extract(ctx, propagation.HeaderCarrier(headers))
bag, err := baggage.FromContext(bctx).SetMember(member)
if err != nil {
return err
}
bctx = baggage.ContextWithBaggage(bctx, bag)
propagation.Baggage{}.Inject(bctx, propagation.HeaderCarrier(headers))
}
return nil
dialer := testOrDefaultDialer(ctx)
client := http.Client{
// Redirects are blocked to prevent misuse.
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
Transport: &http.Transport{
// Disable keep alives as we're usually only making a single
// request, and this triggers goleak in tests
DisableKeepAlives: true,
DialContext: dialer.DialContext,
},
}
var err error
addr := fmt.Sprintf("%s:%d", opts.host, opts.httpPort)
conn, err = workspacesdk.DialTCPUpgrade(ctx, addr, &client, headers, workspacesdk.AgentStandardSSHPort)
if errors.Is(err, workspacesdk.ErrAgentTCPUpgradeUnsupported) {
// Fall back to dialing the port directly.
addr = fmt.Sprintf("%s:%d", opts.host, opts.tcpPort)
conn, err = dialer.DialContext(ctx, "tcp", addr)
}
return err
}); err != nil {
return err
}
if err := stack.push("tcp conn", conn); err != nil {
if err := opts.stack.push("tcp conn", conn); err != nil {
return err
}
agentssh.Bicopy(ctx, conn, &StdioRwc{
Reader: stdin,
Writer: stdout,
Reader: opts.stdin,
Writer: opts.stdout,
})
return nil
cli/ssh_internal_test.go
+27 -8
View File
@@ -6,6 +6,7 @@ import (
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"os"
"sync"
@@ -278,20 +279,38 @@ func TestCoderConnectStdio(t *testing.T) {
}
}()
server := newSSHServer("127.0.0.1:0")
ln, err := net.Listen("tcp", server.server.Addr)
require.NoError(t, err)
sshSrv := newSSHServer("127.0.0.1:0")
defer sshSrv.Close()
sshLn, err := net.Listen("tcp", sshSrv.server.Addr)
require.NoError(t, err)
go func() {
_ = server.Serve(ln)
_ = sshSrv.Serve(sshLn)
}()
t.Cleanup(func() {
_ = server.Close()
})
sshAddr, ok := sshLn.Addr().(*net.TCPAddr)
require.True(t, ok)
// Force using the fallback.
httpSrv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusNotFound)
}))
defer httpSrv.Close()
httpAddr, ok := httpSrv.Listener.Addr().(*net.TCPAddr)
require.True(t, ok)
stdioDone := make(chan struct{})
go func() {
err = runCoderConnectStdio(ctx, ln.Addr().String(), clientOutput, serverInput, stack, logger)
err = runCoderConnectStdio(ctx, coderConnectOpts{
host: "127.0.0.1",
httpPort: uint16(httpAddr.Port), // nolint:gosec // ports fit within uint16
tcpPort: uint16(sshAddr.Port), // nolint:gosec // ports fit within uint16
stdin: clientOutput,
stdout: serverInput,
stack: stack,
logger: logger,
})
assert.NoError(t, err)
close(stdioDone)
}()
cli/ssh_test.go
+266 -82
View File
@@ -46,12 +46,14 @@ import (
"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/cli/cliui"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/connectionlog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/provisioner/echo"
"github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/pty"
@@ -79,32 +81,60 @@ func setupWorkspaceForAgent(t *testing.T, mutations ...func([]*proto.Agent) []*p
func TestSSH(t *testing.T) {
t.Parallel()
t.Run("ImmediateExit", func(t *testing.T) {
t.Parallel()
logger := testutil.Logger(t)
client, workspace, agentToken := setupWorkspaceForAgent(t)
inv, root := clitest.New(t, "ssh", workspace.Name)
clitest.SetupConfig(t, client, root)
stdout := expecter.NewAttachedToInvocation(t, inv)
stdin := testutil.NewWriterAttachedToInvocation(t, logger.Named("stdin"), inv)
tests := []struct {
name string
id string
}{
{name: "GenerateSessionID"},
{name: "ClientSessionID", id: "0123456789abcdef0123456789abcdef"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
connLogger := connectionlog.NewFake()
client, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
ConnectionLogger: connLogger,
})
client.SetLogger(testutil.Logger(t).Named("client"))
first := coderdtest.CreateFirstUser(t, client)
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
r.Username = "myuser"
})
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
Name: "myworkspace",
OrganizationID: first.OrganizationID,
OwnerID: user.ID,
}).WithAgent().Do()
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
inv, root := clitest.New(t, "ssh", r.Workspace.Name)
clitest.SetupConfig(t, userClient, root)
if tc.id != "" {
inv.Environ.Set("CODER_TRACE_SESSION_ID", tc.id)
}
stdout := expecter.NewAttachedToInvocation(t, inv)
stdin := testutil.NewWriterAttachedToInvocation(t, testutil.Logger(t).Named("stdin"), inv)
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
stdout.ExpectMatch(ctx, "Waiting")
_ = agenttest.New(t, client.URL, r.AgentToken)
_ = coderdtest.NewWorkspaceAgentWaiter(t, userClient, r.Workspace.ID).Wait()
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
stdin.WriteLine("exit")
<-cmdDone
assertConnLog(t, connLogger, r.Workspace, database.ConnectionTypeSsh, tc.id)
})
stdout.ExpectMatch(ctx, "Waiting")
}
_ = agenttest.New(t, client.URL, agentToken)
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
stdin.WriteLine("exit")
<-cmdDone
})
t.Run("WorkspaceNameInput", func(t *testing.T) {
t.Parallel()
@@ -445,66 +475,6 @@ func TestSSH(t *testing.T) {
}
})
t.Run("Stdio", func(t *testing.T) {
t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t)
_, _ = tGoContext(t, func(ctx context.Context) {
// Run this async so the SSH command has to wait for
// the build and agent to connect!
_ = agenttest.New(t, client.URL, agentToken)
<-ctx.Done()
})
clientOutput, clientInput := io.Pipe()
serverOutput, serverInput := io.Pipe()
defer func() {
for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
_ = c.Close()
}
}()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name)
clitest.SetupConfig(t, client, root)
inv.Stdin = clientOutput
inv.Stdout = serverInput
inv.Stderr = io.Discard
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
conn, channels, requests, err := ssh.NewClientConn(&testutil.ReaderWriterConn{
Reader: serverOutput,
Writer: clientInput,
}, "", &ssh.ClientConfig{
// #nosec
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
})
require.NoError(t, err)
defer conn.Close()
sshClient := ssh.NewClient(conn, channels, requests)
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
command := "sh -c exit"
if runtime.GOOS == "windows" {
command = "cmd.exe /c exit"
}
err = session.Run(command)
require.NoError(t, err)
err = sshClient.Close()
require.NoError(t, err)
_ = clientOutput.Close()
<-cmdDone
})
t.Run("DeterministicHostKey", func(t *testing.T) {
t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t)
@@ -1827,6 +1797,98 @@ func TestSSH(t *testing.T) {
})
}
func TestSSH_Stdio(t *testing.T) {
t.Parallel()
tests := []struct {
name string
id string
useOldAgent bool
}{
{name: "GenerateSessionID"},
{name: "ClientSessionID", id: "0123456789abcdef0123456789abcdef"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
connLogger := connectionlog.NewFake()
client, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
ConnectionLogger: connLogger,
})
client.SetLogger(testutil.Logger(t).Named("client"))
first := coderdtest.CreateFirstUser(t, client)
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
r.Username = "myuser"
})
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
Name: "myworkspace",
OrganizationID: first.OrganizationID,
OwnerID: user.ID,
}).WithAgent().Do()
_, _ = tGoContext(t, func(ctx context.Context) {
// Run this async so the SSH command has to wait for
// the build and agent to connect!
_ = agenttest.New(t, userClient.URL, r.AgentToken)
<-ctx.Done()
})
clientOutput, clientInput := io.Pipe()
serverOutput, serverInput := io.Pipe()
defer func() {
for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
_ = c.Close()
}
}()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
inv, root := clitest.New(t, "ssh", "--stdio", r.Workspace.Name)
clitest.SetupConfig(t, userClient, root)
if tc.id != "" {
inv.Environ.Set("CODER_TRACE_SESSION_ID", tc.id)
}
inv.Stdin = clientOutput
inv.Stdout = serverInput
inv.Stderr = io.Discard
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
conn, channels, requests, err := ssh.NewClientConn(&testutil.ReaderWriterConn{
Reader: serverOutput,
Writer: clientInput,
}, "", &ssh.ClientConfig{
// #nosec
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
})
require.NoError(t, err)
defer conn.Close()
sshClient := ssh.NewClient(conn, channels, requests)
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
command := "sh -c exit"
if runtime.GOOS == "windows" {
command = "cmd.exe /c exit"
}
err = session.Run(command)
require.NoError(t, err)
err = sshClient.Close()
require.NoError(t, err)
_ = clientOutput.Close()
<-cmdDone
assertConnLog(t, connLogger, r.Workspace, database.ConnectionTypeSsh, tc.id)
})
}
}
//nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel.
func TestSSH_ForwardGPG(t *testing.T) {
if runtime.GOOS == "windows" {
@@ -2168,6 +2230,99 @@ func TestSSH_Container(t *testing.T) {
func TestSSH_CoderConnect(t *testing.T) {
t.Parallel()
tests := []struct {
name string
id string
useOldAgent bool
}{
{name: "GenerateSessionID"},
{name: "ClientSessionID", id: "0123456789abcdef0123456789abcdef"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
connLogger := connectionlog.NewFake()
client, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
ConnectionLogger: connLogger,
})
client.SetLogger(testutil.Logger(t).Named("client"))
first := coderdtest.CreateFirstUser(t, client)
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
r.Username = "myuser"
})
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
Name: "myworkspace",
OrganizationID: first.OrganizationID,
OwnerID: user.ID,
}).WithAgent().Do()
_ = agenttest.New(t, userClient.URL, r.AgentToken)
resources := coderdtest.NewWorkspaceAgentWaiter(t, userClient, r.Workspace.ID).Wait()
clientOutput, clientInput := io.Pipe()
serverOutput, serverInput := io.Pipe()
defer func() {
for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
_ = c.Close()
}
}()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
defer cancel()
// Get an agent conn to use as the Coder Connect dialer.
agentConn, err := workspacesdk.New(userClient).DialAgent(ctx, resources[0].Agents[0].ID, nil)
require.NoError(t, err)
defer agentConn.Close()
require.True(t, agentConn.AwaitReachable(ctx))
inv, root := clitest.New(t, "ssh", r.Workspace.Name, "--stdio")
clitest.SetupConfig(t, userClient, root)
if tc.id != "" {
inv.Environ.Set("CODER_TRACE_SESSION_ID", tc.id)
}
ctx = cli.WithTestOnlyCoderConnectDialer(ctx, agentConn)
ctx = withCoderConnectRunning(ctx)
inv.Stdin = clientOutput
inv.Stdout = serverInput
inv.Stderr = io.Discard
cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
conn, channels, requests, err := ssh.NewClientConn(&testutil.ReaderWriterConn{
Reader: serverOutput,
Writer: clientInput,
}, "", &ssh.ClientConfig{
// #nosec
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
})
require.NoError(t, err)
defer conn.Close()
sshClient := ssh.NewClient(conn, channels, requests)
session, err := sshClient.NewSession()
require.NoError(t, err)
defer session.Close()
command := "sh -c exit"
if runtime.GOOS == "windows" {
command = "cmd.exe /c exit"
}
err = session.Run(command)
require.NoError(t, err)
err = sshClient.Close()
require.NoError(t, err)
_ = clientOutput.Close()
<-cmdDone
assertConnLog(t, connLogger, r.Workspace, database.ConnectionTypeSsh, tc.id)
})
}
t.Run("Enabled", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
@@ -2195,7 +2350,7 @@ func TestSSH_CoderConnect(t *testing.T) {
err := testutil.TryReceive(ctx, t, errCh)
// Our mock dialer will always fail with this error, if it was called
require.ErrorContains(t, err, "dial coder connect host \"dev.myworkspace.myuser.coder:22\" over tcp")
require.ErrorContains(t, err, "dial coder connect host \"dev.myworkspace.myuser.coder:4\" over tcp")
// The network info file should be created since we passed `--stdio`
entries, err := afero.ReadDir(fs, "/net")
@@ -2603,3 +2758,32 @@ func TestSSH_Completion(t *testing.T) {
require.Empty(t, output)
})
}
func assertConnLog(t *testing.T,
connLogger *connectionlog.FakeConnectionLogger,
workspace database.WorkspaceTable,
expectedType database.ConnectionType,
clientSessionID string,
) {
t.Helper()
// Wait until we get at least one related log.
var logs []database.UpsertConnectionLogParams
require.Eventually(t, func() bool {
for _, log := range connLogger.ConnectionLogs() {
if log.Type == expectedType &&
log.WorkspaceID == workspace.ID {
logs = append(logs, log)
}
}
return len(logs) > 0
}, testutil.WaitShort, testutil.IntervalFast)
for _, log := range logs {
switch {
case clientSessionID != "":
require.Equal(t, clientSessionID, log.ClientSessionID.String)
default: // Should have generated a client session ID.
require.NotEmpty(t, log.ClientSessionID.String)
}
}
}
cli/vscodessh.go
+1 -1
View File
@@ -182,7 +182,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
})
defer closeUsage()
rawSSH, err := agentConn.SSH(ctx)
rawSSH, err := agentConn.SSHTCPConn(ctx)
if err != nil {
return err
}
coderd/tracing/httpmw.go
+12
View File
@@ -28,6 +28,17 @@ import (
// 32-character hexadecimal string.
const SessionIDBaggageKey = "client_session_id"
type clientSessionIDContextKey struct{}
// ClientSessionID returns the client session ID tracing middleware handler.
func ClientSessionID(r *http.Request) string {
user, ok := r.Context().Value(clientSessionIDContextKey{}).(string)
if !ok {
panic("developer error: tracing middleware not provided")
}
return user
}
// DefaultRoutePatterns are the route patterns coderd and wsproxy trace. The
// middleware runs high in the stack so it can capture the entire request, but
// only acts on these patterns.
@@ -66,6 +77,7 @@ func Middleware(
// context. This is done even when tracing is disabled so that logs can
// always be correlated by client_session_id.
sessionID := sessionIDFromRequest(r)
r = r.WithContext(context.WithValue(r.Context(), clientSessionIDContextKey{}, sessionID))
if sessionID != "" {
r = r.WithContext(slog.With(r.Context(), slog.F("client_session_id", sessionID)))
}
coderd/workspaceapps/proxy.go
+15 -1
View File
@@ -753,7 +753,13 @@ func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
if !ok {
return
}
log := s.Logger.With(slog.F("agent_id", appToken.AgentID))
clientSessionID := tracing.ClientSessionID(r)
log := s.Logger.With(
slog.F("agent_id", appToken.AgentID),
slog.F("client_session_id", clientSessionID),
)
log.Debug(ctx, "resolved PTY request")
values := r.URL.Query()
@@ -802,6 +808,14 @@ func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
}
defer release()
log.Debug(ctx, "dialed workspace agent")
// Pass through the client session ID to the agent via the header.
if clientSessionID != "" {
agentConn.SetExtraHeaders(http.Header{
"baggage": []string{tracing.SessionIDBaggageKey + "=" + clientSessionID},
})
}
// #nosec G115 - Safe conversion for terminal height/width which are expected to be within uint16 range (0-65535)
ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect, uint16(height), uint16(width), r.URL.Query().Get("command"), func(arp *workspacesdk.AgentReconnectingPTYInit) {
arp.Container = container
codersdk/agentsdk/agentsdk.go
+13
View File
@@ -362,6 +362,19 @@ func (c *Client) ConnectRPC211WithRole(ctx context.Context, role string) (
return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
}
// ConnectRPC212WithRole returns a dRPC client to the Agent API v2.12, which
// adds the ability to attach a client session ID to connections. Pass role
// "agent" for workspace agents to enable connection monitoring.
func (c *Client) ConnectRPC212WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient212, tailnetproto.DRPCTailnetClient28, error,
) {
conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 12), role)
if err != nil {
return nil, nil, err
}
return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
}
// ConnectRPC connects to the workspace agent API and tailnet API.
// It does not send a role query parameter, so the server will apply
// its default behavior (currently: enable connection monitoring for
codersdk/workspacesdk/agentconn.go
+137 -4
View File
@@ -125,7 +125,9 @@ type AgentConn interface {
UploadChatFile(ctx context.Context, req UploadChatFileRequest) (AgentUploadChatFileResponse, error)
EditFiles(ctx context.Context, edits FileEditRequest) (FileEditResponse, error)
BundleFiles(ctx context.Context, req BundleFilesRequest) ([]byte, error)
// Deprecated: Use SSHTCPConn instead.
SSH(ctx context.Context) (*gonet.TCPConn, error)
SSHTCPConn(ctx context.Context) (TCPConn, error)
SSHClient(ctx context.Context) (*ssh.Client, error)
SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error)
SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error)
@@ -249,7 +251,16 @@ func (c *agentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, w
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
}
conn, err := c.DialContextTCP(ctx, netip.AddrPortFrom(c.agentAddress(), AgentReconnectingPTYPort))
c.headersMu.RLock()
extraHeaders := c.extraHeaders.Clone()
c.headersMu.RUnlock()
addr := netip.AddrPortFrom(c.agentAddress(), AgentHTTPAPIServerPort)
conn, err := DialTCPUpgrade(ctx, addr.String(), c.apiClient(ctx), extraHeaders, AgentReconnectingPTYPort)
if errors.Is(err, ErrAgentTCPUpgradeUnsupported) {
// Fall back to dialing the port directly.
conn, err = c.DialContextTCP(ctx, netip.AddrPortFrom(c.agentAddress(), AgentReconnectingPTYPort))
}
if err != nil {
return nil, err
}
@@ -279,12 +290,38 @@ func (c *agentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, w
return conn, nil
}
// SSH pipes the SSH protocol over the returned net.Conn.
// SSH pipes the SSH protocol over the returned gonet.TCPConn.
// This connects to the built-in SSH server in the workspace agent.
func (c *agentConn) SSH(ctx context.Context) (*gonet.TCPConn, error) {
return c.SSHOnPort(ctx, AgentSSHPort)
}
// SSHTCPConn makes an HTTP request with the client session ID that then
// upgrades into an SSH connection. If the agent does not support the endpoint,
// falls back to dialing the port directly.
func (c *agentConn) SSHTCPConn(ctx context.Context) (TCPConn, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
if !c.AwaitReachable(ctx) {
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
}
c.SendConnectedTelemetry(c.agentAddress(), tailnet.TelemetryApplicationSSH)
c.headersMu.RLock()
extraHeaders := c.extraHeaders.Clone()
c.headersMu.RUnlock()
addr := netip.AddrPortFrom(c.agentAddress(), AgentHTTPAPIServerPort)
conn, err := DialTCPUpgrade(ctx, addr.String(), c.apiClient(ctx), extraHeaders, AgentStandardSSHPort)
if errors.Is(err, ErrAgentTCPUpgradeUnsupported) {
// Fall back to dialing the port directly.
return c.DialContextTCP(ctx, netip.AddrPortFrom(c.agentAddress(), AgentStandardSSHPort))
}
return conn, err
}
// SSHOnPort pipes the SSH protocol over the returned net.Conn.
// This connects to the built-in SSH server in the workspace agent on the specified port.
func (c *agentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error) {
@@ -299,9 +336,29 @@ func (c *agentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn,
return c.DialContextTCP(ctx, netip.AddrPortFrom(c.agentAddress(), port))
}
// SSHClient calls SSH to create a client
// SSHClient makes an HTTP request with the client session ID that then upgrades
// into an SSH client connection. If the agent does not support the endpoint,
// falls back to dialing the port directly.
func (c *agentConn) SSHClient(ctx context.Context) (*ssh.Client, error) {
return c.SSHClientOnPort(ctx, AgentSSHPort)
ctx, span := tracing.StartSpan(ctx)
defer span.End()
netConn, err := c.SSHTCPConn(ctx)
if err != nil {
return nil, xerrors.Errorf("ssh: %w", err)
}
sshConn, channels, requests, err := ssh.NewClientConn(netConn, "localhost:22", &ssh.ClientConfig{
// SSH host validation isn't helpful, because obtaining a peer
// connection already signifies user-intent to dial a workspace.
// #nosec
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
})
if err != nil {
return nil, xerrors.Errorf("ssh conn: %w", err)
}
return ssh.NewClient(sshConn, channels, requests), nil
}
// SSHClientOnPort calls SSH to create a client on a specific port
@@ -327,6 +384,82 @@ func (c *agentConn) SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Clie
return ssh.NewClient(sshConn, channels, requests), nil
}
type TCPConn interface {
net.Conn
CloseWrite() error
}
var ErrAgentTCPUpgradeUnsupported = xerrors.New("agent does not support the tcp upgrade endpoint")
// DialTCPUpgrade dials the agent's /tcp HTTP endpoint which allows sending
// extra data to the agent via headers. These must include the client session
// ID or the connection will be refused. The agent will then upgrade the
// connection based on the port. Returns ErrAgentTCPUpgradeUnsupported if the
// agent does not support the endpoint.
func DialTCPUpgrade(ctx context.Context, addr string, client *http.Client, headers http.Header, port uint16) (TCPConn, error) {
apiURL := fmt.Sprintf("http://%s/api/v0/tcp/%d", addr, port)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, nil)
if err != nil {
return nil, xerrors.Errorf("new http api request to %q: %w", apiURL, err)
}
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Upgrade", "tcp")
// The client session ID baggage is found in the extra headers.
for k, v := range headers {
req.Header[k] = v
}
//nolint:bodyclose // On success the caller is responsible for closing.
resp, err := client.Do(req)
if err != nil {
return nil, xerrors.Errorf("do upgrade request to %q: %w", apiURL, err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
_ = resp.Body.Close()
return nil, ErrAgentTCPUpgradeUnsupported
}
respBody := resp.Body
conn, ok := respBody.(rawConn)
if !ok {
_ = respBody.Close()
return nil, xerrors.Errorf("response body is not a rawConn: %T", respBody)
}
return &upgradedConn{
rawConn: conn,
remoteAddr: upgradeAddr{"tcp", addr},
// For now setting a fake local address.
localAddr: upgradeAddr{"ssh", "ssh"},
}, nil
}
type rawConn interface {
io.ReadWriteCloser
CloseWrite() error
}
type upgradedConn struct {
rawConn
remoteAddr net.Addr
localAddr net.Addr
}
func (u *upgradedConn) LocalAddr() net.Addr { return u.localAddr }
func (u *upgradedConn) RemoteAddr() net.Addr { return u.remoteAddr }
func (*upgradedConn) SetDeadline(_ time.Time) error { return nil }
func (*upgradedConn) SetReadDeadline(_ time.Time) error { return nil }
func (*upgradedConn) SetWriteDeadline(_ time.Time) error { return nil }
type upgradeAddr struct {
network string
addr string
}
func (u upgradeAddr) Network() string { return u.network }
func (u upgradeAddr) String() string { return u.addr }
// Speedtest runs a speedtest against the workspace agent.
func (c *agentConn) Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
ctx, span := tracing.StartSpan(ctx)
codersdk/workspacesdk/agentconn_test.go
+111
View File
@@ -2,8 +2,10 @@ package workspacesdk_test
import (
"context"
"encoding/binary"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/netip"
@@ -185,6 +187,115 @@ func TestAgentConnAppHTTPClientRefusesRedirects(t *testing.T) {
require.ErrorIs(t, client.CheckRedirect(req, nil), http.ErrUseLastResponse)
}
func TestAgent_SSHUpgrade_Fallback(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
clientID := uuid.New()
agentID := uuid.New()
clientConn, _ := newTailnetConn(t, derpMap, clientID, "client")
agentConn, agentIP := newTailnetConn(t, derpMap, agentID, "agent")
stitchTailnet(t, map[uuid.UUID]*tailnet.Conn{
clientID: clientConn,
agentID: agentConn,
})
// Simulate an older agent that does not support the tcp upgrade endpoint.
var upgradeHit atomic.Bool
upgradeHandler := http.NewServeMux()
upgradeHandler.HandleFunc("/api/v0/tcp/", func(rw http.ResponseWriter, r *http.Request) {
upgradeHit.Store(true)
rw.WriteHeader(http.StatusNotFound)
})
serveTailnetHTTP(t, agentConn, upgradeHandler)
ln, err := agentConn.Listen("tcp", fmt.Sprintf(":%d", workspacesdk.AgentStandardSSHPort))
require.NoError(t, err)
defer ln.Close()
go func() {
conn, err := ln.Accept()
assert.NoError(t, err)
defer conn.Close()
_, err = conn.Write([]byte("hello world"))
assert.NoError(t, err)
}()
conn := workspacesdk.NewAgentConn(clientConn, workspacesdk.AgentConnOptions{
AgentID: agentID,
})
require.True(t, clientConn.AwaitReachable(ctx, agentIP))
sshConn, err := conn.SSHTCPConn(ctx)
require.NoError(t, err)
defer sshConn.Close()
b, err := io.ReadAll(sshConn)
require.NoError(t, err)
require.True(t, upgradeHit.Load(), "should hit upgrade")
require.Equal(t, "hello world", string(b))
}
func TestAgent_ReconnectingPTYUpgrade_Fallback(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
clientID := uuid.New()
agentID := uuid.New()
clientConn, _ := newTailnetConn(t, derpMap, clientID, "client")
agentConn, agentIP := newTailnetConn(t, derpMap, agentID, "agent")
stitchTailnet(t, map[uuid.UUID]*tailnet.Conn{
clientID: clientConn,
agentID: agentConn,
})
// Simulate an older agent that does not support the tcp upgrade endpoint.
var upgradeHit atomic.Bool
upgradeHandler := http.NewServeMux()
upgradeHandler.HandleFunc("/api/v0/tcp/", func(rw http.ResponseWriter, r *http.Request) {
upgradeHit.Store(true)
rw.WriteHeader(http.StatusNotFound)
})
serveTailnetHTTP(t, agentConn, upgradeHandler)
ln, err := agentConn.Listen("tcp", fmt.Sprintf(":%d", workspacesdk.AgentReconnectingPTYPort))
require.NoError(t, err)
defer ln.Close()
go func() {
conn, err := ln.Accept()
assert.NoError(t, err)
defer conn.Close()
// Read in the init message.
rawLen := make([]byte, 2)
_, err = io.ReadFull(conn, rawLen)
assert.NoError(t, err)
_, err = io.ReadFull(conn, make([]byte, binary.LittleEndian.Uint16(rawLen)))
assert.NoError(t, err)
// Reply with some bytes.
_, err = conn.Write([]byte("hello world"))
assert.NoError(t, err)
}()
conn := workspacesdk.NewAgentConn(clientConn, workspacesdk.AgentConnOptions{
AgentID: agentID,
})
require.True(t, clientConn.AwaitReachable(ctx, agentIP))
ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "bash")
require.NoError(t, err)
defer ptyConn.Close()
b, err := io.ReadAll(ptyConn)
require.NoError(t, err)
require.True(t, upgradeHit.Load(), "should hit upgrade")
require.Equal(t, "hello world", string(b))
}
func newTailnetConn(t *testing.T, derpMap *tailcfg.DERPMap, id uuid.UUID, name string) (*tailnet.Conn, netip.Addr) {
t.Helper()
codersdk/workspacesdk/agentconnmock/agentconnmock.go
+15
View File
@@ -559,6 +559,21 @@ func (mr *MockAgentConnMockRecorder) SSHOnPort(ctx, port any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSHOnPort", reflect.TypeOf((*MockAgentConn)(nil).SSHOnPort), ctx, port)
}
// SSHTCPConn mocks base method.
func (m *MockAgentConn) SSHTCPConn(ctx context.Context) (workspacesdk.TCPConn, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SSHTCPConn", ctx)
ret0, _ := ret[0].(workspacesdk.TCPConn)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// SSHTCPConn indicates an expected call of SSHTCPConn.
func (mr *MockAgentConnMockRecorder) SSHTCPConn(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSHTCPConn", reflect.TypeOf((*MockAgentConn)(nil).SSHTCPConn), ctx)
}
// SetExtraHeaders mocks base method.
func (m *MockAgentConn) SetExtraHeaders(h http.Header) {
m.ctrl.T.Helper()