portfwd: create separate gRPC streams for each UDP client

The UDP port forwarder previously used a single gRPC stream for all
clients, which could cause responses from the guest to be sent to the
wrong client on the host.

This occurred because the stream was created before client connections
were demultiplexed by `gvisor-tap-vsock`'s `UDPProxy`.

The root cause is the interaction with `gvisor-tap-vsock`'s `UDPProxy`,
which handles client demultiplexing internally based on the source
address of incoming datagrams. It expects its `dialer` function to
return a new `net.Conn` for each new client it detects.

This commit moves the gRPC stream creation into the `UDPProxy` dialer
function. This ensures a new, dedicated stream is created for each new
client, fixing the incorrect response routing.

Signed-off-by: Viktor Oreshkin <imselfish@stek29.rocks>

Author: stek29

pkg/portfwd/client.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync/atomic"
 	"time"
 
 	"github.com/containers/gvisor-tap-vsock/pkg/services/forwarder"
@@ -40,33 +41,35 @@ func HandleTCPConnection(ctx context.Context, client *guestagentclient.GuestAgen
 }
 
 func HandleUDPConnection(ctx context.Context, client *guestagentclient.GuestAgentClient, conn net.PacketConn, guestAddr string) {
-	id := fmt.Sprintf("udp-%s", conn.LocalAddr().String())
-
-	stream, err := client.Tunnel(ctx)
-	if err != nil {
-		logrus.Errorf("could not open udp tunnel for id: %s error:%v", id, err)
-		return
-	}
-
-	// Handshake message to start tunnel
-	if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: "udp", GuestAddr: guestAddr}); err != nil {
-		logrus.Errorf("could not start udp tunnel for id: %s error:%v", id, err)
-		return
-	}
+	var udpConnectionCounter atomic.Uint32
+	initialID := fmt.Sprintf("udp-%s", conn.LocalAddr().String())
 
+	// 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)
+		if err != nil {
+			return nil, fmt.Errorf("could not open udp tunnel for id: %s error:%w", id, 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)
+		}
 		rw := &GrpcClientRW{stream: stream, id: id, addr: guestAddr, protocol: "udp"}
 		return rw, nil
 	})
 	if err != nil {
-		logrus.Errorf("error in udp tunnel proxy for id: %s error:%v", id, err)
+		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", id, err)
+			logrus.Errorf("error in closing udp tunnel proxy for id: %s error:%v", initialID, err)
 		}
 	}()
 	proxy.Run()