add udp support for desktop client

api/file_provider.go

@@ -24,5 +24,5 @@ func (f FileConfigProvider) Load() ([]byte, error) {
 
 // Save stores data to the file
 func (f FileConfigProvider) Save(data []byte) error {
-	return os.WriteFile(string(f), data, 0600)
+	return os.WriteFile(string(f), data, 0o600)
 }

api/listener_server.go

@@ -11,6 +11,8 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/rs/zerolog/log"
+
 	pb "github.com/pomerium/cli/proto"
 )
 
@@ -70,6 +72,14 @@ func (s *server) connectTunnelLocked(id string) (net.Addr, error) {
 		return nil, err
 	}
 
+	if rec.GetConn().GetProtocol() == pb.Protocol_UDP {
+		return s.connectUDPTunnelLocked(id, tun, listenAddr)
+	}
+
+	return s.connectTCPTunnelLocked(id, tun, listenAddr)
+}
+
+func (s *server) connectTCPTunnelLocked(id string, tun Tunnel, listenAddr string) (net.Addr, error) {
 	ctx, cancel := context.WithCancel(context.Background())
 	lc := new(net.ListenConfig)
 	li, err := lc.Listen(ctx, "tcp", listenAddr)
@@ -98,6 +108,47 @@ func (s *server) connectTunnelLocked(id string) (net.Addr, error) {
 	return li.Addr(), nil
 }
 
+func (s *server) connectUDPTunnelLocked(id string, tun Tunnel, listenAddr string) (net.Addr, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	addr, err := net.ResolveUDPAddr("udp", listenAddr)
+	if err != nil {
+		_ = s.EventBroadcaster.Update(ctx, &pb.ConnectionStatusUpdate{
+			Id:        id,
+			LastError: proto.String(fmt.Errorf("ResolveUDPAddr: %w", err).Error()),
+			Ts:        timestamppb.Now(),
+		})
+		cancel()
+		return nil, err
+	}
+
+	conn, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		_ = s.EventBroadcaster.Update(ctx, &pb.ConnectionStatusUpdate{
+			Id:        id,
+			LastError: proto.String(fmt.Errorf("ListenUDP: %w", err).Error()),
+			Ts:        timestamppb.Now(),
+		})
+		cancel()
+		return nil, err
+	}
+	context.AfterFunc(ctx, func() { _ = conn.Close() })
+
+	go func() {
+		defer cancel()
+		evt := &tunnelEvents{EventBroadcaster: s.EventBroadcaster, id: id}
+		defer evt.onTunnelClosed()
+		evt.onListening(ctx)
+
+		err := tun.RunUDPSessionManager(ctx, conn, evt)
+		if err != nil {
+			log.Ctx(ctx).Error().Err(err).Msg("error serving local connection")
+		}
+	}()
+
+	return addr, nil
+}
+
 func onContextCancel(ctx context.Context, cl io.Closer) {
 	<-ctx.Done()
 	_ = cl.Close()

api/server.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net"
 	"sync"
 
 	"github.com/golang/groupcache/lru"
@@ -38,6 +39,7 @@ type ListenerStatus interface {
 // Tunnel is abstraction over tunnel.Tunnel to allow mocking
 type Tunnel interface {
 	Run(context.Context, io.ReadWriter, tunnel.EventSink) error
+	RunUDPSessionManager(ctx context.Context, conn *net.UDPConn, eventSink tunnel.EventSink) error
 }
 
 // Server implements both config and listener interfaces