pkg/hostagent: Support port forwarding to direct guest IP

Change to use `github.com/inetaf/tcpproxy`:
- pkg/driver/vz/vsock_forwarder.go
- pkg/portfwd/client.go
- pkg/portfwdserver/server.go

pkg/portfwd:
- Use `dialContext` instead of `client` to use `net.Conn` other than `GrpcClientRW`.
- Change to create proxies from `dialContext` parameters instead of `conn`.
- Add `DialContextToGRPCTunnel()` to return `DialContext` function that connects to GRPC tunnel.

pkg/hostagent:
- Add `HostAgent.DialContextToGuestIP()` to return `DialContext` function that connects to the guest IP directly.
  If the guest IP is not known, it returns nil.
- Prefer `HostAgent.DialContextToGuestIP()` over `portfwd.DialContextToGRPCTunnel()`.

Signed-off-by: Norio Nomura <norio.nomura@gmail.com>

pkg/hostagent: Use `HostAgent.DialContextToGuestIP()` if the IP is accessible directly.

Signed-off-by: Norio Nomura <norio.nomura@gmail.com>

Revert to "github.com/containers/gvisor-tap-vsock/pkg/tcpproxy"

Signed-off-by: Norio Nomura <norio.nomura@gmail.com>

# Conflicts:
#	pkg/hostagent/hostagent.go

Author: norio-nomura

pkg/hostagent/hostagent.go

@@ -901,7 +901,19 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 		if useSSHFwd {
 			a.portForwarder.OnEvent(ctx, ev)
 		} else {
-			a.grpcPortForwarder.OnEvent(ctx, client, ev)
+			dialContext := func(ctx context.Context, network, guestAddress string) (net.Conn, error) {
+				guestIP := a.GuestIP()
+				if guestIP == nil {
+					return portfwd.DialContextToGRPCTunnel(client)(ctx, network, guestAddress)
+				}
+				if host, _, err := net.SplitHostPort(guestAddress); err != nil {
+					return nil, err
+				} else if ip := net.ParseIP(host); ip.IsUnspecified() || ip.Equal(guestIP) {
+					return DialContextToGuestIP(guestIP)(ctx, network, guestAddress)
+				}
+				return portfwd.DialContextToGRPCTunnel(client)(ctx, network, guestAddress)
+			}
+			a.grpcPortForwarder.OnEvent(ctx, dialContext, ev)
 		}
 	}
 
@@ -914,6 +926,24 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 	return io.EOF
 }
 
+// DialContextToGuestIP returns a DialContext function that connects to the guest IP directly.
+// If the guest IP is not known, it returns nil.
+func DialContextToGuestIP(guestIP net.IP) func(ctx context.Context, network, address string) (net.Conn, error) {
+	if guestIP == nil {
+		return nil
+	}
+	return func(ctx context.Context, network, address string) (net.Conn, error) {
+		var d net.Dialer
+		_, port, err := net.SplitHostPort(address)
+		if err != nil {
+			return nil, err
+		}
+		// Host part of address is ignored, because it already has been checked by forwarding rules
+		// and we want to connect to the guest IP directly.
+		return d.DialContext(ctx, network, net.JoinHostPort(guestIP.String(), port))
+	}
+}
+
 const (
 	verbForward = "forward"
 	verbCancel  = "cancel"

pkg/portfwd/client.go

@@ -18,61 +18,50 @@ import (
 	guestagentclient "github.com/lima-vm/lima/v2/pkg/guestagent/api/client"
 )
 
-func HandleTCPConnection(ctx context.Context, client *guestagentclient.GuestAgentClient, conn net.Conn, guestAddr string) {
-	id := fmt.Sprintf("tcp-%s-%s", conn.LocalAddr().String(), conn.RemoteAddr().String())
+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)
+}
 
