Merge pull request #4237 from norio-nomura/fix-closing-grpc-stream-in-guestagent

pkg/portfwdserver: Close `stream` by returning from the handler method

Author: AkihiroSuda

.github/workflows/test.yml

@@ -177,7 +177,7 @@ jobs:
     - name: Integration tests (WSL2, Windows host)
       run: |
         $env:PATH = "$pwd\_output\bin;" + 'C:\msys64\usr\bin;' + $env:PATH
-        pacman -Sy --noconfirm openbsd-netcat diffutils socat
+        pacman -Sy --noconfirm openbsd-netcat diffutils socat w3m
         $env:MSYS2_ENV_CONV_EXCL = 'HOME_HOST;HOME_GUEST;_LIMA_WINDOWS_EXTRA_PATH'
         $env:HOME_HOST = $(cygpath.exe "$env:USERPROFILE")
         $env:HOME_GUEST = "/mnt$env:HOME_HOST"
@@ -206,7 +206,7 @@ jobs:
     - name: Integration tests (QEMU, Windows host)
       run: |
         $env:PATH = "$pwd\_output\bin;" + 'C:\msys64\usr\bin;' + 'C:\Program Files\QEMU;' + $env:PATH
-        pacman -Sy --noconfirm openbsd-netcat diffutils socat
+        pacman -Sy --noconfirm openbsd-netcat diffutils socat w3m
         $env:MSYS2_ENV_CONV_EXCL = 'HOME_HOST;HOME_GUEST;_LIMA_WINDOWS_EXTRA_PATH'
         $env:HOME_HOST = $(cygpath.exe "$env:USERPROFILE")
         $env:HOME_GUEST = "$env:HOME_HOST"

hack/test-templates.sh

@@ -489,7 +489,7 @@ if [[ -n ${CHECKS["port-forwards"]} ]]; then
 				limactl shell "$NAME" $sudo $CONTAINER_ENGINE rm -f nginx
 			fi
 		fi
-		if [[ ${NAME} != "alpine"* && ${NAME} != "wsl2"* ]] && command -v w3m >/dev/null; then
+		if [[ ${NAME} != "alpine"* ]] && command -v w3m >/dev/null; then
 			INFO "Testing https://github.com/lima-vm/lima/issues/3685 ([gRPC portfwd] client connection is not closed immediately when server closed the connection)"
 			# Skip the test on Alpine, as systemd-run is missing
 			# Skip the test on WSL2, as port forwarding is half broken https://github.com/lima-vm/lima/pull/3686#issuecomment-3034842616

pkg/portfwd/client.go

