@@ -11,68 +11,57 @@ import (
1111 "time"
1212
1313 "github.com/containers/gvisor-tap-vsock/pkg/services/forwarder"
14- "github.com/containers/gvisor-tap-vsock/pkg /tcpproxy"
14+ "github.com/inetaf /tcpproxy"
1515 "github.com/sirupsen/logrus"
1616
1717 "github.com/lima-vm/lima/v2/pkg/guestagent/api"
1818 guestagentclient "github.com/lima-vm/lima/v2/pkg/guestagent/api/client"
1919)
2020
21- func HandleTCPConnection (ctx context.Context , client * guestagentclient.GuestAgentClient , conn net.Conn , guestAddr string ) {
22- id := fmt .Sprintf ("tcp-%s-%s" , conn .LocalAddr ().String (), conn .RemoteAddr ().String ())
21+ func HandleTCPConnection (_ context.Context , dialContext func (ctx context.Context , network string , addr string ) (net.Conn , error ), conn net.Conn , guestAddr string ) {
22+ proxy := tcpproxy.DialProxy {Addr : guestAddr , DialContext : dialContext }
23+ proxy .HandleConn (conn )
24+ }
2325
24- stream , err := client .Tunnel (ctx )
26+ func HandleUDPConnection (ctx context.Context , dialContext func (ctx context.Context , network string , addr string ) (net.Conn , error ), conn net.PacketConn , guestAddr string ) {
27+ proxy , err := forwarder .NewUDPProxy (conn , func () (net.Conn , error ) {
28+ return dialContext (ctx , "udp" , guestAddr )
29+ })
2530 if err != nil {
26- logrus .Errorf ("could not open tcp tunnel for id: %s error:%v" , id , err )
27- return
28- }
29-
30- // Handshake message to start tunnel
31- if err := stream .Send (& api.TunnelMessage {Id : id , Protocol : "tcp" , GuestAddr : guestAddr }); err != nil {
32- logrus .Errorf ("could not start tcp tunnel for id: %s error:%v" , id , err )
31+ logrus .WithError (err ).Error ("error in udp tunnel proxy" )
3332 return
3433 }
3534
36- rw := & GrpcClientRW {stream : stream , id : id , addr : guestAddr , protocol : "tcp" }
37- proxy := tcpproxy.DialProxy {DialContext : func (_ context.Context , _ , _ string ) (net.Conn , error ) {
38- return conn , nil
39- }}
40- proxy .HandleConn (rw )
35+ defer func () {
36+ err := proxy .Close ()
37+ if err != nil {
38+ logrus .WithError (err ).Error ("error in closing udp tunnel proxy" )
39+ }
40+ }()
41+ proxy .Run ()
4142}
4243
43- func HandleUDPConnection (ctx context.Context , client * guestagentclient.GuestAgentClient , conn net.PacketConn , guestAddr string ) {
44- var udpConnectionCounter atomic.Uint32
45- initialID := fmt .Sprintf ("udp-%s" , conn .LocalAddr ().String ())
46-
44+ func DialContextToGRPCTunnel (client * guestagentclient.GuestAgentClient ) func (ctx context.Context , network , addr string ) (net.Conn , error ) {
4745 // gvisor-tap-vsock's UDPProxy demultiplexes client connections internally based on their source address.
4846 // It calls this dialer function only when it receives a datagram from a new, unrecognized client.
4947 // For each new client, we must return a new net.Conn, which in our case is a new gRPC stream.
5048 // The atomic counter ensures that each stream has a unique ID to distinguish them on the server side.
51- proxy , err := forwarder .NewUDPProxy (conn , func () (net.Conn , error ) {
52- id := fmt .Sprintf ("%s-%d" , initialID , udpConnectionCounter .Add (1 ))
53- stream , err := client .Tunnel (ctx )
49+ var connectionCounter atomic.Uint32
50+ return func (_ context.Context , network , addr string ) (net.Conn , error ) {
51+ // Passed context.Context is used for timeout on initiate connection, not for the lifetime of the connection.
52+ // We use context.Background() here to avoid unexpected cancellation.
53+ stream , err := client .Tunnel (context .Background ())
5454 if err != nil {
55- return nil , fmt .Errorf ("could not open udp tunnel for id : %s error:%w" , id , err )
55+ return nil , fmt .Errorf ("could not open tunnel for addr : %s error:%w" , addr , err )
5656 }
5757 // 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 )
58+ id := fmt .Sprintf ("%s-%s-%d" , network , addr , connectionCounter .Add (1 ))
59+ if err := stream .Send (& api.TunnelMessage {Id : id , Protocol : network , GuestAddr : addr }); err != nil {
60+ return nil , fmt .Errorf ("could not start tunnel for id: %s addr: %s error:%w" , id , addr , err )
6061 }
61- rw := & GrpcClientRW {stream : stream , id : id , addr : guestAddr , protocol : "udp" }
62+ rw := & GrpcClientRW {stream : stream , id : id , addr : addr , protocol : network }
6263 return rw , nil
63- })
64- if err != nil {
65- logrus .Errorf ("error in udp tunnel proxy for id: %s error:%v" , initialID , err )
66- return
6764 }
68-
69- defer func () {
70- err := proxy .Close ()
71- if err != nil {
72- logrus .Errorf ("error in closing udp tunnel proxy for id: %s error:%v" , initialID , err )
73- }
74- }()
75- proxy .Run ()
7665}
7766
7867type GrpcClientRW struct {
0 commit comments