Skip to content

Commit 5ba96b5

Browse files
committed
Add interface for swapping
1 parent 38f4004 commit 5ba96b5

File tree

6 files changed

+33
-19
lines changed

6 files changed

+33
-19
lines changed

agent/agent_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -439,7 +439,7 @@ func setupSSHSession(t *testing.T, options agent.Metadata) *ssh.Session {
439439
return session
440440
}
441441

442-
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) *agent.Conn {
442+
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) agent.Conn {
443443
client, server := provisionersdk.TransportPipe()
444444
closer := agent.New(agent.Options{
445445
FetchMetadata: func(ctx context.Context) (agent.Metadata, error) {
@@ -468,7 +468,7 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
468468
_ = conn.Close()
469469
})
470470

471-
return &agent.Conn{
471+
return &agent.WebRTCConn{
472472
Negotiator: api,
473473
Conn: conn,
474474
}

agent/conn.go

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@ import (
44
"context"
55
"encoding/json"
66
"fmt"
7+
"io"
78
"net"
89
"net/url"
910
"strings"
11+
"time"
1012

1113
"golang.org/x/crypto/ssh"
1214
"golang.org/x/xerrors"
@@ -23,9 +25,21 @@ type ReconnectingPTYRequest struct {
2325
Width uint16 `json:"width"`
2426
}
2527

28+
// Conn is a temporary interface while we switch from WebRTC to Wireguard networking.
29+
type Conn interface {
30+
io.Closer
31+
Closed() <-chan struct{}
32+
Ping() (time.Duration, error)
33+
CloseWithError(err error) error
34+
ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error)
35+
SSH() (net.Conn, error)
36+
SSHClient() (*ssh.Client, error)
37+
DialContext(ctx context.Context, network string, addr string) (net.Conn, error)
38+
}
39+
2640
// Conn wraps a peer connection with helper functions to
2741
// communicate with the agent.
28-
type Conn struct {
42+
type WebRTCConn struct {
2943
// Negotiator is responsible for exchanging messages.
3044
Negotiator proto.DRPCPeerBrokerClient
3145

@@ -36,7 +50,7 @@ type Conn struct {
3650
// be reconnected to via ID.
3751
//
3852
// The command is optional and defaults to start a shell.
39-
func (c *Conn) ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error) {
53+
func (c *WebRTCConn) ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error) {
4054
channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d:%s", id, height, width, command), &peer.ChannelOptions{
4155
Protocol: ProtocolReconnectingPTY,
4256
})
@@ -47,7 +61,7 @@ func (c *Conn) ReconnectingPTY(id string, height, width uint16, command string)
4761
}
4862

4963
// SSH dials the built-in SSH server.
50-
func (c *Conn) SSH() (net.Conn, error) {
64+
func (c *WebRTCConn) SSH() (net.Conn, error) {
5165
channel, err := c.CreateChannel(context.Background(), "ssh", &peer.ChannelOptions{
5266
Protocol: ProtocolSSH,
5367
})
@@ -59,7 +73,7 @@ func (c *Conn) SSH() (net.Conn, error) {
5973

6074
// SSHClient calls SSH to create a client that uses a weak cipher
6175
// for high throughput.
62-
func (c *Conn) SSHClient() (*ssh.Client, error) {
76+
func (c *WebRTCConn) SSHClient() (*ssh.Client, error) {
6377
netConn, err := c.SSH()
6478
if err != nil {
6579
return nil, xerrors.Errorf("ssh: %w", err)
@@ -78,7 +92,7 @@ func (c *Conn) SSHClient() (*ssh.Client, error) {
7892

7993
// DialContext dials an arbitrary protocol+address from inside the workspace and
8094
// proxies it through the provided net.Conn.
81-
func (c *Conn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
95+
func (c *WebRTCConn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
8296
u := &url.URL{
8397
Scheme: network,
8498
}
@@ -112,7 +126,7 @@ func (c *Conn) DialContext(ctx context.Context, network string, addr string) (ne
112126
return channel.NetConn(), nil
113127
}
114128

115-
func (c *Conn) Close() error {
129+
func (c *WebRTCConn) Close() error {
116130
_ = c.Negotiator.DRPCConn().Close()
117131
return c.Conn.Close()
118132
}

cli/portforward.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ func portForward() *cobra.Command {
157157
return cmd
158158
}
159159

160-
func listenAndPortForward(ctx context.Context, cmd *cobra.Command, conn *coderagent.Conn, wg *sync.WaitGroup, spec portForwardSpec) (net.Listener, error) {
160+
func listenAndPortForward(ctx context.Context, cmd *cobra.Command, conn coderagent.Conn, wg *sync.WaitGroup, spec portForwardSpec) (net.Listener, error) {
161161
_, _ = fmt.Fprintf(cmd.OutOrStderr(), "Forwarding '%v://%v' locally to '%v://%v' in the workspace\n", spec.listenNetwork, spec.listenAddress, spec.dialNetwork, spec.dialAddress)
162162

163163
var (

coderd/workspaceagents.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,7 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
465465
// dialWorkspaceAgent connects to a workspace agent by ID. Only rely on
466466
// r.Context() for cancellation if it's use is safe or r.Hijack() has
467467
// not been performed.
468-
func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.Conn, error) {
468+
func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (agent.Conn, error) {
469469
client, server := provisionersdk.TransportPipe()
470470
ctx, cancelFunc := context.WithCancel(context.Background())
471471
go func() {
@@ -524,7 +524,7 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
524524
<-peerConn.Closed()
525525
cancelFunc()
526526
}()
527-
return &agent.Conn{
527+
return &agent.WebRTCConn{
528528
Negotiator: peerClient,
529529
Conn: peerConn,
530530
}, nil

coderd/wsconncache/wsconncache_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ func TestCache(t *testing.T) {
3737
t.Parallel()
3838
t.Run("Same", func(t *testing.T) {
3939
t.Parallel()
40-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
40+
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (agent.Conn, error) {
4141
return setupAgent(t, agent.Metadata{}, 0), nil
4242
}, 0)
4343
t.Cleanup(func() {
@@ -52,7 +52,7 @@ func TestCache(t *testing.T) {
5252
t.Run("Expire", func(t *testing.T) {
5353
t.Parallel()
5454
called := atomic.NewInt32(0)
55-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
55+
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (agent.Conn, error) {
5656
called.Add(1)
5757
return setupAgent(t, agent.Metadata{}, 0), nil
5858
}, time.Microsecond)
@@ -71,7 +71,7 @@ func TestCache(t *testing.T) {
7171
})
7272
t.Run("NoExpireWhenLocked", func(t *testing.T) {
7373
t.Parallel()
74-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
74+
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (agent.Conn, error) {
7575
return setupAgent(t, agent.Metadata{}, 0), nil
7676
}, time.Microsecond)
7777
t.Cleanup(func() {
@@ -103,7 +103,7 @@ func TestCache(t *testing.T) {
103103
})
104104
go server.Serve(random)
105105

106-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
106+
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (agent.Conn, error) {
107107
return setupAgent(t, agent.Metadata{}, 0), nil
108108
}, time.Microsecond)
109109
t.Cleanup(func() {
@@ -139,7 +139,7 @@ func TestCache(t *testing.T) {
139139
})
140140
}
141141

142-
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) *agent.Conn {
142+
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) agent.Conn {
143143
client, server := provisionersdk.TransportPipe()
144144
closer := agent.New(agent.Options{
145145
FetchMetadata: func(ctx context.Context) (agent.Metadata, error) {
@@ -171,7 +171,7 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
171171
_ = conn.Close()
172172
})
173173

174-
return &agent.Conn{
174+
return &agent.WebRTCConn{
175175
Negotiator: api,
176176
Conn: conn,
177177
}

codersdk/workspaceagents.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ func (c *Client) WorkspaceAgentNodeBroker(ctx context.Context) (agent.NodeBroker
303303
}
304304

305305
// DialWorkspaceAgent creates a connection to the specified resource.
306-
func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, options *peer.ConnOptions) (*agent.Conn, error) {
306+
func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, options *peer.ConnOptions) (agent.Conn, error) {
307307
serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/dial", agentID.String()))
308308
if err != nil {
309309
return nil, xerrors.Errorf("parse url: %w", err)
@@ -368,7 +368,7 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
368368
if err != nil {
369369
return nil, xerrors.Errorf("dial peer: %w", err)
370370
}
371-
return &agent.Conn{
371+
return &agent.WebRTCConn{
372372
Negotiator: client,
373373
Conn: peerConn,
374374
}, nil

0 commit comments

Comments
 (0)