@@ -21,6 +21,7 @@ import (
 func HandleTCPConnection(_ context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), conn net.Conn, guestAddr string) {
 	proxy := tcpproxy.DialProxy{Addr: guestAddr, DialContext: dialContext}
 	proxy.HandleConn(conn)
+	logrus.Debugf("tcp proxy for guestAddr: %s closed", guestAddr)
 }
 
 func HandleUDPConnection(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), conn net.PacketConn, guestAddr string) {
@@ -39,6 +40,7 @@ func HandleUDPConnection(ctx context.Context, dialContext func(ctx context.Conte
 		}
 	}()
 	proxy.Run()
+	logrus.Debugf("udp proxy for guestAddr: %s closed", guestAddr)
 }
 
 func DialContextToGRPCTunnel(client *guestagentclient.GuestAgentClient) func(ctx context.Context, network, addr string) (net.Conn, error) {
@@ -97,6 +99,7 @@ func (g *GrpcClientRW) Read(p []byte) (n int, err error) {
 }
 
 func (g *GrpcClientRW) Close() error {
+	logrus.Debugf("closing GrpcClientRW for id: %s", g.id)
 	return g.stream.CloseSend()
 }
 

pkg/portfwdserver/server.go

@@ -8,13 +8,11 @@ import (
 	"errors"
 	"io"
 	"net"
-	"os"
-	"strings"
 	"time"
 
 	"github.com/containers/gvisor-tap-vsock/pkg/tcpproxy"
+	"github.com/sirupsen/logrus"
 
-	"github.com/lima-vm/lima/v2/pkg/bicopy"
 	"github.com/lima-vm/lima/v2/pkg/guestagent/api"
 )
 
@@ -41,30 +39,30 @@ func (s *TunnelServer) Start(stream api.GuestService_TunnelServer) error {
 	if err != nil {
 		return err
 	}
-	rw := &GRPCServerRW{stream: stream, id: in.Id}
-
-	// FIXME: consolidate bicopy and tcpproxy into one
-	//
-	// The bicopy package does not seem to work with `w3m -dump`:
-	// https://github.com/lima-vm/lima/issues/3685
-	//
-	// However, the tcpproxy package can't pass the CI for WSL2 (experimental):
-	// https://github.com/lima-vm/lima/pull/3686#issuecomment-3034842616
-	if wsl2, _ := seemsWSL2(); wsl2 {
-		bicopy.Bicopy(rw, conn, nil)
-	} else {
-		proxy := tcpproxy.DialProxy{DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-			return conn, nil
-		}}
-		proxy.HandleConn(rw)
-	}
+	rw := &GRPCServerRW{stream: stream, id: in.Id, closeCh: make(chan any, 1)}
+	go func() {
+		<-ctx.Done()
+		rw.Close()
+	}()
+
+	proxy := tcpproxy.DialProxy{DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+		return conn, nil
+	}}
+	go proxy.HandleConn(rw)
+
+	// The stream will be closed when this function returns.
+	// Wait here until rw.Close(), rw.CloseRead(), or rw.CloseWrite() is called.
+	// We can't close rw.closeCh since the calling order of Close* methods is not guaranteed.
+	<-rw.closeCh
+	logrus.Debugf("closed GRPCServerRW for id: %s", in.Id)
 
 	return nil
 }
 
 type GRPCServerRW struct {
-	id     string
-	stream api.GuestService_TunnelServer
+	id      string
+	stream  api.GuestService_TunnelServer
+	closeCh chan any
 }
 
 var _ net.Conn = (*GRPCServerRW)(nil)
@@ -84,6 +82,23 @@ func (g *GRPCServerRW) Read(p []byte) (n int, err error) {
 }
 
 func (g *GRPCServerRW) Close() error {
+	logrus.Debugf("closing GRPCServerRW for id: %s", g.id)
+	g.closeCh <- struct{}{}
+	return nil
+}
+
+// By adding CloseRead and CloseWrite methods, GRPCServerRW can work with
+// other than containers/gvisor-tap-vsock/pkg/tcpproxy, e.g., inetaf/tcpproxy, bicopy.Bicopy.
+
+func (g *GRPCServerRW) CloseRead() error {
+	logrus.Debugf("closing read GRPCServerRW for id: %s", g.id)
+	g.closeCh <- struct{}{}
+	return nil
+}
+
+func (g *GRPCServerRW) CloseWrite() error {
+	logrus.Debugf("closing write GRPCServerRW for id: %s", g.id)
+	g.closeCh <- struct{}{}
 	return nil
 }
 
@@ -106,13 +121,3 @@ func (g *GRPCServerRW) SetReadDeadline(_ time.Time) error {
 func (g *GRPCServerRW) SetWriteDeadline(_ time.Time) error {
 	return nil
 }
-
-// seemsWSL2 returns whether lima.env contains LIMA_CIDATA_VMTYPE=wsl2 .
-// This is a temporary workaround and has to be removed.
-func seemsWSL2() (bool, error) {
-	b, err := os.ReadFile("/mnt/lima-cidata/lima.env")
-	if err != nil {
-		return false, err
-	}
-	return strings.Contains(string(b), "LIMA_CIDATA_VMTYPE=wsl2"), nil
-}