-	stream, err := client.Tunnel(ctx)
+func HandleUDPConnection(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), conn net.PacketConn, guestAddr string) {
+	proxy, err := forwarder.NewUDPProxy(conn, func() (net.Conn, error) {
+		return dialContext(ctx, "udp", guestAddr)
+	})
 	if err != nil {
-		logrus.Errorf("could not open tcp tunnel for id: %s error:%v", id, err)
-		return
-	}
-
-	// Handshake message to start tunnel
-	if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: "tcp", GuestAddr: guestAddr}); err != nil {
-		logrus.Errorf("could not start tcp tunnel for id: %s error:%v", id, err)
+		logrus.WithError(err).Error("error in udp tunnel proxy")
 		return
 	}
 
-	rw := &GrpcClientRW{stream: stream, id: id, addr: guestAddr, protocol: "tcp"}
-	proxy := tcpproxy.DialProxy{DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-		return conn, nil
-	}}
-	proxy.HandleConn(rw)
+	defer func() {
+		err := proxy.Close()
+		if err != nil {
+			logrus.WithError(err).Error("error in closing udp tunnel proxy")
+		}
+	}()
+	proxy.Run()
 }
 
-func HandleUDPConnection(ctx context.Context, client *guestagentclient.GuestAgentClient, conn net.PacketConn, guestAddr string) {
-	var udpConnectionCounter atomic.Uint32
-	initialID := fmt.Sprintf("udp-%s", conn.LocalAddr().String())
-
+func DialContextToGRPCTunnel(client *guestagentclient.GuestAgentClient) func(ctx context.Context, network, addr string) (net.Conn, error) {
 	// gvisor-tap-vsock's UDPProxy demultiplexes client connections internally based on their source address.
 	// It calls this dialer function only when it receives a datagram from a new, unrecognized client.
 	// For each new client, we must return a new net.Conn, which in our case is a new gRPC stream.
 	// The atomic counter ensures that each stream has a unique ID to distinguish them on the server side.
-	proxy, err := forwarder.NewUDPProxy(conn, func() (net.Conn, error) {
-		id := fmt.Sprintf("%s-%d", initialID, udpConnectionCounter.Add(1))
-		stream, err := client.Tunnel(ctx)
+	var connectionCounter atomic.Uint32
+	return func(_ context.Context, network, addr string) (net.Conn, error) {
+		// Passed context.Context is used for timeout on initiate connection, not for the lifetime of the connection.
+		// We use context.Background() here to avoid unexpected cancellation.
+		stream, err := client.Tunnel(context.Background())
 		if err != nil {
-			return nil, fmt.Errorf("could not open udp tunnel for id: %s error:%w", id, err)
+			return nil, fmt.Errorf("could not open tunnel for addr: %s error:%w", addr, err)
 		}
 		// Handshake message to start tunnel
-		if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: "udp", GuestAddr: guestAddr}); err != nil {
-			return nil, fmt.Errorf("could not start udp tunnel for id: %s error:%w", id, err)
+		id := fmt.Sprintf("%s-%s-%d", network, addr, connectionCounter.Add(1))
+		if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: network, GuestAddr: addr}); err != nil {
+			return nil, fmt.Errorf("could not start tunnel for id: %s addr: %s error:%w", id, addr, err)
 		}
-		rw := &GrpcClientRW{stream: stream, id: id, addr: guestAddr, protocol: "udp"}
+		rw := &GrpcClientRW{stream: stream, id: id, addr: addr, protocol: network}
 		return rw, nil
-	})
-	if err != nil {
-		logrus.Errorf("error in udp tunnel proxy for id: %s error:%v", initialID, err)
-		return
 	}
-
-	defer func() {
-		err := proxy.Close()
-		if err != nil {
-			logrus.Errorf("error in closing udp tunnel proxy for id: %s error:%v", initialID, err)
-		}
-	}()
-	proxy.Run()
 }
 
 type GrpcClientRW struct {

pkg/portfwd/forward.go

@@ -11,7 +11,6 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/lima-vm/lima/v2/pkg/guestagent/api"
-	guestagentclient "github.com/lima-vm/lima/v2/pkg/guestagent/api/client"
 	"github.com/lima-vm/lima/v2/pkg/limatype"
 	"github.com/lima-vm/lima/v2/pkg/limayaml"
 )
