Merge pull request #3183 from akhilmhdh/feat/connector

feat: fixed cli issues in gateway

Author: maidul98

cli/packages/gateway/connection.go

@@ -61,14 +61,14 @@ func handleStream(stream quic.Stream, quicConn quic.Connection) {
 
 		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 {
 				log.Error().Msgf("Failed to connect to target: %v", err)
 				return
 			}
 			defer destTarget.Close()
+			log.Info().Msgf("Starting secure transmission between %s->%s", quicConn.LocalAddr().String(), destTarget.LocalAddr().String())
 
 			// Handle buffered data
 			buffered := reader.Buffered()
@@ -87,6 +87,7 @@ func handleStream(stream quic.Stream, quicConn quic.Connection) {
 			}
 
 			CopyDataFromQuicToTcp(stream, destTarget)
+			log.Info().Msgf("Ending secure transmission between %s->%s", quicConn.LocalAddr().String(), destTarget.LocalAddr().String())
 			return
 		case "PING":
 			if _, err := stream.Write([]byte("PONG\n")); err != nil {

cli/packages/gateway/gateway.go

@@ -6,11 +6,13 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net"
+	"os"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/api"
+	"github.com/Infisical/infisical-merge/packages/systemd"
 	"github.com/go-resty/resty/v2"
 	"github.com/pion/logging"
 	"github.com/pion/turn/v4"
@@ -75,14 +77,18 @@ func (g *Gateway) ConnectWithRelay() error {
 
 	// 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
+	logger := logging.NewDefaultLoggerFactory()
+	if os.Getenv("LOG_LEVEL") == "debug" {
+		logger.DefaultLogLevel = logging.LogLevelDebug
+	}
 	cfg := &turn.ClientConfig{
 		STUNServerAddr: relayDetails.TurnServerAddress,
 		TURNServerAddr: relayDetails.TurnServerAddress,
 		Conn:           turn.NewSTUNConn(conn),
 		Username:       relayDetails.TurnServerUsername,
 		Password:       relayDetails.TurnServerPassword,
 		Realm:          relayDetails.TurnServerRealm,
-		LoggerFactory:  logging.NewDefaultLoggerFactory(),
+		LoggerFactory:  logger,
 	}
 
 	client, err := turn.NewClient(cfg)
@@ -96,10 +102,6 @@ func (g *Gateway) ConnectWithRelay() error {
 		TurnServerAddress:  relayDetails.TurnServerAddress,
 		InfisicalStaticIp:  relayDetails.InfisicalStaticIp,
 	}
-	// if port not specific allow all port
-	if relayDetails.InfisicalStaticIp != "" && !strings.Contains(relayDetails.InfisicalStaticIp, ":") {
-		g.config.InfisicalStaticIp = g.config.InfisicalStaticIp + ":0"
-	}
 
 	g.client = client
 	return nil
@@ -144,7 +146,10 @@ func (g *Gateway) Listen(ctx context.Context) error {
 	errCh := make(chan error, 1)
 	shutdownCh := make(chan bool, 1)
 
-	g.registerPermissionRefresh(ctx, errCh)
+	if err = g.createPermissionForStaticIps(g.config.InfisicalStaticIp); err != nil {
+		return err
+	}
+
 	g.registerHeartBeat(ctx, errCh)
 
 	cert, err := tls.X509KeyPair([]byte(gatewayCert.Certificate), []byte(gatewayCert.PrivateKey))
@@ -171,8 +176,7 @@ func (g *Gateway) Listen(ctx context.Context) error {
 		KeepAlivePeriod: 2 * time.Second,
 	}
 
-	g.registerRelayIsActive(ctx, relayUdpConnection.LocalAddr().String(), tlsConfig, quicConfig, errCh)
-
+	g.registerRelayIsActive(ctx, relayUdpConnection.LocalAddr().String(), errCh)
 	quicListener, err := quic.Listen(relayUdpConnection, tlsConfig, quicConfig)
 	if err != nil {
 		return fmt.Errorf("Failed to listen for QUIC: %w", err)
@@ -234,6 +238,8 @@ func (g *Gateway) Listen(ctx context.Context) error {
 		}
 	}()
 
+	// make this compatiable with systemd notify mode
+	systemd.SdNotify(false, systemd.SdNotifyReady)
 	select {
 	case <-ctx.Done():
 		log.Info().Msg("Shutting down gateway...")
@@ -282,90 +288,76 @@ func (g *Gateway) registerHeartBeat(ctx context.Context, errCh chan error) {
 	}()
 }
 
-func (g *Gateway) registerRelayIsActive(ctx context.Context, serverAddr string, tlsConf *tls.Config, quicConf *quic.Config, errCh chan error) {
-	ticker := time.NewTicker(5 * time.Second)
-	maxFailures := 3
-	failures := 0
+func (g *Gateway) createPermissionForStaticIps(staticIps string) error {
+	if staticIps == "" {
+		return fmt.Errorf("Missing Infisical static ips for permission")
+	}
 
-	go func() {
-		time.Sleep(2 * time.Second)
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-ticker.C:
-				conn, err := quic.DialAddr(ctx, serverAddr, tlsConf, quicConf)
-				if conn != nil {
-					failures = 0
-					conn.CloseWithError(0, "connection closed")
-				}
+	splittedIps := strings.Split(staticIps, ",")
+	resolvedIps := make([]net.Addr, 0)
+	for _, ip := range splittedIps {
+		ip = strings.TrimSpace(ip)
+		if ip == "" {
+			continue
+		}
 
-				if err != nil && !strings.Contains(err.Error(), "tls: failed to verify certificate") {
-					failures++
-					log.Warn().Err(err).Int("failures", failures).Msg("Relay connection check failed")
+		// if port not specific allow all port
+		if !strings.Contains(ip, ":") {
+			ip = ip + ":0"
+		}
 
-					if failures >= maxFailures {
-						errCh <- fmt.Errorf("relay connection check failed: %w", err)
-					}
-				}
-			}
+		peerAddr, err := net.ResolveUDPAddr("udp", ip)
+		if err != nil {
+			return fmt.Errorf("Failed to resolve static ip for permission: %w", err)
 		}
-	}()
-}
 
-func (g *Gateway) registerPermissionRefresh(ctx context.Context, errCh chan error) {
-	if g.config.InfisicalStaticIp == "" {
-		return
+		resolvedIps = append(resolvedIps, peerAddr)
 	}
 
-	log.Info().Msg("Starting TURN permission refresh routine")
-
-	go func() {
-		ticker := time.NewTicker(30 * time.Second)
-		defer ticker.Stop()
+	if err := g.client.CreatePermission(resolvedIps...); err != nil {
+		return fmt.Errorf("Failed to set ip permission: %w", err)
+	}
+	return nil
+}
 
-		g.refreshPermission(errCh)
+func (g *Gateway) registerRelayIsActive(ctx context.Context, relayAddress string, errCh chan error) error {
+	ticker := time.NewTicker(10 * time.Second)
+	maxFailures := 3
+	failures := 0
 
+	go func() {
+		time.Sleep(2 * time.Second)
 		for {
 			select {
 			case <-ctx.Done():
-				log.Info().Msg("Context cancelled, stopping TURN permission refresh")
 				return
 			case <-ticker.C:
-				g.refreshPermission(errCh)
+				// Configure TLS to skip verification
+				tlsConfig := &tls.Config{
+					InsecureSkipVerify: true,
+					NextProtos:         []string{"infisical-gateway"},
+				}
+				quicConfig := &quic.Config{
+					EnableDatagrams: true,
+				}
+				func() {
+					checkCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
+					defer cancel()
+					conn, err := quic.DialAddr(checkCtx, relayAddress, tlsConfig, quicConfig)
+					if err != nil {
+						failures++
+						log.Warn().Err(err).Int("failures", failures).Msg("Relay connection check failed")
+						if failures >= maxFailures {
+							errCh <- fmt.Errorf("relay connection check failed: %w", err)
+						}
+					}
+					if conn != nil {
+						conn.CloseWithError(0, "closed")
+					}
+				}()
 			}
 		}
 	}()
-}
 
-func (g *Gateway) refreshPermission(errCh chan error) {
-	log.Info().Msg("Attempting to refresh TURN permission")
-	maxRetries := 3
-	retryDelay := 5 * time.Second
-
-	var lastErr error
-	for i := 0; i < maxRetries; i++ {
-		peerAddr, err := net.ResolveUDPAddr("udp", g.config.InfisicalStaticIp)
-		if err != nil {
-			log.Error().Err(err).Msg("Failed to resolve static IP for permission refresh")
-			continue
-		}
-
-		if err := g.client.CreatePermission(peerAddr); err != nil {
-			lastErr = err
-			log.Warn().Err(err).Int("attempt", i+1).Msg("Failed to refresh TURN permission, retrying...")
-			time.Sleep(retryDelay)
-			continue
-		}
-
-		log.Info().Msg("Successfully refreshed TURN permission")
-		return
-	}
-
-	if lastErr != nil {
-		log.Error().Err(lastErr).Msg("Failed to refresh TURN permission after retries")
-		if reconnectErr := g.ConnectWithRelay(); reconnectErr != nil {
-			errCh <- fmt.Errorf("failed to refresh permissions and reconnect: %w", reconnectErr)
-		}
-	}
+	return nil
 }

cli/packages/gateway/relay.go

@@ -13,6 +13,7 @@ import (
 	"syscall"
 
 	udplistener "github.com/Infisical/infisical-merge/packages/gateway/udp_listener"
+	"github.com/Infisical/infisical-merge/packages/systemd"
 	"github.com/pion/logging"
 	"github.com/pion/turn/v4"
 	"github.com/rs/zerolog/log"
@@ -164,6 +165,9 @@ func (g *GatewayRelay) Run() error {
 	}
 
 	log.Info().Msgf("Relay listening on %s\n", connAddress)
+
+	// make this compatiable with systemd notify mode
+	systemd.SdNotify(false, systemd.SdNotifyReady)
 	// Block until user sends SIGINT or SIGTERM
 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

cli/packages/systemd/daemon.go

@@ -0,0 +1,84 @@
+// Copyright 2014 Docker, Inc.
+// Copyright 2015-2018 CoreOS, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+// Package daemon provides a Go implementation of the sd_notify protocol.
+// It can be used to inform systemd of service start-up completion, watchdog
+// events, and other status changes.
+//
+// https://www.freedesktop.org/software/systemd/man/sd_notify.html#Description
+package systemd
+
+import (
+	"net"
+	"os"
+)
+
+const (
+	// SdNotifyReady tells the service manager that service startup is finished
+	// or the service finished loading its configuration.
+	SdNotifyReady = "READY=1"
+
+	// SdNotifyStopping tells the service manager that the service is beginning
+	// its shutdown.
+	SdNotifyStopping = "STOPPING=1"
+
+	// SdNotifyReloading tells the service manager that this service is
+	// reloading its configuration. Note that you must call SdNotifyReady when
+	// it completed reloading.
+	SdNotifyReloading = "RELOADING=1"
+
+	// SdNotifyWatchdog tells the service manager to update the watchdog
+	// timestamp for the service.
+	SdNotifyWatchdog = "WATCHDOG=1"
+)
+
+// SdNotify sends a message to the init daemon. It is common to ignore the error.
+// If `unsetEnvironment` is true, the environment variable `NOTIFY_SOCKET`
+// will be unconditionally unset.
+//
+// It returns one of the following:
+// (false, nil) - notification not supported (i.e. NOTIFY_SOCKET is unset)
+// (false, err) - notification supported, but failure happened (e.g. error connecting to NOTIFY_SOCKET or while sending data)
+// (true, nil) - notification supported, data has been sent
+func SdNotify(unsetEnvironment bool, state string) (bool, error) {
+	socketAddr := &net.UnixAddr{
+		Name: os.Getenv("NOTIFY_SOCKET"),
+		Net:  "unixgram",
+	}
+
+	// NOTIFY_SOCKET not set
+	if socketAddr.Name == "" {
+		return false, nil
+	}
+
+	if unsetEnvironment {
+		if err := os.Unsetenv("NOTIFY_SOCKET"); err != nil {
+			return false, err
+		}
+	}
+
+	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)
+	// Error connecting to NOTIFY_SOCKET
+	if err != nil {
+		return false, err
+	}
+	defer conn.Close()
+
+	if _, err = conn.Write([]byte(state)); err != nil {
+		return false, err
+	}
+	return true, nil
+}