fix(port-forward): close dropped websocket tunnels

Author: onutc

api/port_forward.go

@@ -33,6 +33,8 @@ type portForwardControlMessage struct {
 	Type string `json:"type"`
 }
 
+var errPortForwardHalfClose = errors.New("port-forward half-close")
+
 func newPortForwardConfig() portForwardConfig {
 	return portForwardConfig{
 		enabled:         parseBoolEnv("SPRITZ_PORT_FORWARD_ENABLED", true),
@@ -164,22 +166,35 @@ func proxyWebSocketNetConn(ws *websocket.Conn, upstream net.Conn) error {
 		errCh <- copyNetConnToWebSocket(upstream, ws)
 	}()
 
-	var firstErr error
+	halfClosed := 0
 	for completed := 0; completed < 2; completed++ {
 		err := <-errCh
-		if err == nil || errors.Is(err, io.EOF) || websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
-			continue
-		}
-		if ne, ok := err.(net.Error); ok && ne.Timeout() {
-			continue
-		}
-		if firstErr == nil {
-			firstErr = err
+		switch {
+		case err == nil:
+			closeAll()
+			return nil
+		case errors.Is(err, errPortForwardHalfClose):
+			halfClosed++
+			if halfClosed == 2 {
+				closeAll()
+				return nil
+			}
+		case errors.Is(err, io.EOF), errors.Is(err, context.Canceled), websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway):
+			closeAll()
+			return nil
+		case func() bool {
+			ne, ok := err.(net.Error)
+			return ok && ne.Timeout()
+		}():
+			closeAll()
+			return nil
+		default:
 			closeAll()
+			return err
 		}
 	}
 	closeAll()
-	return firstErr
+	return nil
 }
 
 func (s *server) findPortForwardPod(ctx context.Context, namespace, name, container string) (*corev1.Pod, error) {
@@ -204,7 +219,7 @@ func copyWebSocketToNetConn(ws *websocket.Conn, upstream net.Conn) error {
 				if err := closeConnWrite(upstream); err != nil {
 					return err
 				}
-				return nil
+				return errPortForwardHalfClose
 			}
 			continue
 		}
@@ -234,6 +249,7 @@ func copyNetConnToWebSocket(upstream net.Conn, ws *websocket.Conn) error {
 				if writeErr := ws.WriteMessage(websocket.TextMessage, mustMarshalPortForwardControl(portForwardControlMessage{Type: "eof"})); writeErr != nil {
 					return writeErr
 				}
+				return errPortForwardHalfClose
 			}
 			return err
 		}

api/port_forward_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"io"
 	"net"
 	"net/http"
@@ -240,3 +241,92 @@ func TestOpenPortForwardPreservesEOFFramedExchange(t *testing.T) {
 		t.Fatalf("upstream exchange failed: %v", err)
 	}
 }
+
+func TestOpenPortForwardClosesUpstreamWhenWebSocketCloses(t *testing.T) {
+	scheme := newTestSpritzScheme(t)
+	spritz := &spritzv1.Spritz{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "tidal-falcon",
+			Namespace: "spritz-test",
+		},
+		Spec: spritzv1.SpritzSpec{
+			Owner: spritzv1.SpritzOwner{ID: "user-1"},
+		},
+	}
+
+	upstreamListener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen upstream: %v", err)
+	}
+	defer upstreamListener.Close()
+
+	upstreamDone := make(chan error, 1)
+	go func() {
+		conn, err := upstreamListener.Accept()
+		if err != nil {
+			upstreamDone <- err
+			return
+		}
+		defer conn.Close()
+		_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
+		buffer := make([]byte, 1)
+		_, err = conn.Read(buffer)
+		if err == nil || !errors.Is(err, io.EOF) {
+			upstreamDone <- err
+			return
+		}
+		upstreamDone <- nil
+	}()
+
+	s := &server{
+		client: ctrlclientfake.NewClientBuilder().
+			WithScheme(scheme).
+			WithObjects(spritz).
+			Build(),
+		scheme:    scheme,
+		namespace: "spritz-test",
+		auth: authConfig{
+			mode:              authModeHeader,
+			headerID:          "X-Spritz-User-Id",
+			headerDefaultType: principalTypeHuman,
+		},
+		internalAuth: internalAuthConfig{enabled: false},
+		portForward:  portForwardConfig{enabled: true, containerName: "spritz"},
+		findRunningPodFunc: func(ctx context.Context, namespace, name, container string) (*corev1.Pod, error) {
+			return &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "tidal-falcon-pod",
+					Namespace: namespace,
+				},
+			}, nil
+		},
+		openPodPortForwardFunc: func(ctx context.Context, pod *corev1.Pod, remotePort uint32) (net.Conn, io.Closer, error) {
+			conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", upstreamListener.Addr().String())
+			if err != nil {
+				return nil, nil, err
+			}
+			return conn, closeFunc(func() error { return nil }), nil
+		},
+	}
+
+	e := echo.New()
+	e.GET("/api/spritzes/:name/port-forward", s.openPortForward)
+	srv := httptest.NewServer(e)
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/api/spritzes/tidal-falcon/port-forward?port=3000"
+	headers := http.Header{}
+	headers.Set("X-Spritz-User-Id", "user-1")
+	headers.Set("Origin", srv.URL)
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
+	if err != nil {
+		t.Fatalf("dial websocket: %v", err)
+	}
+
+	if err := conn.Close(); err != nil {
+		t.Fatalf("close websocket: %v", err)
+	}
+	if err := <-upstreamDone; err != nil {
+		t.Fatalf("expected upstream to close after websocket exit: %v", err)
+	}
+}

