diff --git a/backend/src/ee/services/dynamic-secret/providers/sql-database.ts b/backend/src/ee/services/dynamic-secret/providers/sql-database.ts index 32bac4e6a..df1e412ae 100644 --- a/backend/src/ee/services/dynamic-secret/providers/sql-database.ts +++ b/backend/src/ee/services/dynamic-secret/providers/sql-database.ts @@ -51,7 +51,6 @@ export const SqlDatabaseProvider = ({ gatewayService }: TSqlDatabaseProviderDTO) user: providerInputs.username, password: providerInputs.password, ssl, - pool: { min: 0, max: 1 }, // @ts-expect-error this is because of knexjs type signature issue. This is directly passed to driver // https://github.com/knex/knex/blob/b6507a7129d2b9fafebf5f831494431e64c6a8a0/lib/dialects/mssql/index.js#L66 // https://github.com/tediousjs/tedious/blob/ebb023ed90969a7ec0e4b036533ad52739d921f7/test/config.ci.ts#L19 @@ -106,10 +105,8 @@ export const SqlDatabaseProvider = ({ gatewayService }: TSqlDatabaseProviderDTO) }; if (providerInputs.projectGatewayId) { - console.log(">>>>>> inside gateway"); await gatewayProxyWrapper(providerInputs, gatewayCallback); } else { - console.log(">>>>>> outside gateway"); await gatewayCallback(); } return isConnected; @@ -121,24 +118,27 @@ export const SqlDatabaseProvider = ({ gatewayService }: TSqlDatabaseProviderDTO) const password = generatePassword(providerInputs.client); const gatewayCallback = async (host = providerInputs.host, port = providerInputs.port) => { const db = await $getClient({ ...providerInputs, port, host }); - const { database } = providerInputs; - const expiration = new Date(expireAt).toISOString(); + try { + const { database } = providerInputs; + const expiration = new Date(expireAt).toISOString(); - const creationStatement = handlebars.compile(providerInputs.creationStatement, { noEscape: true })({ - username, - password, - expiration, - database - }); + const creationStatement = handlebars.compile(providerInputs.creationStatement, { noEscape: true })({ + username, + password, + expiration, + database + }); - const queries = creationStatement.toString().split(";").filter(Boolean); - await db.transaction(async (tx) => { - for (const query of queries) { - // eslint-disable-next-line - await tx.raw(query); - } - }); - await db.destroy(); + const queries = creationStatement.toString().split(";").filter(Boolean); + await db.transaction(async (tx) => { + for (const query of queries) { + // eslint-disable-next-line + await tx.raw(query); + } + }); + } finally { + await db.destroy(); + } }; if (providerInputs.projectGatewayId) { await gatewayProxyWrapper(providerInputs, gatewayCallback); @@ -154,16 +154,18 @@ export const SqlDatabaseProvider = ({ gatewayService }: TSqlDatabaseProviderDTO) const { database } = providerInputs; const gatewayCallback = async (host = providerInputs.host, port = providerInputs.port) => { const db = await $getClient({ ...providerInputs, port, host }); - const revokeStatement = handlebars.compile(providerInputs.revocationStatement)({ username, database }); - const queries = revokeStatement.toString().split(";").filter(Boolean); - await db.transaction(async (tx) => { - for (const query of queries) { - // eslint-disable-next-line - await tx.raw(query); - } - }); - - await db.destroy(); + try { + const revokeStatement = handlebars.compile(providerInputs.revocationStatement)({ username, database }); + const queries = revokeStatement.toString().split(";").filter(Boolean); + await db.transaction(async (tx) => { + for (const query of queries) { + // eslint-disable-next-line + await tx.raw(query); + } + }); + } finally { + await db.destroy(); + } }; if (providerInputs.projectGatewayId) { await gatewayProxyWrapper(providerInputs, gatewayCallback); @@ -187,18 +189,19 @@ export const SqlDatabaseProvider = ({ gatewayService }: TSqlDatabaseProviderDTO) expiration, database }); - - if (renewStatement) { - const queries = renewStatement.toString().split(";").filter(Boolean); - await db.transaction(async (tx) => { - for (const query of queries) { - // eslint-disable-next-line - await tx.raw(query); - } - }); + try { + if (renewStatement) { + const queries = renewStatement.toString().split(";").filter(Boolean); + await db.transaction(async (tx) => { + for (const query of queries) { + // eslint-disable-next-line + await tx.raw(query); + } + }); + } + } finally { + await db.destroy(); } - - await db.destroy(); }; if (providerInputs.projectGatewayId) { await gatewayProxyWrapper(providerInputs, gatewayCallback); diff --git a/cli/packages/cmd/gateway.go b/cli/packages/cmd/gateway.go index e8591ca68..a0b664664 100644 --- a/cli/packages/cmd/gateway.go +++ b/cli/packages/cmd/gateway.go @@ -53,34 +53,57 @@ var gatewayCmd = &cobra.Command{ <-sigCh close(sigStopCh) cancel() + + // If we get a second signal, force exit + <-sigCh + log.Warn().Msgf("Force exit triggered") + os.Exit(1) }() + // Main gateway retry loop with proper context handling + retryTicker := time.NewTicker(5 * time.Second) + defer retryTicker.Stop() + for { - select { - case <-sigStopCh: + if ctx.Err() != nil { log.Info().Msg("Shutting down gateway") return - default: - gatewayInstance, err := gateway.NewGateway(token.Token) - if err != nil { - util.HandleError(err) - } + } + gatewayInstance, err := gateway.NewGateway(token.Token) + if err != nil { + util.HandleError(err) + } - if err = gatewayInstance.ConnectWithRelay(); err != nil { - log.Error().Msgf("Gateway connection error with relay: %s", err) - log.Info().Msg("Restarting gateway...") - time.Sleep(5 * time.Second) - continue - } - err = gatewayInstance.Listen(ctx) - if err == nil { - // meaning everything went smooth and we are exiting + if err = gatewayInstance.ConnectWithRelay(); err != nil { + if ctx.Err() != nil { + log.Info().Msg("Shutting down gateway") return } - log.Error().Msgf("Gateway listen error: %s", err) - log.Info().Msg("Restarting gateway...") - time.Sleep(5 * time.Second) + log.Error().Msgf("Gateway connection error with relay: %s", err) + log.Info().Msg("Retrying connection in 5 seconds...") + select { + case <-retryTicker.C: + continue + case <-ctx.Done(): + log.Info().Msg("Shutting down gateway") + return + } + } + + err = gatewayInstance.Listen(ctx) + if ctx.Err() != nil { + log.Info().Msg("Gateway shutdown complete") + return + } + log.Error().Msgf("Gateway listen error: %s", err) + log.Info().Msg("Retrying connection in 5 seconds...") + select { + case <-retryTicker.C: + continue + case <-ctx.Done(): + log.Info().Msg("Shutting down gateway") + return } } }, diff --git a/cli/packages/gateway/gateway.go b/cli/packages/gateway/gateway.go index e5ec03293..5960e4ed5 100644 --- a/cli/packages/gateway/gateway.go +++ b/cli/packages/gateway/gateway.go @@ -112,6 +112,7 @@ func (g *Gateway) Listen(ctx context.Context) error { } log.Info().Msg("Connected with relay") + // Allocate a relay socket on the TURN server. On success, it // will return a net.PacketConn which represents the remote // socket. @@ -139,7 +140,7 @@ func (g *Gateway) Listen(ctx context.Context) error { g.config.Certificate = gatewayCert.Certificate g.config.CertificateChain = gatewayCert.CertificateChain - done := make(chan bool, 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) @@ -150,7 +151,7 @@ func (g *Gateway) Listen(ctx context.Context) error { g.registerPermissionLifecycle(func() error { err := relayNonTlsConn.CreatePermissions(peerAddr) return err - }, done) + }, shutdownCh) } cert, err := tls.X509KeyPair([]byte(gatewayCert.Certificate), []byte(gatewayCert.PrivateKey)) @@ -170,21 +171,32 @@ func (g *Gateway) Listen(ctx context.Context) error { errCh := make(chan error, 1) log.Info().Msg("Gateway started successfully") - g.registerHeartBeat(errCh, done) - g.registerRelayIsActive(relayNonTlsConn.Addr().String(), errCh, done) + g.registerHeartBeat(errCh, shutdownCh) + g.registerRelayIsActive(relayNonTlsConn.Addr().String(), 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 <-done: + case <-ctx.Done(): + return + case <-shutdownCh: return default: // Accept new relay connection conn, err := relayConn.Accept() 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) } @@ -198,7 +210,11 @@ func (g *Gateway) Listen(ctx context.Context) error { 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() @@ -212,34 +228,54 @@ func (g *Gateway) Listen(ctx context.Context) error { commonName := state.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() continue } } // Handle the connection in a goroutine wg.Add(1) - go func() { + go func(c net.Conn) { defer wg.Done() - handleConnection(conn) - }() + defer c.Close() + + // Monitor parent context to close this connection when needed + go func() { + select { + case <-ctx.Done(): + c.Close() // Force close connection when context is canceled + case <-shutdownCh: + c.Close() // Force close connection when accepting loop is done + } + }() + + handleConnection(c) + }(conn) } } }() - var isShutdown bool select { case <-ctx.Done(): log.Info().Msg("Shutting down gateway...") - isShutdown = true case err = <-errCh: } // Signal the accept loop to stop - close(done) - wg.Wait() + close(shutdownCh) - if isShutdown { - log.Info().Msg("Gateway shutdown complete") + // Set a timeout for waiting on connections to close + waitCh := make(chan struct{}) + go func() { + wg.Wait() + close(waitCh) + }() + + select { + case <-waitCh: + // All connections closed normally + case <-time.After(5 * time.Second): + log.Warn().Msg("Timeout waiting for connections to close gracefully") } return err