@@ -38,7 +37,7 @@ func (fw *Forwarder) Close() error {
 	return fw.closableListeners.Close()
 }
 
-func (fw *Forwarder) OnEvent(ctx context.Context, client *guestagentclient.GuestAgentClient, ev *api.Event) {
+func (fw *Forwarder) OnEvent(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), ev *api.Event) {
 	for _, f := range ev.AddedLocalPorts {
 		// Before forwarding, check if any static rule matches this port otherwise it will be forwarded twice and cause a port conflict
 		if fw.isPortStaticallyForwarded(f) {
@@ -55,7 +54,7 @@ func (fw *Forwarder) OnEvent(ctx context.Context, client *guestagentclient.Guest
 			continue
 		}
 		logrus.Infof("Forwarding %s from %s to %s", strings.ToUpper(f.Protocol), remote, local)
-		fw.closableListeners.Forward(ctx, client, f.Protocol, local, remote)
+		fw.closableListeners.Forward(ctx, dialContext, f.Protocol, local, remote)
 	}
 	for _, f := range ev.RemovedLocalPorts {
 		local, remote := fw.forwardingAddresses(f)

pkg/portfwd/listener.go

@@ -14,8 +14,6 @@ import (
 	"sync"
 
 	"github.com/sirupsen/logrus"
-
-	guestagentclient "github.com/lima-vm/lima/v2/pkg/guestagent/api/client"
 )
 
 type ClosableListeners struct {
@@ -59,14 +57,14 @@ func (p *ClosableListeners) Close() error {
 	return errors.Join(errs...)
 }
 
-func (p *ClosableListeners) Forward(ctx context.Context, client *guestagentclient.GuestAgentClient,
+func (p *ClosableListeners) Forward(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error),
 	protocol string, hostAddress string, guestAddress string,
 ) {
 	switch protocol {
 	case "tcp", "tcp6":
-		go p.forwardTCP(ctx, client, hostAddress, guestAddress)
+		go p.forwardTCP(ctx, dialContext, hostAddress, guestAddress)
 	case "udp", "udp6":
-		go p.forwardUDP(ctx, client, hostAddress, guestAddress)
+		go p.forwardUDP(ctx, dialContext, hostAddress, guestAddress)
 	}
 }
 
@@ -93,7 +91,7 @@ func (p *ClosableListeners) Remove(_ context.Context, protocol, hostAddress, gue
 	}
 }
 
-func (p *ClosableListeners) forwardTCP(ctx context.Context, client *guestagentclient.GuestAgentClient, hostAddress, guestAddress string) {
+func (p *ClosableListeners) forwardTCP(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), hostAddress, guestAddress string) {
 	key := key("tcp", hostAddress, guestAddress)
 
 	p.listenersRW.Lock()
@@ -124,11 +122,11 @@ func (p *ClosableListeners) forwardTCP(ctx context.Context, client *guestagentcl
 			}
 			return
 		}
-		go HandleTCPConnection(ctx, client, conn, guestAddress)
+		go HandleTCPConnection(ctx, dialContext, conn, guestAddress)
 	}
 }
 
-func (p *ClosableListeners) forwardUDP(ctx context.Context, client *guestagentclient.GuestAgentClient, hostAddress, guestAddress string) {
+func (p *ClosableListeners) forwardUDP(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), hostAddress, guestAddress string) {
 	key := key("udp", hostAddress, guestAddress)
 	defer p.Remove(ctx, "udp", hostAddress, guestAddress)
 
@@ -148,7 +146,7 @@ func (p *ClosableListeners) forwardUDP(ctx context.Context, client *guestagentcl
 	p.udpListeners[key] = udpConn
 	p.udpListenersRW.Unlock()
 
-	HandleUDPConnection(ctx, client, udpConn, guestAddress)
+	HandleUDPConnection(ctx, dialContext, udpConn, guestAddress)
 }
 
 func key(protocol, hostAddress, guestAddress string) string {