cli/src/index.ts

@@ -1507,10 +1507,7 @@ async function bridgePortForwardSocket(socket: net.Socket, url: string, headers:
       }
       writePortForwardOutput(socket, data);
     };
-    const onWsClose = () => {
-      wsEnded = true;
-      maybeFinish();
-    };
+    const onWsClose = () => finish();
     const onWsError = (err: Error) => {
       if (!opened) {
         finish(err);

cli/test/port-forward.test.ts

@@ -410,6 +410,70 @@ test('port-forward preserves EOF-framed exchanges over websocket', async (t) =>
   assert.equal(exitCode, 0, `spz port-forward should exit cleanly: ${stderr}`);
 });
 
+test('port-forward closes the local client when the websocket tunnel drops after startup', async (t) => {
+  const localPort = await getFreePort();
+  const server = http.createServer();
+  const wss = new WebSocketServer({ noServer: true });
+  await listen(server);
+  t.after(() => {
+    wss.close();
+    server.close();
+  });
+  const address = server.address();
+  assert.ok(address && typeof address === 'object');
+
+  let connections = 0;
+  server.on('upgrade', (req, socket, head) => {
+    wss.handleUpgrade(req, socket, head, (ws) => {
+      wss.emit('connection', ws, req);
+    });
+  });
+  wss.on('connection', (ws) => {
+    connections += 1;
+    if (connections === 1) {
+      return;
+    }
+    ws.close(1011, 'boom');
+  });
+
+  const child = spawnCli(
+    ['port-forward', 'devbox1', '--transport', 'ws', '--namespace', 'spritz', '--local', String(localPort), '--remote', '4000'],
+    buildTestEnv(`http://127.0.0.1:${address.port}/api`),
+  );
+  let stderr = '';
+  const stderrBuffer = { value: '' };
+  child.stderr.on('data', (chunk) => {
+    const text = chunk.toString();
+    stderr += text;
+    stderrBuffer.value += text;
+  });
+  t.after(() => {
+    child.kill('SIGTERM');
+  });
+
+  await waitForPattern(stderrBuffer, new RegExp(`forwarding 127\\.0\\.0\\.1:${localPort}`));
+
+  const client = net.connect(localPort, '127.0.0.1');
+  const clientClosed = new Promise<void>((resolve, reject) => {
+    const timer = setTimeout(() => reject(new Error('timed out waiting for local socket to close')), 2000);
+    client.on('error', () => {
+      clearTimeout(timer);
+      resolve();
+    });
+    client.on('close', () => {
+      clearTimeout(timer);
+      resolve();
+    });
+  });
+
+  await clientClosed;
+  client.destroy();
+
+  child.kill('SIGTERM');
+  const exitCode = await new Promise<number | null>((resolve) => child.on('exit', resolve));
+  assert.equal(exitCode, 0, `spz port-forward should exit cleanly: ${stderr}`);
+});
+
 test('port-forward falls back to SSH when websocket startup validation is rejected by default', async (t) => {
   const tempDir = mkdtempSync(path.join(os.tmpdir(), 'spz-port-forward-'));
   const fakeKeygen = path.join(tempDir, 'ssh-keygen');