Visitar URL original
Fix concurrent map writes in codespace port forwarding (#13313) · cli/cli@74e7791 · GitHub
Skip to content

Commit 74e7791

Browse files
williammartinCopilotBagToad
authored
Fix concurrent map writes in codespace port forwarding (#13313)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Kynan Ware <47394200+BagToad@users.noreply.github.com>
1 parent 6dae307 commit 74e7791

3 files changed

Lines changed: 109 additions & 78 deletions

File tree

‎internal/codespaces/connection/connection.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@ type CodespaceConnection struct {
3030
Options *tunnels.TunnelRequestOptions
3131
Tunnel *tunnels.Tunnel
3232
AllowedPortPrivacySettings []string
33+
34+
// ManagerMu serializes access to TunnelManager operations which mutate
35+
// shared state on the Tunnel object and are not goroutine-safe.
36+
ManagerMu sync.Mutex
3337
}
3438

3539
// NewCodespaceConnection initializes a connection to a codespace.

‎internal/codespaces/portforwarder/port_forwarder.go‎

Lines changed: 36 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ type ForwardPortOpts struct {
3636
}
3737

3838
type CodespacesPortForwarder struct {
39-
connection connection.CodespaceConnection
39+
connection *connection.CodespaceConnection
4040
keepAliveReason chan string
4141
}
4242

@@ -54,7 +54,7 @@ type PortForwarder interface {
5454
// NewPortForwarder returns a new PortForwarder for the specified codespace.
5555
func NewPortForwarder(ctx context.Context, codespaceConnection *connection.CodespaceConnection) (fwd PortForwarder, err error) {
5656
return &CodespacesPortForwarder{
57-
connection: *codespaceConnection,
57+
connection: codespaceConnection,
5858
keepAliveReason: make(chan string, 1),
5959
}, nil
6060
}
@@ -108,6 +108,32 @@ func (fwd *CodespacesPortForwarder) ForwardPort(ctx context.Context, opts Forwar
108108
return fmt.Errorf("error converting port: %w", err)
109109
}
110110

111+
if err := fwd.createTunnelPort(ctx, port, opts); err != nil {
112+
return err
113+
}
114+
115+
// Connect to the tunnel
116+
err = fwd.connection.Connect(ctx)
117+
if err != nil {
118+
return fmt.Errorf("connect failed: %v", err)
119+
}
120+
121+
// Inform the host that we've forwarded the port locally
122+
err = fwd.connection.TunnelClient.RefreshPorts(ctx)
123+
if err != nil {
124+
return fmt.Errorf("refresh ports failed: %v", err)
125+
}
126+
127+
return nil
128+
}
129+
130+
// createTunnelPort creates a tunnel port while holding the manager mutex.
131+
// TunnelManager operations mutate shared state on the Tunnel object and are
132+
// not goroutine-safe, so all calls are serialized under ManagerMu.
133+
func (fwd *CodespacesPortForwarder) createTunnelPort(ctx context.Context, port uint16, opts ForwardPortOpts) error {
134+
fwd.connection.ManagerMu.Lock()
135+
defer fwd.connection.ManagerMu.Unlock()
136+
111137
// In v0.0.25 of dev-tunnels, the dev-tunnel manager `CreateTunnelPort` would "accept" requests that
112138
// change the port protocol but they would not result in any actual change. This has changed, resulting in
113139
// an error `Invalid arguments. The tunnel port protocol cannot be changed.`. It's not clear why the previous
@@ -166,18 +192,6 @@ func (fwd *CodespacesPortForwarder) ForwardPort(ctx context.Context, opts Forwar
166192
return fmt.Errorf("create tunnel port failed: %v", err)
167193
}
168194

169-
// Connect to the tunnel
170-
err = fwd.connection.Connect(ctx)
171-
if err != nil {
172-
return fmt.Errorf("connect failed: %v", err)
173-
}
174-
175-
// Inform the host that we've forwarded the port locally
176-
err = fwd.connection.TunnelClient.RefreshPorts(ctx)
177-
if err != nil {
178-
return fmt.Errorf("refresh ports failed: %v", err)
179-
}
180-
181195
return nil
182196
}
183197

@@ -243,6 +257,9 @@ func (fwd *CodespacesPortForwarder) ConnectToForwardedPort(ctx context.Context,
243257

244258
// ListPorts fetches the list of ports that are currently forwarded.
245259
func (fwd *CodespacesPortForwarder) ListPorts(ctx context.Context) (ports []*tunnels.TunnelPort, err error) {
260+
fwd.connection.ManagerMu.Lock()
261+
defer fwd.connection.ManagerMu.Unlock()
262+
246263
ports, err = fwd.connection.TunnelManager.ListTunnelPorts(ctx, fwd.connection.Tunnel, fwd.connection.Options)
247264
if err != nil {
248265
return nil, fmt.Errorf("error listing ports: %w", err)
@@ -253,22 +270,27 @@ func (fwd *CodespacesPortForwarder) ListPorts(ctx context.Context) (ports []*tun
253270

254271
// UpdatePortVisibility changes the visibility (private, org, public) of the specified port.
255272
func (fwd *CodespacesPortForwarder) UpdatePortVisibility(ctx context.Context, remotePort int, visibility string) error {
273+
fwd.connection.ManagerMu.Lock()
256274
tunnelPort, err := fwd.connection.TunnelManager.GetTunnelPort(ctx, fwd.connection.Tunnel, remotePort, fwd.connection.Options)
257275
if err != nil {
276+
fwd.connection.ManagerMu.Unlock()
258277
return fmt.Errorf("error getting tunnel port: %w", err)
259278
}
260279

261280
// If the port visibility isn't changing, don't do anything
262281
if AccessControlEntriesToVisibility(tunnelPort.AccessControl.Entries) == visibility {
282+
fwd.connection.ManagerMu.Unlock()
263283
return nil
264284
}
265285

266286
// Delete the existing tunnel port to update
267287
port, err := convertIntToUint16(remotePort)
268288
if err != nil {
289+
fwd.connection.ManagerMu.Unlock()
269290
return fmt.Errorf("error converting port: %w", err)
270291
}
271292
err = fwd.connection.TunnelManager.DeleteTunnelPort(ctx, fwd.connection.Tunnel, port, fwd.connection.Options)
293+
fwd.connection.ManagerMu.Unlock()
272294
if err != nil {
273295
return fmt.Errorf("error deleting tunnel port: %w", err)
274296
}

‎internal/codespaces/portforwarder/port_forwarder_test.go‎

Lines changed: 69 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ import (
77
"github.com/cli/cli/v2/internal/codespaces/api"
88
"github.com/cli/cli/v2/internal/codespaces/connection"
99
"github.com/microsoft/dev-tunnels/go/tunnels"
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/require"
12+
"golang.org/x/sync/errgroup"
1013
)
1114

1215
func TestNewPortForwarder(t *testing.T) {
@@ -31,26 +34,16 @@ func TestNewPortForwarder(t *testing.T) {
3134

3235
// Create the mock HTTP client
3336
httpClient, err := connection.NewMockHttpClient()
34-
if err != nil {
35-
t.Fatalf("NewHttpClient returned an error: %v", err)
36-
}
37+
require.NoError(t, err)
3738

3839
// Call the function being tested
3940
conn, err := connection.NewCodespaceConnection(ctx, codespace, httpClient)
40-
if err != nil {
41-
t.Fatalf("NewCodespaceConnection returned an error: %v", err)
42-
}
41+
require.NoError(t, err)
4342

4443
// Create the new port forwarder
4544
portForwarder, err := NewPortForwarder(ctx, conn)
46-
if err != nil {
47-
t.Fatalf("NewPortForwarder returned an error: %v", err)
48-
}
49-
50-
// Check that the port forwarder was created successfully
51-
if portForwarder == nil {
52-
t.Fatal("NewPortForwarder returned nil")
53-
}
45+
require.NoError(t, err)
46+
require.NotNil(t, portForwarder)
5447
}
5548

5649
func TestAccessControlEntriesToVisibility(t *testing.T) {
@@ -96,9 +89,7 @@ func TestAccessControlEntriesToVisibility(t *testing.T) {
9689
for _, test := range tests {
9790
t.Run(test.name, func(t *testing.T) {
9891
visibility := AccessControlEntriesToVisibility(test.accessControlEntries)
99-
if visibility != test.expected {
100-
t.Errorf("expected %q, got %q", test.expected, visibility)
101-
}
92+
assert.Equal(t, test.expected, visibility)
10293
})
10394
}
10495
}
@@ -131,9 +122,7 @@ func TestIsInternalPort(t *testing.T) {
131122
for _, test := range tests {
132123
t.Run(test.name, func(t *testing.T) {
133124
isInternal := IsInternalPort(test.port)
134-
if isInternal != test.expected {
135-
t.Errorf("expected %v, got %v", test.expected, isInternal)
136-
}
125+
assert.Equal(t, test.expected, isInternal)
137126
})
138127
}
139128
}
@@ -163,39 +152,70 @@ func TestForwardPortDefaultsToHTTPProtocol(t *testing.T) {
163152
httpClient, err := connection.NewMockHttpClient(
164153
connection.WithSpecificPorts(tunnelPorts),
165154
)
166-
if err != nil {
167-
t.Fatalf("NewMockHttpClient returned an error: %v", err)
168-
}
155+
require.NoError(t, err)
169156

170157
connection, err := connection.NewCodespaceConnection(t.Context(), codespace, httpClient)
171-
if err != nil {
172-
t.Fatalf("NewCodespaceConnection returned an error: %v", err)
173-
}
158+
require.NoError(t, err)
174159

175160
fwd, err := NewPortForwarder(t.Context(), connection)
176-
if err != nil {
177-
t.Fatalf("NewPortForwarder returned an error: %v", err)
178-
}
161+
require.NoError(t, err)
179162

180163
// When we forward a port without an existing one to use for a protocol, it should default to HTTP.
181-
if err := fwd.ForwardPort(t.Context(), ForwardPortOpts{
164+
err = fwd.ForwardPort(t.Context(), ForwardPortOpts{
182165
Port: 1337,
183-
}); err != nil {
184-
t.Fatalf("ForwardPort returned an error: %v", err)
185-
}
166+
})
167+
require.NoError(t, err)
186168

187169
ports, err := fwd.ListPorts(t.Context())
188-
if err != nil {
189-
t.Fatalf("ListPorts returned an error: %v", err)
190-
}
170+
require.NoError(t, err)
171+
require.Len(t, ports, 1)
172+
assert.Equal(t, string(tunnels.TunnelProtocolHttp), ports[0].Protocol)
173+
}
191174

192-
if len(ports) != 1 {
193-
t.Fatalf("expected 1 port, got %d", len(ports))
175+
func TestConcurrentForwardPortDoesNotRace(t *testing.T) {
176+
codespace := &api.Codespace{
177+
Name: "codespace-name",
178+
State: api.CodespaceStateAvailable,
179+
Connection: api.CodespaceConnection{
180+
TunnelProperties: api.TunnelProperties{
181+
ConnectAccessToken: "tunnel access-token",
182+
ManagePortsAccessToken: "manage-ports-token",
183+
ServiceUri: "http://global.rel.tunnels.api.visualstudio.com/",
184+
TunnelId: "tunnel-id",
185+
ClusterId: "usw2",
186+
Domain: "domain.com",
187+
},
188+
},
189+
RuntimeConstraints: api.RuntimeConstraints{
190+
AllowedPortPrivacySettings: []string{"public", "private"},
191+
},
194192
}
195193

196-
if ports[0].Protocol != string(tunnels.TunnelProtocolHttp) {
197-
t.Fatalf("expected port protocol to be http, got %s", ports[0].Protocol)
194+
tunnelPorts := map[int]tunnels.TunnelPort{}
195+
196+
httpClient, err := connection.NewMockHttpClient(
197+
connection.WithSpecificPorts(tunnelPorts),
198+
)
199+
require.NoError(t, err)
200+
201+
conn, err := connection.NewCodespaceConnection(t.Context(), codespace, httpClient)
202+
require.NoError(t, err)
203+
204+
// Forward multiple ports concurrently from the same connection,
205+
// mirroring what ForwardPorts does in ports.go.
206+
group, ctx := errgroup.WithContext(t.Context())
207+
for port := 3000; port < 3010; port++ {
208+
fwd, err := NewPortForwarder(ctx, conn)
209+
require.NoError(t, err)
210+
211+
group.Go(func() error {
212+
return fwd.ForwardPort(ctx, ForwardPortOpts{
213+
Port: port,
214+
})
215+
})
198216
}
217+
218+
require.NoError(t, group.Wait())
199219
}
200220

201221
func TestForwardPortRespectsProtocolOfExistingTunneledPorts(t *testing.T) {
@@ -230,38 +250,23 @@ func TestForwardPortRespectsProtocolOfExistingTunneledPorts(t *testing.T) {
230250
httpClient, err := connection.NewMockHttpClient(
231251
connection.WithSpecificPorts(tunnelPorts),
232252
)
233-
if err != nil {
234-
t.Fatalf("NewMockHttpClient returned an error: %v", err)
235-
}
253+
require.NoError(t, err)
236254

237255
connection, err := connection.NewCodespaceConnection(t.Context(), codespace, httpClient)
238-
if err != nil {
239-
t.Fatalf("NewCodespaceConnection returned an error: %v", err)
240-
}
256+
require.NoError(t, err)
241257

242258
fwd, err := NewPortForwarder(t.Context(), connection)
243-
if err != nil {
244-
t.Fatalf("NewPortForwarder returned an error: %v", err)
245-
}
259+
require.NoError(t, err)
246260

247261
// When we forward a port, it would typically default to HTTP, to which the mock server would respond with a 400,
248262
// but it should respect the existing port's protocol and forward it as HTTPS.
249-
if err := fwd.ForwardPort(t.Context(), ForwardPortOpts{
263+
err = fwd.ForwardPort(t.Context(), ForwardPortOpts{
250264
Port: 1337,
251-
}); err != nil {
252-
t.Fatalf("ForwardPort returned an error: %v", err)
253-
}
265+
})
266+
require.NoError(t, err)
254267

255268
ports, err := fwd.ListPorts(t.Context())
256-
if err != nil {
257-
t.Fatalf("ListPorts returned an error: %v", err)
258-
}
259-
260-
if len(ports) != 1 {
261-
t.Fatalf("expected 1 port, got %d", len(ports))
262-
}
263-
264-
if ports[0].Protocol != string(tunnels.TunnelProtocolHttps) {
265-
t.Fatalf("expected port protocol to be https, got %s", ports[0].Protocol)
266-
}
269+
require.NoError(t, err)
270+
require.Len(t, ports, 1)
271+
assert.Equal(t, string(tunnels.TunnelProtocolHttps), ports[0].Protocol)
267272
}

0 commit comments

Comments
 (0)