feat: updated gateway to quic

cli/packages/gateway/connection.go

@@ -3,20 +3,49 @@ package gateway
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"net"
+	"strings"
 	"sync"
 
+	"github.com/quic-go/quic-go"
 	"github.com/rs/zerolog/log"
 )
 
-func handleConnection(conn net.Conn) {
-	defer conn.Close()
-	log.Info().Msgf("New connection from: %s", conn.RemoteAddr().String())
+func handleConnection(ctx context.Context, quicConn quic.Connection) {
+	log.Info().Msgf("New connection from: %s", quicConn.RemoteAddr().String())
+	// Use WaitGroup to track all streams
+	var wg sync.WaitGroup
+	for {
+		// Accept the first stream, which we'll use for commands
+		stream, err := quicConn.AcceptStream(ctx)
+		if err != nil {
+			log.Printf("Failed to accept QUIC stream: %v", err)
+			break
+		}
+		wg.Add(1)
+		go func(stream quic.Stream) {
+			defer wg.Done()
+			defer stream.Close()
+
+			handleStream(stream, quicConn)
+		}(stream)
+	}
+
+	wg.Wait()
+	log.Printf("All streams closed for connection: %s", quicConn.RemoteAddr().String())
+}
+
+func handleStream(stream quic.Stream, quicConn quic.Connection) {
+	streamID := stream.StreamID()
+	log.Printf("New stream %d from: %s", streamID, quicConn.RemoteAddr().String())
 
 	// Use buffered reader for better handling of fragmented data
-	reader := bufio.NewReader(conn)
+	reader := bufio.NewReader(stream)
+	defer stream.Close()
+
 	for {
 		msg, err := reader.ReadBytes('\n')
 		if err != nil {
@@ -32,6 +61,7 @@ func handleConnection(conn net.Conn) {
 
 		switch string(cmd) {
 		case "FORWARD-TCP":
+			log.Info().Msg("Starting secure connector proxy...")
 			proxyAddress := string(bytes.Split(args, []byte(" "))[0])
 			destTarget, err := net.Dial("tcp", proxyAddress)
 			if err != nil {
@@ -56,10 +86,10 @@ func handleConnection(conn net.Conn) {
 				}
 			}
 
-			CopyData(conn, destTarget)
+			CopyDataFromQuicToTcp(stream, destTarget)
 			return
 		case "PING":
-			if _, err := conn.Write([]byte("PONG")); err != nil {
+			if _, err := stream.Write([]byte("PONG\n")); err != nil {
 				log.Error().Msgf("Error writing PONG response: %v", err)
 			}
 			return
@@ -74,34 +104,38 @@ type CloseWrite interface {
 	CloseWrite() error
 }
 
-func CopyData(src, dst net.Conn) {
+func CopyDataFromQuicToTcp(quicStream quic.Stream, tcpConn net.Conn) {
+	// Create a WaitGroup to wait for both copy operations
 	var wg sync.WaitGroup
 	wg.Add(2)
 
-	copyAndClose := func(dst, src net.Conn, done chan<- bool) {
+	// Start copying from QUIC stream to TCP
+	go func() {
 		defer wg.Done()
-		_, err := io.Copy(dst, src)
-		if err != nil && !errors.Is(err, io.EOF) {
-			log.Error().Msgf("Copy error: %v", err)
+		if _, err := io.Copy(tcpConn, quicStream); err != nil {
+			log.Error().Msgf("Error copying quic->postgres: %v", err)
 		}
 
-		// Signal we're done writing
-		done <- true
-
-		// Half close the connection if possible
-		if c, ok := dst.(CloseWrite); ok {
-			c.CloseWrite()
+		if e, ok := tcpConn.(CloseWrite); ok {
+			log.Debug().Msg("Closing TCP write end")
+			e.CloseWrite()
+		} else {
+			log.Debug().Msg("TCP connection does not support CloseWrite")
 		}
-	}
+	}()
 
-	done1 := make(chan bool, 1)
-	done2 := make(chan bool, 1)
-
-	go copyAndClose(dst, src, done1)
-	go copyAndClose(src, dst, done2)
+	// Start copying from TCP to QUIC stream
+	go func() {
+		defer wg.Done()
+		if _, err := io.Copy(quicStream, tcpConn); err != nil {
+			log.Debug().Msgf("Error copying postgres->quic: %v", err)
+		}
+		// Close the write side of the QUIC stream
+		if err := quicStream.Close(); err != nil && !strings.Contains(err.Error(), "close called for canceled stream") {
+			log.Error().Msgf("Error closing QUIC stream write: %v", err)
+		}
+	}()
 
 	// Wait for both copies to complete
-	<-done1
-	<-done2
 	wg.Wait()
 }

cli/packages/gateway/gateway.go

@@ -15,6 +15,8 @@ import (
 	"github.com/pion/logging"
 	"github.com/pion/turn/v4"
 	"github.com/rs/zerolog/log"
+
+	"github.com/quic-go/quic-go"
 )
 
 type GatewayConfig struct {
@@ -56,13 +58,12 @@ func (g *Gateway) ConnectWithRelay() error {
 	if relayPort == "5349" {
 		log.Info().Msgf("Provided relay port %s. Using TLS", relayPort)
 		conn, err = tls.Dial("tcp", relayDetails.TurnServerAddress, &tls.Config{
-			InsecureSkipVerify: false,
-			ServerName:         relayAddress,
+			ServerName: relayAddress,
 		})
 	} else {
 		log.Info().Msgf("Provided relay port %s. Using non TLS connection.", relayPort)
-		peerAddr, err := net.ResolveTCPAddr("tcp", relayDetails.TurnServerAddress)
-		if err != nil {
+		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)
@@ -116,20 +117,20 @@ func (g *Gateway) Listen(ctx context.Context) error {
 	// Allocate a relay socket on the TURN server. On success, it
 	// will return a net.PacketConn which represents the remote
 	// socket.
-	relayNonTlsConn, err := g.client.AllocateTCP()
+	relayUdpConnection, err := g.client.Allocate()
 	if err != nil {
 		return fmt.Errorf("Failed to allocate relay connection: %w", err)
 	}
 
-	log.Info().Msg(relayNonTlsConn.Addr().String())
+	log.Info().Msg(relayUdpConnection.LocalAddr().String())
 	defer func() {
-		if closeErr := relayNonTlsConn.Close(); closeErr != nil {
+		if closeErr := relayUdpConnection.Close(); closeErr != nil {
 			log.Error().Msgf("Failed to close connection: %s", closeErr)
 		}
 	}()
 
 	gatewayCert, err := api.CallExchangeRelayCertV1(g.httpClient, api.ExchangeRelayCertRequestV1{
-		RelayAddress: relayNonTlsConn.Addr().String(),
+		RelayAddress: relayUdpConnection.LocalAddr().String(),
 	})
 	if err != nil {
 		return err
@@ -144,45 +145,58 @@ func (g *Gateway) Listen(ctx context.Context) error {
 
 	if g.config.InfisicalStaticIp != "" {
 		log.Info().Msgf("Found static ip from Infisical: %s. Creating permission IP lifecycle", g.config.InfisicalStaticIp)
-		peerAddr, err := net.ResolveTCPAddr("tcp", g.config.InfisicalStaticIp)
+		peerAddr, err := net.ResolveUDPAddr("udp", g.config.InfisicalStaticIp)
 		if err != nil {
 			return fmt.Errorf("Failed to parse infisical static ip: %w", err)
 		}
-		g.registerPermissionLifecycle(func() error {
-			err := relayNonTlsConn.CreatePermissions(peerAddr)
-			return err
-		}, shutdownCh)
+		err = g.client.CreatePermission(peerAddr)
+		if err != nil {
+			return fmt.Errorf("Failed to set permission: %w", err)
+		}
 	}
 
 	cert, err := tls.X509KeyPair([]byte(gatewayCert.Certificate), []byte(gatewayCert.PrivateKey))
 	if err != nil {
-		return fmt.Errorf("failed to parse cert: %s", err)
+		return fmt.Errorf("failed to parse cert: %w", err)
 	}
 
 	caCertPool := x509.NewCertPool()
 	caCertPool.AppendCertsFromPEM([]byte(gatewayCert.CertificateChain))
 
-	relayConn := tls.NewListener(relayNonTlsConn, &tls.Config{
+	// Setup QUIC server
+	tlsConfig := &tls.Config{
 		Certificates: []tls.Certificate{cert},
 		MinVersion:   tls.VersionTLS12,
 		ClientCAs:    caCertPool,
 		ClientAuth:   tls.RequireAndVerifyClientCert,
-	})
+		NextProtos:   []string{"infisical-gateway"},
+	}
+
+	// Setup QUIC listener on the relayConn
+	quicConfig := &quic.Config{
+		EnableDatagrams: true,
+		MaxIdleTimeout:  30 * time.Second,
+		KeepAlivePeriod: 15 * time.Second,
+	}
+
+	quicListener, err := quic.Listen(relayUdpConnection, tlsConfig, quicConfig)
+	if err != nil {
+		return fmt.Errorf("Failed to listen for QUIC: %w", err)
+	}
+	defer quicListener.Close()
+
+	log.Printf("Listener started on %s", quicListener.Addr())
 
 	errCh := make(chan error, 1)
 	log.Info().Msg("Gateway started successfully")
 	g.registerHeartBeat(errCh, shutdownCh)
-	g.registerRelayIsActive(relayNonTlsConn.Addr().String(), errCh, shutdownCh)
+	g.registerRelayIsActive(relayUdpConnection.LocalAddr().String(), tlsConfig, quicConfig, errCh, shutdownCh)
 
 	// Create a WaitGroup to track active connections
 	var wg sync.WaitGroup
 
 	go func() {
 		for {
-			if relayDeadlineConn, ok := relayConn.(*net.TCPListener); ok {
-				relayDeadlineConn.SetDeadline(time.Now().Add(1 * time.Second))
-			}
-
 			select {
 			case <-ctx.Done():
 				return
@@ -190,67 +204,42 @@ func (g *Gateway) Listen(ctx context.Context) error {
 				return
 			default:
 				// Accept new relay connection
-				conn, err := relayConn.Accept()
+				quicConn, err := quicListener.Accept(context.Background())
 				if err != nil {
-					// Check if it's a timeout error (which we expect due to our deadline)
-					if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
-						continue
-					}
-
-					if !strings.Contains(err.Error(), "data contains incomplete STUN or TURN frame") {
-						log.Error().Msgf("Failed to accept connection: %v", err)
-					}
+					log.Printf("Failed to accept QUIC connection: %v", err)
 					continue
 				}
 
-				tlsConn, ok := conn.(*tls.Conn)
-				if !ok {
-					log.Error().Msg("Failed to convert to TLS connection")
-					conn.Close()
-					continue
-				}
-
-				// Set a deadline for the handshake to prevent hanging
-				tlsConn.SetDeadline(time.Now().Add(10 * time.Second))
-				err = tlsConn.Handshake()
-				// Clear the deadline after handshake
-				tlsConn.SetDeadline(time.Time{})
-				if err != nil {
-					log.Error().Msgf("TLS handshake failed: %v", err)
-					conn.Close()
-					continue
-				}
-
-				// Get connection state which contains certificate information
-				state := tlsConn.ConnectionState()
-				if len(state.PeerCertificates) > 0 {
-					organizationUnit := state.PeerCertificates[0].Subject.OrganizationalUnit
-					commonName := state.PeerCertificates[0].Subject.CommonName
+				tlsState := quicConn.ConnectionState().TLS
+				if len(tlsState.PeerCertificates) > 0 {
+					organizationUnit := tlsState.PeerCertificates[0].Subject.OrganizationalUnit
+					commonName := tlsState.PeerCertificates[0].Subject.CommonName
 					if organizationUnit[0] != "gateway-client" || commonName != "cloud" {
-						log.Error().Msgf("Client certificate verification failed. Received %s, %s", organizationUnit, commonName)
-						conn.Close()
+						errMsg := fmt.Sprintf("Client certificate verification failed. Received %s, %s", organizationUnit, commonName)
+						log.Error().Msg(errMsg)
+						quicConn.CloseWithError(1, errMsg)
 						continue
 					}
 				}
 
 				// Handle the connection in a goroutine
 				wg.Add(1)
-				go func(c net.Conn) {
+				go func(c quic.Connection) {
 					defer wg.Done()
-					defer c.Close()
+					defer c.CloseWithError(0, "connection closed")
 
 					// Monitor parent context to close this connection when needed
 					go func() {
 						select {
 						case <-ctx.Done():
-							c.Close() // Force close connection when context is canceled
+							c.CloseWithError(0, "connection closed") // Force close connection when context is canceled
 						case <-shutdownCh:
-							c.Close() // Force close connection when accepting loop is done
+							c.CloseWithError(0, "connection closed") // Force close connection when accepting loop is done
 						}
 					}()
 
-					handleConnection(c)
-				}(conn)
+					handleConnection(ctx, c)
+				}(quicConn)
 			}
 		}
 	}()
@@ -282,7 +271,7 @@ func (g *Gateway) Listen(ctx context.Context) error {
 }
 
 func (g *Gateway) registerHeartBeat(errCh chan error, done chan bool) {
-	ticker := time.NewTicker(1 * time.Hour)
+	ticker := time.NewTicker(30 * time.Minute)
 
 	go func() {
 		time.Sleep(10 * time.Second)
@@ -306,26 +295,7 @@ func (g *Gateway) registerHeartBeat(errCh chan error, done chan bool) {
 	}()
 }
 
-func (g *Gateway) registerPermissionLifecycle(permissionFn func() error, done chan bool) {
-	ticker := time.NewTicker(3 * time.Minute)
-
-	go func() {
-		// wait for 5 mins
-		permissionFn()
-		log.Printf("Created permission for incoming connections")
-		for {
-			select {
-			case <-done:
-				ticker.Stop()
-				return
-			case <-ticker.C:
-				permissionFn()
-			}
-		}
-	}()
-}
-
-func (g *Gateway) registerRelayIsActive(serverAddr string, errCh chan error, done chan bool) {
+func (g *Gateway) registerRelayIsActive(serverAddr string, tlsConf *tls.Config, quicConf *quic.Config, errCh chan error, done chan bool) {
 	ticker := time.NewTicker(10 * time.Second)
 
 	go func() {
@@ -336,14 +306,19 @@ func (g *Gateway) registerRelayIsActive(serverAddr string, errCh chan error, don
 				ticker.Stop()
 				return
 			case <-ticker.C:
-				conn, err := net.Dial("tcp", serverAddr)
-				if err != nil {
-					errCh <- err
-					return
-				}
-				if conn != nil {
-					conn.Close()
-				}
+				func() {
+					ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) // 3s handshake timeout
+					defer cancel()
+					conn, err := quic.DialAddr(ctx, serverAddr, tlsConf, quicConf)
+					if conn != nil {
+						conn.CloseWithError(0, "connection closed")
+					}
+					// this error means quic connection is alive
+					if err != nil && !strings.Contains(err.Error(), "tls: failed to verify certificate") {
+						errCh <- err
+						return
+					}
+				}()
 			}
 		}
 	}()

cli/packages/gateway/relay.go

@@ -0,0 +1,186 @@
+package gateway
+
+import (
+	"context"
+	"crypto/tls"
+	"errors"
+	"fmt"
+	"net"
+	"os"
+	"os/signal"
+	"runtime"
+	"strconv"
+	"syscall"
+
+	"github.com/pion/logging"
+	"github.com/pion/turn/v4"
+	"github.com/rs/zerolog/log"
+	"golang.org/x/sys/unix"
+	"gopkg.in/yaml.v2"
+)
+
+var (
+	errMissingTlsCert = errors.New("Missing TLS files")
+)
+
+type GatewayRelay struct {
+	Config *GatewayRelayConfig
+}
+
+type GatewayRelayConfig struct {
+	PublicIP          string `json:"public_ip"`
+	Port              int    `json:"port"`
+	Realm             string `json:"realm"`
+	AuthSecret        string `json:"auth_secret"`
+	RelayMinPort      uint16 `json:"relay_min_port"`
+	RelayMaxPort      uint16 `json:"relay_max_port"`
+	TlsCertPath       string `json:"tls_cert_path"`
+	TlsPrivateKeyPath string `json:"tls_private_key_path"`
+
+	tls          tls.Certificate
+	isTlsEnabled bool
+}
+
+func NewGatewayRelay(configFilePath string) (*GatewayRelay, error) {
+	cfgFile, err := os.ReadFile(configFilePath)
+	if err != nil {
+		return nil, err
+	}
+	var cfg GatewayRelayConfig
+	if err := yaml.Unmarshal(cfgFile, &cfg); err != nil {
+		return nil, err
+	}
+
+	if cfg.PublicIP == "" {
+		return nil, fmt.Errorf("Missing public ip")
+	}
+
+	if cfg.AuthSecret == "" {
+		return nil, fmt.Errorf("Missing auth secret")
+	}
+
+	if cfg.Realm == "" {
+		cfg.Realm = "infisical.org"
+	}
+
+	if cfg.RelayMinPort == 0 {
+		cfg.RelayMinPort = 49152
+	}
+
+	if cfg.RelayMaxPort == 0 {
+		cfg.RelayMaxPort = 65535
+	}
+
+	if cfg.Port == 0 {
+		cfg.Port = 3478
+	} else if cfg.Port == 5349 {
+		if cfg.TlsCertPath == "" || cfg.TlsPrivateKeyPath == "" {
+			return nil, errMissingTlsCert
+		}
+
+		tlsCertFile, err := os.ReadFile(cfg.TlsCertPath)
+		if err != nil {
+			return nil, err
+		}
+		tlsPrivateKeyFile, err := os.ReadFile(cfg.TlsPrivateKeyPath)
+		if err != nil {
+			return nil, err
+		}
+
+		cert, err := tls.LoadX509KeyPair(string(tlsCertFile), string(tlsPrivateKeyFile))
+		if err != nil {
+			return nil, err
+		}
+		cfg.tls = cert
+		cfg.isTlsEnabled = true
+	}
+
+	return &GatewayRelay{
+		Config: &cfg,
+	}, nil
+}
+
+func (g *GatewayRelay) Run() error {
+	addr, err := net.ResolveTCPAddr("tcp", "0.0.0.0:"+strconv.Itoa(g.Config.Port))
+	if err != nil {
+		return fmt.Errorf("Failed to parse server address: %s", err)
+	}
+
+	// NewLongTermAuthHandler takes a pion.LeveledLogger. This allows you to intercept messages
+	// 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 := &net.ListenConfig{
+		Control: func(network, address string, conn syscall.RawConn) error { // nolint: revive
+			var operr error
+			if err = conn.Control(func(fd uintptr) {
+				operr = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1)
+			}); err != nil {
+				return err
+			}
+
+			return operr
+		},
+	}
+
+	publicIP := g.Config.PublicIP
+	relayAddressGenerator := &turn.RelayAddressGeneratorPortRange{
+		RelayAddress: net.ParseIP(publicIP), // Claim that we are listening on IP passed by user
+		Address:      "0.0.0.0",             // But actually be listening on every interface
+		MinPort:      g.Config.RelayMinPort,
+		MaxPort:      g.Config.RelayMaxPort,
+	}
+
+	threadNum := runtime.NumCPU()
+	listenerConfigs := make([]turn.ListenerConfig, threadNum)
+	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 {
+			listenerConfigs[i].Listener = tls.NewListener(conn, &tls.Config{
+				Certificates: []tls.Certificate{g.Config.tls},
+			})
+		} else {
+			listenerConfigs[i].Listener = conn
+		}
+
+		log.Printf("Server %d listening on %s\n", i, conn.Addr().String())
+	}
+
+	loggerF := logging.NewDefaultLoggerFactory()
+	loggerF.DefaultLogLevel = logging.LogLevelDebug
+
+	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,
+	})
+
+	if err != nil {
+		return fmt.Errorf("Failed to start server: %w", err)
+	}
+
+	// Block until user sends SIGINT or SIGTERM
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	<-sigs
+
+	if err = server.Close(); err != nil {
+		return fmt.Errorf("Failed to close server: %w", err)
+	}
+	return nil
+}