Skip to content

Commit 0bda974

Browse files
kylecarbskylecarbs
committed
feat: Add config-ssh and tests for resiliency
1 parent 7ed283c commit 0bda974

File tree

10 files changed

+121
-27
lines changed

10 files changed

+121
-27
lines changed

‎agent/agent_test.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,10 @@ func TestAgent(t *testing.T) {
3939
t.Cleanup(func() {
4040
_ = conn.Close()
4141
})
42-
client := agent.Conn{conn}
42+
client := agent.Conn{
43+
Negotiator: api,
44+
Conn: conn,
45+
}
4346
sshClient, err := client.SSHClient()
4447
require.NoError(t, err)
4548
session, err := sshClient.NewSession()
@@ -65,7 +68,10 @@ func TestAgent(t *testing.T) {
6568
t.Cleanup(func() {
6669
_ = conn.Close()
6770
})
68-
client := &agent.Conn{conn}
71+
client := &agent.Conn{
72+
Negotiator: api,
73+
Conn: conn,
74+
}
6975
sshClient, err := client.SSHClient()
7076
require.NoError(t, err)
7177
session, err := sshClient.NewSession()

‎agent/conn.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,15 @@ import (
88
"golang.org/x/xerrors"
99

1010
"github.com/coder/coder/peer"
11+
"github.com/coder/coder/peerbroker/proto"
1112
)
1213

1314
// Conn wraps a peer connection with helper functions to
1415
// communicate with the agent.
1516
type Conn struct {
17+
// Negotiator is responsible for exchanging messages.
18+
Negotiator proto.DRPCPeerBrokerClient
19+
1620
*peer.Conn
1721
}
1822

@@ -48,3 +52,8 @@ func (c *Conn) SSHClient() (*ssh.Client, error) {
4852
}
4953
return ssh.NewClient(sshConn, channels, requests), nil
5054
}
55+
56+
func (c *Conn) Close() error {
57+
_ = c.Negotiator.DRPCConn().Close()
58+
return c.Conn.Close()
59+
}

‎cli/root.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ func Root() *cobra.Command {
6464
projects(),
6565
users(),
6666
workspaces(),
67-
workspaceSSH(),
67+
ssh(),
6868
workspaceTunnel(),
6969
)
7070

‎cli/ssh.go

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,15 @@
11
package cli
22

33
import (
4-
"os"
5-
64
"github.com/spf13/cobra"
7-
"golang.org/x/crypto/ssh"
8-
"golang.org/x/term"
5+
gossh "golang.org/x/crypto/ssh"
96
"golang.org/x/xerrors"
107

118
"github.com/coder/coder/coderd/database"
129
"github.com/coder/coder/codersdk"
1310
)
1411

15-
func workspaceSSH() *cobra.Command {
12+
func ssh() *cobra.Command {
1613
cmd := &cobra.Command{
1714
Use: "ssh <workspace> [resource]",
1815
RunE: func(cmd *cobra.Command, args []string) error {
@@ -68,6 +65,7 @@ func workspaceSSH() *cobra.Command {
6865
if err != nil {
6966
return err
7067
}
68+
defer conn.Close()
7169
sshClient, err := conn.SSHClient()
7270
if err != nil {
7371
return err
@@ -77,16 +75,16 @@ func workspaceSSH() *cobra.Command {
7775
if err != nil {
7876
return err
7977
}
80-
_, _ = term.MakeRaw(int(os.Stdin.Fd()))
81-
err = sshSession.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
82-
ssh.OCRNL: 1,
78+
79+
err = sshSession.RequestPty("xterm-256color", 128, 128, gossh.TerminalModes{
80+
gossh.OCRNL: 1,
8381
})
8482
if err != nil {
8583
return err
8684
}
87-
sshSession.Stdin = os.Stdin
88-
sshSession.Stdout = os.Stdout
89-
sshSession.Stderr = os.Stderr
85+
sshSession.Stdin = cmd.InOrStdin()
86+
sshSession.Stdout = cmd.OutOrStdout()
87+
sshSession.Stderr = cmd.OutOrStdout()
9088
err = sshSession.Shell()
9189
if err != nil {
9290
return err

‎cli/ssh_test.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package cli_test
2+
3+
import (
4+
"testing"
5+
6+
"github.com/google/uuid"
7+
"github.com/stretchr/testify/require"
8+
9+
"cdr.dev/slog"
10+
"cdr.dev/slog/sloggers/slogtest"
11+
"github.com/coder/coder/agent"
12+
"github.com/coder/coder/cli/clitest"
13+
"github.com/coder/coder/coderd/coderdtest"
14+
"github.com/coder/coder/codersdk"
15+
"github.com/coder/coder/peer"
16+
"github.com/coder/coder/provisioner/echo"
17+
"github.com/coder/coder/provisionersdk/proto"
18+
"github.com/coder/coder/pty/ptytest"
19+
)
20+
21+
func TestSSH(t *testing.T) {
22+
t.Parallel()
23+
t.Run("Echo", func(t *testing.T) {
24+
t.Parallel()
25+
client := coderdtest.New(t, nil)
26+
user := coderdtest.CreateFirstUser(t, client)
27+
daemonCloser := coderdtest.NewProvisionerDaemon(t, client)
28+
agentToken := uuid.NewString()
29+
version := coderdtest.CreateProjectVersion(t, client, user.OrganizationID, &echo.Responses{
30+
Parse: echo.ParseComplete,
31+
ProvisionDryRun: echo.ProvisionComplete,
32+
Provision: []*proto.Provision_Response{{
33+
Type: &proto.Provision_Response_Complete{
34+
Complete: &proto.Provision_Complete{
35+
Resources: []*proto.Resource{{
36+
Name: "dev",
37+
Type: "google_compute_instance",
38+
Agent: &proto.Agent{
39+
Id: uuid.NewString(),
40+
Auth: &proto.Agent_Token{
41+
Token: agentToken,
42+
},
43+
},
44+
}},
45+
},
46+
},
47+
}},
48+
})
49+
coderdtest.AwaitProjectVersionJob(t, client, version.ID)
50+
project := coderdtest.CreateProject(t, client, user.OrganizationID, version.ID)
51+
workspace := coderdtest.CreateWorkspace(t, client, "", project.ID)
52+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
53+
daemonCloser.Close()
54+
agentClient := codersdk.New(client.URL)
55+
agentClient.SessionToken = agentToken
56+
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
57+
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
58+
})
59+
defer agentCloser.Close()
60+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
61+
62+
cmd, root := clitest.New(t, "ssh", workspace.Name)
63+
clitest.SetupConfig(t, client, root)
64+
doneChan := make(chan struct{})
65+
pty := ptytest.New(t)
66+
cmd.SetIn(pty.Input())
67+
cmd.SetOut(pty.Output())
68+
go func() {
69+
defer close(doneChan)
70+
err := cmd.Execute()
71+
require.NoError(t, err)
72+
}()
73+
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
74+
pty.ExpectMatch("ish")
75+
pty.WriteLine("exit")
76+
<-doneChan
77+
})
78+
}

‎cli/workspaces.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ func workspaces() *cobra.Command {
1818
cmd.AddCommand(workspaceShow())
1919
cmd.AddCommand(workspaceStop())
2020
cmd.AddCommand(workspaceStart())
21-
cmd.AddCommand(workspaceSSH())
21+
cmd.AddCommand(ssh())
2222
cmd.AddCommand(workspaceUpdate())
2323

2424
return cmd

‎coderd/coderdtest/coderdtest.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@ func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, build uuid.UUID
251251
if resource.Agent == nil {
252252
continue
253253
}
254+
// fmt.Printf("resources: %+v\n", resource.Agent)
254255
if resource.Agent.FirstConnectedAt == nil {
255256
return false
256257
}

‎codersdk/workspaceresources.go

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -133,15 +133,9 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, resource uuid.UUID, ice
133133
if err != nil {
134134
return nil, xerrors.Errorf("dial peer: %w", err)
135135
}
136-
go func() {
137-
// The stream is kept alive to renegotiate the RTC connection
138-
// if need-be. The calling context can be canceled to end
139-
// the negotiation stream, but not the peer connection.
140-
<-peerConn.Closed()
141-
_ = conn.Close(websocket.StatusNormalClosure, "")
142-
}()
143136
return &agent.Conn{
144-
Conn: peerConn,
137+
Negotiator: client,
138+
Conn: peerConn,
145139
}, nil
146140
}
147141

‎peerbroker/dial.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package peerbroker
22

33
import (
4+
"context"
5+
"errors"
6+
"io"
47
"reflect"
58

69
"github.com/pion/webrtc/v3"
@@ -54,6 +57,11 @@ func Dial(stream proto.DRPCPeerBroker_NegotiateConnectionClient, iceServers []we
5457
for {
5558
serverToClientMessage, err := stream.Recv()
5659
if err != nil {
60+
// p2p connections should never die if this stream does due
61+
// to proper closure or context cancellation!
62+
if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
63+
return
64+
}
5765
_ = peerConn.CloseWithError(xerrors.Errorf("recv: %w", err))
5866
return
5967
}

‎peerbroker/listen.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -166,8 +166,10 @@ func (b *peerBrokerService) NegotiateConnection(stream proto.DRPCPeerBroker_Nego
166166
for {
167167
clientToServerMessage, err := stream.Recv()
168168
if err != nil {
169-
if errors.Is(err, io.EOF) {
170-
break
169+
// p2p connections should never die if this stream does due
170+
// to proper closure or context cancellation!
171+
if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
172+
return nil
171173
}
172174
return peerConn.CloseWithError(xerrors.Errorf("recv: %w", err))
173175
}
@@ -186,6 +188,4 @@ func (b *peerBrokerService) NegotiateConnection(stream proto.DRPCPeerBroker_Nego
186188
return peerConn.CloseWithError(xerrors.Errorf("unhandled message: %s", reflect.TypeOf(clientToServerMessage).String()))
187189
}
188190
}
189-
190-
return nil
191191
}

0 commit comments

Comments
 (0)