Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 1a10358

Browse files
kylecarbskylecarbs
committed
feat: Add support for VS Code and JetBrains Gateway via SSH
This fixes various bugs that made this not work: - Incorrect max message size in `peer` - Incorrect reader buffer size in `peer` - Lack of SFTP support in `agent` - Lack of direct-tcpip support in `agent` - Misuse of command from session. It should always use the shell - Blocking on SSH session, only allowing one at a time Fixes #833 too.
1 parent 8fecb67 commit 1a10358

File tree

5 files changed

+71
-44
lines changed

5 files changed

+71
-44
lines changed

‎agent/agent.go

Lines changed: 35 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ import (
2121
"github.com/coder/coder/pty"
2222
"github.com/coder/retry"
2323

24+
"github.com/pkg/sftp"
25+
2426
"github.com/gliderlabs/ssh"
2527
gossh "golang.org/x/crypto/ssh"
2628
"golang.org/x/xerrors"
@@ -120,7 +122,7 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
120122

121123
switch channel.Protocol() {
122124
case "ssh":
123-
a.sshServer.HandleConn(channel.NetConn())
125+
go a.sshServer.HandleConn(channel.NetConn())
124126
default:
125127
a.options.Logger.Warn(ctx, "unhandled protocol from channel",
126128
slog.F("protocol", channel.Protocol()),
@@ -145,7 +147,10 @@ func (a *agent) init(ctx context.Context) {
145147
sshLogger := a.options.Logger.Named("ssh-server")
146148
forwardHandler := &ssh.ForwardedTCPHandler{}
147149
a.sshServer = &ssh.Server{
148-
ChannelHandlers: ssh.DefaultChannelHandlers,
150+
ChannelHandlers: map[string]ssh.ChannelHandler{
151+
"direct-tcpip": ssh.DirectTCPIPHandler,
152+
"session": ssh.DefaultSessionHandler,
153+
},
149154
ConnectionFailedCallback: func(conn net.Conn, err error) {
150155
sshLogger.Info(ctx, "ssh connection ended", slog.Error(err))
151156
},
@@ -184,61 +189,54 @@ func (a *agent) init(ctx context.Context) {
184189
NoClientAuth: true,
185190
}
186191
},
192+
SubsystemHandlers: map[string]ssh.SubsystemHandler{
193+
"sftp": func(session ssh.Session) {
194+
server, err := sftp.NewServer(session)
195+
if err != nil {
196+
a.options.Logger.Debug(session.Context(), "initialize sftp server", slog.Error(err))
197+
return
198+
}
199+
defer server.Close()
200+
err = server.Serve()
201+
if errors.Is(err, io.EOF) {
202+
return
203+
}
204+
a.options.Logger.Debug(session.Context(), "sftp server exited with error", slog.Error(err))
205+
},
206+
},
187207
}
188208

189209
go a.run(ctx)
190210
}
191211

192212
func (a *agent) handleSSHSession(session ssh.Session) error {
193-
var (
194-
command string
195-
args = []string{}
196-
err error
197-
)
198-
199213
currentUser, err := user.Current()
200214
if err != nil {
201215
return xerrors.Errorf("get current user: %w", err)
202216
}
203217
username := currentUser.Username
204218

219+
shell, err := usershell.Get(username)
220+
if err != nil {
221+
return xerrors.Errorf("get user shell: %w", err)
222+
}
223+
205224
// gliderlabs/ssh returns a command slice of zero
206225
// when a shell is requested.
226+
command := session.RawCommand()
207227
if len(session.Command()) == 0 {
208-
command, err = usershell.Get(username)
209-
if err != nil {
210-
return xerrors.Errorf("get user shell: %w", err)
211-
}
212-
} else {
213-
command = session.Command()[0]
214-
if len(session.Command()) > 1 {
215-
args = session.Command()[1:]
216-
}
228+
command = shell
217229
}
218230

219-
signals := make(chan ssh.Signal)
220-
breaks := make(chan bool)
221-
defer close(signals)
222-
defer close(breaks)
223-
go func() {
224-
for {
225-
select {
226-
case <-session.Context().Done():
227-
return
228-
// Ignore signals and breaks for now!
229-
case <-signals:
230-
case <-breaks:
231-
}
232-
}
233-
}()
234-
235-
cmd := exec.CommandContext(session.Context(), command, args...)
231+
// OpenSSH executes all commands with the users current shell.
232+
// We replicate that behavior for IDE support.
233+
cmd := exec.CommandContext(session.Context(), shell, "-c", command)
236234
cmd.Env = append(os.Environ(), session.Environ()...)
237235
executablePath, err := os.Executable()
238236
if err != nil {
239237
return xerrors.Errorf("getting os executable: %w", err)
240238
}
241-
cmd.Env = append(session.Environ(), fmt.Sprintf(`GIT_SSH_COMMAND="%s gitssh --"`, executablePath))
239+
cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_SSH_COMMAND="%s gitssh --"`, executablePath))
242240

243241
sshPty, windowSize, isPty := session.Pty()
244242
if isPty {
@@ -267,7 +265,7 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
267265
}
268266

269267
cmd.Stdout = session
270-
cmd.Stderr = session
268+
cmd.Stderr = session.Stderr()
271269
// This blocks forever until stdin is received if we don't
272270
// use StdinPipe. It's unknown what causes this.
273271
stdinPipe, err := cmd.StdinPipe()
@@ -281,8 +279,7 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
281279
if err != nil {
282280
return xerrors.Errorf("start: %w", err)
283281
}
284-
_ = cmd.Wait()
285-
return nil
282+
return cmd.Wait()
286283
}
287284

288285
// isClosed returns whether the API is closed or not.

‎agent/agent_test.go

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"io"
77
"net"
8+
"os"
89
"os/exec"
910
"path/filepath"
1011
"runtime"
@@ -13,6 +14,7 @@ import (
1314
"testing"
1415

1516
"github.com/pion/webrtc/v3"
17+
"github.com/pkg/sftp"
1618
"github.com/stretchr/testify/require"
1719
"go.uber.org/goleak"
1820
"golang.org/x/crypto/ssh"
@@ -114,13 +116,31 @@ func TestAgent(t *testing.T) {
114116
conn.Close()
115117
<-done
116118
})
119+
120+
t.Run("SFTP", func(t *testing.T) {
121+
t.Parallel()
122+
sshClient, err := setupAgent(t).SSHClient()
123+
require.NoError(t, err)
124+
client, err := sftp.NewClient(sshClient)
125+
require.NoError(t, err)
126+
tempFile := filepath.Join(t.TempDir(), "sftp")
127+
file, err := client.Create(tempFile)
128+
require.NoError(t, err)
129+
err = file.Close()
130+
require.NoError(t, err)
131+
_, err = os.Stat(tempFile)
132+
require.NoError(t, err)
133+
})
117134
}
118135

119136
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
120137
agentConn := setupAgent(t)
121-
socket := filepath.Join(t.TempDir(), "ssh")
122-
listener, err := net.Listen("unix", socket)
138+
139+
listener, err := net.Listen("tcp", "127.0.0.1:0")
123140
require.NoError(t, err)
141+
t.Cleanup(func() {
142+
_ = listener.Close()
143+
})
124144
go func() {
125145
for {
126146
conn, err := listener.Accept()
@@ -136,7 +156,7 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
136156
t.Cleanup(func() {
137157
_ = listener.Close()
138158
})
139-
args := append(beforeArgs, "-o", "ProxyCommand socat - UNIX-CLIENT:"+socket, "-o", "StrictHostKeyChecking=no", "host")
159+
args := append(beforeArgs, "-o", "ProxyCommand socat - TCP4:"+listener.Addr().String(), "-o", "StrictHostKeyChecking=no", "host")
140160
args = append(args, afterArgs...)
141161
return exec.Command("ssh", args...)
142162
}

‎go.mod

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,12 @@ require (
9595

9696
require github.com/go-chi/httprate v0.5.3
9797

98-
require github.com/jedib0t/go-pretty/v6 v6.3.0
98+
require (
99+
github.com/jedib0t/go-pretty/v6 v6.3.0
100+
github.com/pkg/sftp v1.13.4
101+
)
102+
103+
require github.com/kr/fs v0.1.0 // indirect
99104

100105
require (
101106
github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect

‎go.sum

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1103,6 +1103,7 @@ github.com/klauspost/crc32 v1.2.0/go.mod h1:+ZoRqAPRLkC4NPOvfYeR5KNOrY6TD+/sAC3H
11031103
github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
11041104
github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
11051105
github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
1106+
github.com/kr/fs v0.1.0 h1:Jskdu9ieNAYnjxsi0LbQp1ulIKZV1LAFgK1tWhpZgl8=
11061107
github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
11071108
github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
11081109
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
@@ -1438,6 +1439,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
14381439
github.com/pkg/profile v1.2.1/go.mod h1:hJw3o1OdXxsrSjjVksARp5W95eeEaEfptyVZyv6JUPA=
14391440
github.com/pkg/profile v1.6.0/go.mod h1:qBsxPvzyUincmltOk6iyRVxHYg4adc0OFOv72ZdLa18=
14401441
github.com/pkg/sftp v1.13.1/go.mod h1:3HaPG6Dq1ILlpPZRO0HVMrsydcdLt6HRDccSgb87qRg=
1442+
github.com/pkg/sftp v1.13.4 h1:Lb0RYJCmgUcBgZosfoi9Y9sbl6+LJgOIgk/2Y4YjMFg=
1443+
github.com/pkg/sftp v1.13.4/go.mod h1:LzqnAvaD5TWeNBsZpfKxSYn1MbjWwOsCIAFFJbpIsK8=
14411444
github.com/pmezard/go-difflib v0.0.0-20151028094244-d8ed2627bdf0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
14421445
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
14431446
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

‎peer/channel.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ const (
2121
// For some reason messages larger just don't work...
2222
// This shouldn't be a huge deal for real-world usage.
2323
// See: https://github.com/pion/datachannel/issues/59
24-
maxMessageLength = 32 * 1024 // 32 KB
24+
maxMessageLength = 64 * 1024 // 64 KB
2525
)
2626

2727
// newChannel creates a new channel and initializes it.
@@ -145,7 +145,9 @@ func (c *Channel) init() {
145145
if c.opts.Unordered {
146146
c.reader = c.rwc
147147
} else {
148-
c.reader = bufio.NewReader(c.rwc)
148+
// This must be the max message length otherwise a short
149+
// buffer error can occur.
150+
c.reader = bufio.NewReaderSize(c.rwc, maxMessageLength)
149151
}
150152
close(c.opened)
151153
})

0 commit comments

Comments
 (0)