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

The documentation says:
https://pkg.go.dev/google.golang.org/grpc#BidiStreamingServer
> To terminate the stream, return from the handler method and return an error from the status package, or use nil to indicate an OK status code.

- Changed to calling `proxy.HandleConn()`/`bicopy.Bicopy()` as goroutines, then wait `GRPCServerRW.closeCh` in `TunnelServer.Start()` for returning from handler before `HandleConn()`/`bicopy.Bicopy()` finishes.
- Added `CloseRead()` and `CloseWrite()` to `GRPCServerRW`. As a result, `GRPCServerRW` may be expected to pass the test added by #3708 with inetaf/tcpproxy's `tcpproxy.DialProxy()` or `bicopy.Bicopy()`.

Signed-off-by: Norio Nomura <[email protected]>

Author: norio-nomura

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

@@ -13,6 +13,7 @@ import (
 	"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,7 +42,11 @@ func (s *TunnelServer) Start(stream api.GuestService_TunnelServer) error {
 	if err != nil {
 		return err
 	}
-	rw := &GRPCServerRW{stream: stream, id: in.Id}
+	rw := &GRPCServerRW{stream: stream, id: in.Id, closeCh: make(chan any, 1)}
+	go func() {
+		<-ctx.Done()
+		rw.Close()
+	}()
 
 	// FIXME: consolidate bicopy and tcpproxy into one
 	//
@@ -51,20 +56,26 @@ func (s *TunnelServer) Start(stream api.GuestService_TunnelServer) error {
 	// 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)
+		go bicopy.Bicopy(rw, conn, nil)
 	} else {
 		proxy := tcpproxy.DialProxy{DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
 			return conn, nil
 		}}
-		proxy.HandleConn(rw)
+		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 +95,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
 }