|
7 | 7 | "context"
|
8 | 8 | "fmt"
|
9 | 9 | "net"
|
| 10 | + "sync/atomic" |
10 | 11 | "time"
|
11 | 12 |
|
12 | 13 | "github.com/containers/gvisor-tap-vsock/pkg/services/forwarder"
|
@@ -40,33 +41,35 @@ func HandleTCPConnection(ctx context.Context, client *guestagentclient.GuestAgen
|
40 | 41 | }
|
41 | 42 |
|
42 | 43 | func HandleUDPConnection(ctx context.Context, client *guestagentclient.GuestAgentClient, conn net.PacketConn, guestAddr string) {
|
43 |
| - id := fmt.Sprintf("udp-%s", conn.LocalAddr().String()) |
44 |
| - |
45 |
| - stream, err := client.Tunnel(ctx) |
46 |
| - if err != nil { |
47 |
| - logrus.Errorf("could not open udp tunnel for id: %s error:%v", id, err) |
48 |
| - return |
49 |
| - } |
50 |
| - |
51 |
| - // Handshake message to start tunnel |
52 |
| - if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: "udp", GuestAddr: guestAddr}); err != nil { |
53 |
| - logrus.Errorf("could not start udp tunnel for id: %s error:%v", id, err) |
54 |
| - return |
55 |
| - } |
| 44 | + var udpConnectionCounter uint64 |
| 45 | + initialID := fmt.Sprintf("udp-%s", conn.LocalAddr().String()) |
56 | 46 |
|
| 47 | + // gvisor-tap-vsock's UDPProxy demultiplexes client connections internally based on their source address. |
| 48 | + // It calls this dialer function only when it receives a datagram from a new, unrecognized client. |
| 49 | + // For each new client, we must return a new net.Conn, which in our case is a new gRPC stream. |
| 50 | + // The atomic counter ensures that each stream has a unique ID to distinguish them on the server side. |
57 | 51 | proxy, err := forwarder.NewUDPProxy(conn, func() (net.Conn, error) {
|
| 52 | + id := fmt.Sprintf("%s-%d", initialID, atomic.AddUint64(&udpConnectionCounter, 1)) |
| 53 | + stream, err := client.Tunnel(ctx) |
| 54 | + if err != nil { |
| 55 | + return nil, fmt.Errorf("could not open udp tunnel for id: %s error:%w", id, err) |
| 56 | + } |
| 57 | + // Handshake message to start tunnel |
| 58 | + if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: "udp", GuestAddr: guestAddr}); err != nil { |
| 59 | + return nil, fmt.Errorf("could not start udp tunnel for id: %s error:%w", id, err) |
| 60 | + } |
58 | 61 | rw := &GrpcClientRW{stream: stream, id: id, addr: guestAddr, protocol: "udp"}
|
59 | 62 | return rw, nil
|
60 | 63 | })
|
61 | 64 | if err != nil {
|
62 |
| - logrus.Errorf("error in udp tunnel proxy for id: %s error:%v", id, err) |
| 65 | + logrus.Errorf("error in udp tunnel proxy for id: %s error:%v", initialID, err) |
63 | 66 | return
|
64 | 67 | }
|
65 | 68 |
|
66 | 69 | defer func() {
|
67 | 70 | err := proxy.Close()
|
68 | 71 | if err != nil {
|
69 |
| - logrus.Errorf("error in closing udp tunnel proxy for id: %s error:%v", id, err) |
| 72 | + logrus.Errorf("error in closing udp tunnel proxy for id: %s error:%v", initialID, err) |
70 | 73 | }
|
71 | 74 | }()
|
72 | 75 | proxy.Run()
|
|
0 commit comments