From 3badcea95b60624038373cbdeef91ebc805f1eed Mon Sep 17 00:00:00 2001 From: Maidul Islam Date: Wed, 5 Mar 2025 01:07:36 -0500 Subject: [PATCH] added permission refresh and main context --- cli/packages/gateway/gateway.go | 121 ++++++++++++++++++++++---------- 1 file changed, 85 insertions(+), 36 deletions(-) diff --git a/cli/packages/gateway/gateway.go b/cli/packages/gateway/gateway.go index 66c88fd45..846231071 100644 --- a/cli/packages/gateway/gateway.go +++ b/cli/packages/gateway/gateway.go @@ -141,19 +141,11 @@ func (g *Gateway) Listen(ctx context.Context) error { g.config.Certificate = gatewayCert.Certificate g.config.CertificateChain = gatewayCert.CertificateChain + errCh := make(chan error, 1) shutdownCh := make(chan bool, 1) - if g.config.InfisicalStaticIp != "" { - log.Info().Msgf("Found static ip from Infisical: %s. Creating permission IP lifecycle", 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) - } - err = g.client.CreatePermission(peerAddr) - if err != nil { - return fmt.Errorf("Failed to set permission: %w", err) - } - } + g.registerPermissionRefresh(ctx, errCh) + g.registerHeartBeat(ctx, errCh) cert, err := tls.X509KeyPair([]byte(gatewayCert.Certificate), []byte(gatewayCert.PrivateKey)) if err != nil { @@ -175,10 +167,12 @@ func (g *Gateway) Listen(ctx context.Context) error { // Setup QUIC listener on the relayConn quicConfig := &quic.Config{ EnableDatagrams: true, - MaxIdleTimeout: 30 * time.Second, - KeepAlivePeriod: 15 * time.Second, + MaxIdleTimeout: 10 * time.Second, + KeepAlivePeriod: 2 * time.Second, } + g.registerRelayIsActive(ctx, relayUdpConnection.LocalAddr().String(), tlsConfig, quicConfig, errCh) + quicListener, err := quic.Listen(relayUdpConnection, tlsConfig, quicConfig) if err != nil { return fmt.Errorf("Failed to listen for QUIC: %w", err) @@ -187,12 +181,8 @@ func (g *Gateway) Listen(ctx context.Context) error { 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(relayUdpConnection.LocalAddr().String(), tlsConfig, quicConfig, errCh, shutdownCh) - // Create a WaitGroup to track active connections var wg sync.WaitGroup go func() { @@ -248,6 +238,7 @@ func (g *Gateway) Listen(ctx context.Context) error { case <-ctx.Done(): log.Info().Msg("Shutting down gateway...") case err = <-errCh: + log.Error().Err(err).Msg("Gateway error occurred") } // Signal the accept loop to stop @@ -270,7 +261,7 @@ func (g *Gateway) Listen(ctx context.Context) error { return err } -func (g *Gateway) registerHeartBeat(errCh chan error, done chan bool) { +func (g *Gateway) registerHeartBeat(ctx context.Context, errCh chan error) { ticker := time.NewTicker(30 * time.Minute) defer ticker.Stop() @@ -283,7 +274,7 @@ func (g *Gateway) registerHeartBeat(errCh chan error, done chan bool) { } select { - case <-done: + case <-ctx.Done(): return case <-ticker.C: } @@ -291,32 +282,90 @@ func (g *Gateway) registerHeartBeat(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) +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 go func() { - time.Sleep(5 * time.Second) + time.Sleep(2 * time.Second) for { select { - case <-done: - ticker.Stop() + case <-ctx.Done(): return case <-ticker.C: - func() { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - conn, err := quic.DialAddr(ctx, serverAddr, tlsConf, quicConf) - if conn != nil { - conn.CloseWithError(0, "connection closed") - } - // If we get a TLS verification error, that means the QUIC connection is working - // Any other error means the connection is not working - if err != nil && !strings.Contains(err.Error(), "tls: failed to verify certificate") { - log.Error().Err(err).Msg("Relay connection check failed") + conn, err := quic.DialAddr(ctx, serverAddr, tlsConf, quicConf) + if conn != nil { + failures = 0 + conn.CloseWithError(0, "connection closed") + } + + 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 failures >= maxFailures { errCh <- fmt.Errorf("relay connection check failed: %w", err) } - }() + } } } }() } + +func (g *Gateway) registerPermissionRefresh(ctx context.Context, errCh chan error) { + if g.config.InfisicalStaticIp == "" { + return + } + + log.Info().Msg("Starting TURN permission refresh routine") + + go func() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + g.refreshPermission(errCh) + + for { + select { + case <-ctx.Done(): + log.Info().Msg("Context cancelled, stopping TURN permission refresh") + return + case <-ticker.C: + g.refreshPermission(errCh) + } + } + }() +} + +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) + } + } +}