feat: updated gateway to use dtls

cli/packages/gateway/gateway.go

@@ -14,6 +14,7 @@ import (
 	"github.com/Infisical/infisical-merge/packages/api"
 	"github.com/Infisical/infisical-merge/packages/systemd"
 	"github.com/go-resty/resty/v2"
+	"github.com/pion/dtls/v3"
 	"github.com/pion/logging"
 	"github.com/pion/turn/v4"
 	"github.com/rs/zerolog/log"
@@ -54,26 +55,6 @@ func (g *Gateway) ConnectWithRelay() error {
 		return err
 	}
 	relayAddress, relayPort := strings.Split(relayDetails.TurnServerAddress, ":")[0], strings.Split(relayDetails.TurnServerAddress, ":")[1]
-	var conn net.Conn
-
-	// Dial TURN Server
-	if relayPort == "5349" {
-		log.Info().Msgf("Provided relay port %s. Using TLS", relayPort)
-		conn, err = tls.Dial("tcp", relayDetails.TurnServerAddress, &tls.Config{
-			ServerName: relayAddress,
-		})
-	} else {
-		log.Info().Msgf("Provided relay port %s. Using non TLS connection.", relayPort)
-		peerAddr, errPeer := net.ResolveTCPAddr("tcp", relayDetails.TurnServerAddress)
-		if errPeer != nil {
-			return fmt.Errorf("Failed to parse turn server address: %w", err)
-		}
-		conn, err = net.DialTCP("tcp", nil, peerAddr)
-	}
-
-	if err != nil {
-		return fmt.Errorf("Failed to connect with relay server: %w", err)
-	}
 
 	// Start a new TURN Client and wrap our net.Conn in a STUNConn
 	// This allows us to simulate datagram based communication over a net.Conn
@@ -81,17 +62,48 @@ func (g *Gateway) ConnectWithRelay() error {
 	if os.Getenv("LOG_LEVEL") == "debug" {
 		logger.DefaultLogLevel = logging.LogLevelDebug
 	}
-	cfg := &turn.ClientConfig{
+
+	turnClientCfg := &turn.ClientConfig{
 		STUNServerAddr: relayDetails.TurnServerAddress,
 		TURNServerAddr: relayDetails.TurnServerAddress,
-		Conn:           turn.NewSTUNConn(conn),
 		Username:       relayDetails.TurnServerUsername,
 		Password:       relayDetails.TurnServerPassword,
 		Realm:          relayDetails.TurnServerRealm,
 		LoggerFactory:  logger,
 	}
 
-	client, err := turn.NewClient(cfg)
+	turnAddr, err := net.ResolveUDPAddr("udp4", relayDetails.TurnServerAddress)
+	if err != nil {
+		return fmt.Errorf("Failed to parse turn server address: %w", err)
+	}
+
+	// Dial TURN Server
+	if relayPort == "5349" {
+		caCertPool := x509.NewCertPool()
+		a, _ := os.ReadFile("/Users/akhilmhdh/Akhi/infisical-gateway-poc/certificates/test/ca-cert.pem")
+		// caCertPool.AppendCertsFromPEM([]byte(g.config.CertificateChain))
+		caCertPool.AppendCertsFromPEM(a)
+
+		log.Info().Msgf("Provided relay port %s. Using TLS", relayPort)
+		conn, err := dtls.Dial("udp", turnAddr, &dtls.Config{
+			ServerName: relayAddress,
+			RootCAs:    caCertPool,
+		})
+		if err != nil {
+			return fmt.Errorf("Failed to connect with relay server: %w", err)
+		}
+		turnClientCfg.Conn = turn.NewSTUNConn(conn)
+	} else {
+		log.Info().Msgf("Provided relay port %s. Using non TLS connection.", relayPort)
+		conn, err := net.ListenPacket("udp4", turnAddr.String())
+		if err != nil {
+			return fmt.Errorf("Failed to connect with relay server: %w", err)
+		}
+
+		turnClientCfg.Conn = conn
+	}
+
+	client, err := turn.NewClient(turnClientCfg)
 	if err != nil {
 		return fmt.Errorf("Failed to create relay client: %w", err)
 	}
@@ -168,7 +180,6 @@ func (g *Gateway) Listen(ctx context.Context) error {
 		ClientAuth:   tls.RequireAndVerifyClientCert,
 		NextProtos:   []string{"infisical-gateway"},
 	}
-
 	// Setup QUIC listener on the relayConn
 	quicConfig := &quic.Config{
 		EnableDatagrams: true,
@@ -176,7 +187,6 @@ func (g *Gateway) Listen(ctx context.Context) error {
 		KeepAlivePeriod: 2 * time.Second,
 	}
 
-	g.registerRelayIsActive(ctx, errCh)
 	quicListener, err := quic.Listen(relayUdpConnection, tlsConfig, quicConfig)
 	if err != nil {
 		return fmt.Errorf("Failed to listen for QUIC: %w", err)
@@ -185,6 +195,8 @@ func (g *Gateway) Listen(ctx context.Context) error {
 
 	log.Printf("Listener started on %s", quicListener.Addr())
 
+	g.registerRelayIsActive(ctx, quicListener.Addr().String(), errCh)
+
 	log.Info().Msg("Gateway started successfully")
 
 	var wg sync.WaitGroup
@@ -320,13 +332,12 @@ func (g *Gateway) createPermissionForStaticIps(staticIps string) error {
 	return nil
 }
 
-func (g *Gateway) registerRelayIsActive(ctx context.Context, errCh chan error) error {
+func (g *Gateway) registerRelayIsActive(ctx context.Context, addr string, errCh chan error) error {
 	ticker := time.NewTicker(15 * time.Second)
 	maxFailures := 3
 	failures := 0
 
 	log.Info().Msg("Starting relay connection health check")
-
 	go func() {
 		time.Sleep(5 * time.Second)
 		for {
@@ -335,36 +346,26 @@ func (g *Gateway) registerRelayIsActive(ctx context.Context, errCh chan error) e
 				log.Info().Msg("Stopping relay connection health check")
 				return
 			case <-ticker.C:
-				func() {
-					log.Debug().Msg("Performing relay connection health check")
-
-					if g.client == nil {
-						failures++
-						log.Warn().Int("failures", failures).Msg("TURN client is nil")
-						if failures >= maxFailures {
-							errCh <- fmt.Errorf("relay connection check failed: TURN client is nil")
-						}
+				log.Debug().Msg("Performing relay connection health check")
+				ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+				defer cancel()
+				// Try to establish a QUIC connection
+				conn, err := quic.DialAddr(ctx, addr, &tls.Config{
+					InsecureSkipVerify: false, // Skip certificate verification
+					NextProtos:         []string{"infisical-gateway"},
+				}, nil)
+				if err != nil && !strings.Contains(err.Error(), "tls:") {
+					failures++
+					log.Warn().Err(err).Int("failures", failures).Msg("Failed to refresh TURN permissions")
+					if failures >= maxFailures {
+						errCh <- fmt.Errorf("relay connection check failed: %w", err)
 						return
 					}
-
-					// we try to refresh permissions - this is a lightweight operation
-					// that will fail immediately if the UDP connection is broken. good for health check
-					log.Debug().Msg("Refreshing TURN permissions to verify connection")
-					if err := g.createPermissionForStaticIps(g.config.InfisicalStaticIp); err != nil {
-						failures++
-						log.Warn().Err(err).Int("failures", failures).Msg("Failed to refresh TURN permissions")
-						if failures >= maxFailures {
-							errCh <- fmt.Errorf("relay connection check failed: %w", err)
-						}
-						return
-					}
-
-					log.Debug().Msg("Successfully refreshed TURN permissions - connection is healthy")
-					if failures > 0 {
-						log.Info().Int("previous_failures", failures).Msg("Relay connection restored")
-						failures = 0
-					}
-				}()
+					continue
+				}
+				if conn != nil {
+					defer conn.CloseWithError(0, "All good")
+				}
 			}
 		}
 	}()

cli/packages/gateway/relay.go

@@ -4,7 +4,6 @@
 package gateway
 
 import (
-	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"errors"
@@ -12,12 +11,13 @@ import (
 	"net"
 	"os"
 	"os/signal"
-	"runtime"
+
+	// "runtime"
 	"strconv"
 	"syscall"
 
-	udplistener "github.com/Infisical/infisical-merge/packages/gateway/udp_listener"
 	"github.com/Infisical/infisical-merge/packages/systemd"
+	"github.com/pion/dtls/v3"
 	"github.com/pion/logging"
 	"github.com/pion/turn/v4"
 	"github.com/rs/zerolog/log"
@@ -108,7 +108,7 @@ func NewGatewayRelay(configFilePath string) (*GatewayRelay, error) {
 }
 
 func (g *GatewayRelay) Run() error {
-	addr, err := net.ResolveTCPAddr("tcp", "0.0.0.0:"+strconv.Itoa(g.Config.Port))
+	addr, err := net.ResolveUDPAddr("udp", "0.0.0.0:"+strconv.Itoa(g.Config.Port))
 	if err != nil {
 		return fmt.Errorf("Failed to parse server address: %s", err)
 	}
@@ -117,13 +117,6 @@ func (g *GatewayRelay) Run() error {
 	// and process them yourself.
 	logger := logging.NewDefaultLeveledLoggerForScope("lt-creds", logging.LogLevelTrace, os.Stdout)
 
-	// Create `numThreads` UDP listeners to pass into pion/turn
-	// pion/turn itself doesn't allocate any UDP sockets, but lets the user pass them in
-	// this allows us to add logging, storage or modify inbound/outbound traffic
-	// UDP listeners share the same local address:port with setting SO_REUSEPORT and the kernel
-	// will load-balance received packets per the IP 5-tuple
-	listenerConfig := udplistener.SetupListenerConfig()
-
 	publicIP := g.Config.PublicIP
 	relayAddressGenerator := &turn.RelayAddressGeneratorPortRange{
 		RelayAddress: net.ParseIP(publicIP), // Claim that we are listening on IP passed by user
@@ -132,49 +125,54 @@ func (g *GatewayRelay) Run() error {
 		MaxPort:      g.Config.RelayMaxPort,
 	}
 
-	threadNum := runtime.NumCPU()
-	listenerConfigs := make([]turn.ListenerConfig, threadNum)
-	var connAddress string
-	for i := 0; i < threadNum; i++ {
-		conn, listErr := listenerConfig.Listen(context.Background(), addr.Network(), addr.String())
-		if listErr != nil {
-			return fmt.Errorf("Failed to allocate TCP listener at %s:%s %s", addr.Network(), addr.String(), listErr)
-		}
-
-		listenerConfigs[i] = turn.ListenerConfig{
-			RelayAddressGenerator: relayAddressGenerator,
-		}
-
-		if g.Config.isTlsEnabled {
-			caCertPool := x509.NewCertPool()
-			caCertPool.AppendCertsFromPEM([]byte(g.Config.tlsCa))
-
-			listenerConfigs[i].Listener = tls.NewListener(conn, &tls.Config{
-				Certificates: []tls.Certificate{g.Config.tls},
-				ClientCAs:    caCertPool,
-			})
-		} else {
-			listenerConfigs[i].Listener = conn
-		}
-		connAddress = conn.Addr().String()
-	}
-
 	loggerF := logging.NewDefaultLoggerFactory()
 	loggerF.DefaultLogLevel = logging.LogLevelDebug
 
+	caCertPool := x509.NewCertPool()
+	caCertPool.AppendCertsFromPEM([]byte(g.Config.tlsCa))
+
+	listenerConfigs := make([]turn.ListenerConfig, 0)
+	packetConfigs := make([]turn.PacketConnConfig, 0)
+
+	if g.Config.isTlsEnabled {
+		caCertPool := x509.NewCertPool()
+		caCertPool.AppendCertsFromPEM([]byte(g.Config.tlsCa))
+		dtlsServer, err := dtls.Listen("udp", addr, &dtls.Config{
+			Certificates: []tls.Certificate{g.Config.tls},
+			ClientCAs:    caCertPool,
+		})
+		if err != nil {
+			return fmt.Errorf("Failed to start dtls server: %w", err)
+		}
+		listenerConfigs = append(listenerConfigs, turn.ListenerConfig{
+			RelayAddressGenerator: relayAddressGenerator,
+			Listener:              dtlsServer,
+		})
+	} else {
+		udpListener, err := net.ListenPacket("udp4", "0.0.0.0:"+strconv.Itoa(g.Config.Port))
+		if err != nil {
+			return fmt.Errorf("Failed to relay udp listener: %w", err)
+		}
+		packetConfigs = append(packetConfigs, turn.PacketConnConfig{
+			RelayAddressGenerator: relayAddressGenerator,
+			PacketConn:            udpListener,
+		})
+	}
+
 	server, err := turn.NewServer(turn.ServerConfig{
 		Realm:       g.Config.Realm,
 		AuthHandler: turn.LongTermTURNRESTAuthHandler(g.Config.AuthSecret, logger),
 		// PacketConnConfigs is a list of UDP Listeners and the configuration around them
-		ListenerConfigs: listenerConfigs,
-		LoggerFactory:   loggerF,
+		ListenerConfigs:   listenerConfigs,
+		PacketConnConfigs: packetConfigs,
+		LoggerFactory:     loggerF,
 	})
 
 	if err != nil {
 		return fmt.Errorf("Failed to start server: %w", err)
 	}
 
-	log.Info().Msgf("Relay listening on %s\n", connAddress)
+	log.Info().Msgf("Relay listening on %d\n", g.Config.Port)
 
 	// make this compatiable with systemd notify mode
 	systemd.SdNotify(false, systemd.SdNotifyReady)