Skip to content

Commit

Permalink
feat: pass transaction to invokeHook, fixing pool exhaustion (supab…
Browse files Browse the repository at this point in the history
…ase#1465)

When invoking a hook, if the invocation is done within an already open
transaction, `invokeHook` should not try to open a new transaction. This
can easily cause connection pool exhaustion and deadlock.
  • Loading branch information
hf authored Mar 3, 2024
1 parent 7185af8 commit b536d36
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 11 deletions.
28 changes: 20 additions & 8 deletions internal/api/hooks.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (
"github.com/supabase/auth/internal/storage"
)

func (a *API) runHook(ctx context.Context, name string, input, output any) ([]byte, error) {
func (a *API) runHook(ctx context.Context, tx *storage.Connection, name string, input, output any) ([]byte, error) {
db := a.db.WithContext(ctx)

request, err := json.Marshal(input)
Expand All @@ -20,7 +20,7 @@ func (a *API) runHook(ctx context.Context, name string, input, output any) ([]by
}

var response []byte
if err := db.Transaction(func(tx *storage.Connection) error {
invokeHookFunc := func(tx *storage.Connection) error {
// We rely on Postgres timeouts to ensure the function doesn't overrun
if terr := tx.RawQuery(fmt.Sprintf("set local statement_timeout TO '%d';", hooks.DefaultTimeout)).Exec(); terr != nil {
return terr
Expand All @@ -36,8 +36,16 @@ func (a *API) runHook(ctx context.Context, name string, input, output any) ([]by
}

return nil
}); err != nil {
return nil, err
}

if tx != nil {
if err := invokeHookFunc(tx); err != nil {
return nil, err
}
} else {
if err := db.Transaction(invokeHookFunc); err != nil {
return nil, err
}
}

if err := json.Unmarshal(response, output); err != nil {
Expand All @@ -47,7 +55,11 @@ func (a *API) runHook(ctx context.Context, name string, input, output any) ([]by
return response, nil
}

func (a *API) invokeHook(ctx context.Context, input, output any) error {
// invokeHook invokes the hook code. tx can be nil, in which case a new
// transaction is opened. If calling invokeHook within a transaction, always
// pass the current transaciton, as pool-exhaustion deadlocks are very easy to
// trigger.
func (a *API) invokeHook(ctx context.Context, tx *storage.Connection, input, output any) error {
config := a.config
switch input.(type) {
case *hooks.MFAVerificationAttemptInput:
Expand All @@ -56,7 +68,7 @@ func (a *API) invokeHook(ctx context.Context, input, output any) error {
panic("output should be *hooks.MFAVerificationAttemptOutput")
}

if _, err := a.runHook(ctx, config.Hook.MFAVerificationAttempt.HookName, input, output); err != nil {
if _, err := a.runHook(ctx, tx, config.Hook.MFAVerificationAttempt.HookName, input, output); err != nil {
return internalServerError("Error invoking MFA verification hook.").WithInternalError(err)
}

Expand All @@ -82,7 +94,7 @@ func (a *API) invokeHook(ctx context.Context, input, output any) error {
panic("output should be *hooks.PasswordVerificationAttemptOutput")
}

if _, err := a.runHook(ctx, config.Hook.PasswordVerificationAttempt.HookName, input, output); err != nil {
if _, err := a.runHook(ctx, tx, config.Hook.PasswordVerificationAttempt.HookName, input, output); err != nil {
return internalServerError("Error invoking password verification hook.").WithInternalError(err)
}

Expand All @@ -108,7 +120,7 @@ func (a *API) invokeHook(ctx context.Context, input, output any) error {
panic("output should be *hooks.CustomAccessTokenOutput")
}

if _, err := a.runHook(ctx, config.Hook.CustomAccessToken.HookName, input, output); err != nil {
if _, err := a.runHook(ctx, tx, config.Hook.CustomAccessToken.HookName, input, output); err != nil {
return internalServerError("Error invoking access token hook.").WithInternalError(err)
}

Expand Down
2 changes: 1 addition & 1 deletion internal/api/mfa.go
Original file line number Diff line number Diff line change
Expand Up @@ -258,7 +258,7 @@ func (a *API) VerifyFactor(w http.ResponseWriter, r *http.Request) error {

output := hooks.MFAVerificationAttemptOutput{}

err := a.invokeHook(ctx, &input, &output)
err := a.invokeHook(ctx, nil, &input, &output)
if err != nil {
return err
}
Expand Down
4 changes: 2 additions & 2 deletions internal/api/token.go
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ func (a *API) ResourceOwnerPasswordGrant(ctx context.Context, w http.ResponseWri
Valid: isValidPassword,
}
output := hooks.PasswordVerificationAttemptOutput{}
err := a.invokeHook(ctx, &input, &output)
err := a.invokeHook(ctx, nil, &input, &output)
if err != nil {
return err
}
Expand Down Expand Up @@ -350,7 +350,7 @@ func (a *API) generateAccessToken(ctx context.Context, tx *storage.Connection, u

output := hooks.CustomAccessTokenOutput{}

err := a.invokeHook(ctx, &input, &output)
err := a.invokeHook(ctx, tx, &input, &output)
if err != nil {
return "", 0, err
}
Expand Down

0 comments on commit b536d36

Please sign in to comment.