Repository: ory/fosite (go)

Issue:
Rollback is not called if there is a panic
### Preflight checklist

- [X] I could not find a solution in the existing issues, docs, nor discussions.
- [X] I agree to follow this project's [Code of Conduct](https://github.com/ory/fosite/blob/master/CODE_OF_CONDUCT.md).
- [X] I have read and am following this repository's [Contribution Guidelines](https://github.com/ory/fosite/blob/master/CONTRIBUTING.md).
- [ ] This issue affects my [Ory Cloud](https://www.ory.sh/) project.
- [X] I have joined the [Ory Community Slack](https://slack.ory.sh).
- [ ] I am signed up to the [Ory Security Patch Newsletter](https://ory.us10.list-manage.com/subscribe?u=ffb1a878e4ec6c0ed312a3480&id=f605a41b53).

### Describe the bug

Current use of `BeginTx`, `CommitTx` and `RollbackTx` in Fosite is problematic because it does not call `RollbackTx` when panic happens. We could assume that Fosite itself does not panic, but that is not necessary true for code it calls into. For example, I had a bug inside my storage implementation which triggered panic. As the panic propagated and was handled higher up in my request handler and logged, transaction itself was not rolled back, leaving the connection over it in an inconsistent state. With connection pooling that brought some stability issues.

### Reproducing the bug

Sadly I do not have a simple reproduction here. One which I was able to do is with Sqlite and setting `SetMaxOpenConns` to 1. Then if a transaction is not rolled back, no other operation can be done on the database and my code hangs. This is how I am now testing that all transactions are properly rolled back.

### Relevant log output

_No response_

### Relevant configuration

_No response_

### Version

observed on current master branch

### On which operating system are you observing this issue?

_No response_

### In which environment are you deploying?

_No response_

### Additional Context

_No response_

Pull request description:
fix: always rollback
## Related Issue or Design Document

Fixes #637.

## Checklist

- [x] I have read the [contributing guidelines](../blob/master/CONTRIBUTING.md)
      and signed the CLA.
- [x] I have referenced an issue containing the design document if my change
      introduces a new feature.
- [x] I have read the [security policy](../security/policy).
- [x] I confirm that this pull request does not address a security
      vulnerability. If this pull request addresses a security vulnerability, I
      confirm that I got green light (please contact
      [security@ory.sh](mailto:security@ory.sh)) from the maintainers to push
      the changes.
- [x] I have added tests that prove my fix is effective or that my feature
      works.
- [ ] I have added necessary documentation within the code base (if
      appropriate).

## Further comments

There is not much change necessary to get this fixed. See the diff.

There is one logical change I have done: before rollback was not called if commit failed, now I made it that it is called if commit fails. It is generally safe to call rollback outside of a transaction in PostgreSQL (you only get warning), but if the connection is in error state, then rollback is required. So just to be safe, I think it is better to rollback on any error, including when commit fails (because we might not know what is the reason for failure of that call and in what state is the connection).

Patch:
diff --git a/handler/oauth2/flow_authorize_code_token.go b/handler/oauth2/flow_authorize_code_token.go
index 986663e..bf14933 100644
--- a/handler/oauth2/flow_authorize_code_token.go
+++ b/handler/oauth2/flow_authorize_code_token.go
@@ -130,7 +130,7 @@ func canIssueRefreshToken(c *AuthorizeExplicitGrantHandler, request fosite.Reque
 	return true
 }
 
-func (c *AuthorizeExplicitGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) error {
+func (c *AuthorizeExplicitGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) (err error) {
 	if !c.CanHandleTokenEndpointRequest(requester) {
 		return errorsx.WithStack(fosite.ErrUnknownRequest)
 	}
@@ -170,22 +170,20 @@ func (c *AuthorizeExplicitGrantHandler) PopulateTokenEndpointResponse(ctx contex
 	if err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
-
-	if err := c.CoreStorage.InvalidateAuthorizeCodeSession(ctx, signature); err != nil {
-		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.CoreStorage); rollBackTxnErr != nil {
-			return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebugf("error: %s; rollback error: %s", err, rollBackTxnErr))
+	defer func() {
+		if err != nil {
+			if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.CoreStorage); rollBackTxnErr != nil {
+				err = errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebugf("error: %s; rollback error: %s", err, rollBackTxnErr))
+			}
 		}
+	}()
+
+	if err = c.CoreStorage.InvalidateAuthorizeCodeSession(ctx, signature); err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
-	} else if err := c.CoreStorage.CreateAccessTokenSession(ctx, accessSignature, requester.Sanitize([]string{})); err != nil {
-		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.CoreStorage); rollBackTxnErr != nil {
-			return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebugf("error: %s; rollback error: %s", err, rollBackTxnErr))
-		}
+	} else if err = c.CoreStorage.CreateAccessTokenSession(ctx, accessSignature, requester.Sanitize([]string{})); err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	} else if refreshSignature != "" {
-		if err := c.CoreStorage.CreateRefreshTokenSession(ctx, refreshSignature, requester.Sanitize([]string{})); err != nil {
-			if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.CoreStorage); rollBackTxnErr != nil {
-				return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebugf("error: %s; rollback error: %s", err, rollBackTxnErr))
-			}
+		if err = c.CoreStorage.CreateRefreshTokenSession(ctx, refreshSignature, requester.Sanitize([]string{})); err != nil {
 			return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 		}
 	}
@@ -198,7 +196,7 @@ func (c *AuthorizeExplicitGrantHandler) PopulateTokenEndpointResponse(ctx contex
 		responder.SetExtra("refresh_token", refresh)
 	}
 
-	if err := storage.MaybeCommitTx(ctx, c.CoreStorage); err != nil {
+	if err = storage.MaybeCommitTx(ctx, c.CoreStorage); err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
 
diff --git a/handler/oauth2/flow_refresh.go b/handler/oauth2/flow_refresh.go
index 1930756..60e58c3 100644
--- a/handler/oauth2/flow_refresh.go
+++ b/handler/oauth2/flow_refresh.go
@@ -121,7 +121,7 @@ func (c *RefreshTokenGrantHandler) HandleTokenEndpointRequest(ctx context.Contex
 }
 
 // PopulateTokenEndpointResponse implements https://tools.ietf.org/html/rfc6749#section-6
-func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) error {
+func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) (err error) {
 	if !c.CanHandleTokenEndpointRequest(requester) {
 		return errorsx.WithStack(fosite.ErrUnknownRequest)
 	}
@@ -142,27 +142,30 @@ func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Con
 	if err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
+	defer func() {
+		err = c.handleRefreshTokenEndpointStorageError(ctx, err)
+	}()
 
 	ts, err := c.TokenRevocationStorage.GetRefreshTokenSession(ctx, signature, nil)
 	if err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+		return err
 	} else if err := c.TokenRevocationStorage.RevokeAccessToken(ctx, ts.GetID()); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+		return err
 	}
 
 	if err := c.TokenRevocationStorage.RevokeRefreshTokenMaybeGracePeriod(ctx, ts.GetID(), signature); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+		return err
 	}
 
 	storeReq := requester.Sanitize([]string{})
 	storeReq.SetID(ts.GetID())
 
-	if err := c.TokenRevocationStorage.CreateAccessTokenSession(ctx, accessSignature, storeReq); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+	if err = c.TokenRevocationStorage.CreateAccessTokenSession(ctx, accessSignature, storeReq); err != nil {
+		return err
 	}
 
-	if err := c.TokenRevocationStorage.CreateRefreshTokenSession(ctx, refreshSignature, storeReq); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+	if err = c.TokenRevocationStorage.CreateRefreshTokenSession(ctx, refreshSignature, storeReq); err != nil {
+		return err
 	}
 
 	responder.SetAccessToken(accessToken)
@@ -171,8 +174,8 @@ func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Con
 	responder.SetScopes(requester.GetGrantedScopes())
 	responder.SetExtra("refresh_token", refreshToken)
 
-	if err := storage.MaybeCommitTx(ctx, c.TokenRevocationStorage); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, false, err)
+	if err = storage.MaybeCommitTx(ctx, c.TokenRevocationStorage); err != nil {
+		return err
 	}
 
 	return nil
@@ -188,37 +191,42 @@ func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Con
 //     attempt the valid refresh token and the access authorization
 //     associated with it are both revoked.
 //
-func (c *RefreshTokenGrantHandler) handleRefreshTokenReuse(ctx context.Context, signature string, req fosite.Requester) error {
-	ctx, err := storage.MaybeBeginTx(ctx, c.TokenRevocationStorage)
+func (c *RefreshTokenGrantHandler) handleRefreshTokenReuse(ctx context.Context, signature string, req fosite.Requester) (err error) {
+	ctx, err = storage.MaybeBeginTx(ctx, c.TokenRevocationStorage)
 	if err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
+	defer func() {
+		err = c.handleRefreshTokenEndpointStorageError(ctx, err)
+	}()
 
-	if err := c.TokenRevocationStorage.DeleteRefreshTokenSession(ctx, signature); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
-	} else if err := c.TokenRevocationStorage.RevokeRefreshToken(
+	if err = c.TokenRevocationStorage.DeleteRefreshTokenSession(ctx, signature); err != nil {
+		return err
+	} else if err = c.TokenRevocationStorage.RevokeRefreshToken(
 		ctx, req.GetID(),
 	); err != nil && !errors.Is(err, fosite.ErrNotFound) {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
-	} else if err := c.TokenRevocationStorage.RevokeAccessToken(
+		return err
+	} else if err = c.TokenRevocationStorage.RevokeAccessToken(
 		ctx, req.GetID(),
 	); err != nil && !errors.Is(err, fosite.ErrNotFound) {
-		return c.handleRefreshTokenEndpointStorageError(ctx, true, err)
+		return err
 	}
 
-	if err := storage.MaybeCommitTx(ctx, c.TokenRevocationStorage); err != nil {
-		return c.handleRefreshTokenEndpointStorageError(ctx, false, err)
+	if err = storage.MaybeCommitTx(ctx, c.TokenRevocationStorage); err != nil {
+		return err
 	}
 
 	return nil
 }
 
-func (c *RefreshTokenGrantHandler) handleRefreshTokenEndpointStorageError(ctx context.Context, rollback bool, storageErr error) (err error) {
+func (c *RefreshTokenGrantHandler) handleRefreshTokenEndpointStorageError(ctx context.Context, storageErr error) (err error) {
+	if storageErr == nil {
+		return nil
+	}
+
 	defer func() {
-		if rollback {
-			if rbErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rbErr != nil {
-				err = errorsx.WithStack(fosite.ErrServerError.WithWrap(rbErr).WithDebug(rbErr.Error()))
-			}
+		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
+			err = errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebugf("error: %s; rollback error: %s", err, rollBackTxnErr))
 		}
 	}()
 


Test patch:
diff --git a/handler/oauth2/flow_authorize_code_token_test.go b/handler/oauth2/flow_authorize_code_token_test.go
index 3e8fb52..8439e95 100644
--- a/handler/oauth2/flow_authorize_code_token_test.go
+++ b/handler/oauth2/flow_authorize_code_token_test.go
@@ -649,6 +649,11 @@ func TestAuthorizeCodeTransactional_HandleTokenEndpointRequest(t *testing.T) {
 					Commit(propagatedContext).
 					Return(errors.New("Whoops, unable to commit transaction!")).
 					Times(1)
+				mockTransactional.
+					EXPECT().
+					Rollback(propagatedContext).
+					Return(nil).
+					Times(1)
 			},
 			expectError: fosite.ErrServerError,
 		},
diff --git a/handler/oauth2/flow_refresh_test.go b/handler/oauth2/flow_refresh_test.go
index 95e86c5..dd0abf1 100644
--- a/handler/oauth2/flow_refresh_test.go
+++ b/handler/oauth2/flow_refresh_test.go
@@ -915,6 +915,11 @@ func TestRefreshFlowTransactional_PopulateTokenEndpointResponse(t *testing.T) {
 					Commit(propagatedContext).
 					Return(errors.New("Could not commit transaction!")).
 					Times(1)
+				mockTransactional.
+					EXPECT().
+					Rollback(propagatedContext).
+					Return(nil).
+					Times(1)
 			},
 			expectError: fosite.ErrServerError,
 		},
@@ -958,6 +963,11 @@ func TestRefreshFlowTransactional_PopulateTokenEndpointResponse(t *testing.T) {
 					Commit(propagatedContext).
 					Return(fosite.ErrSerializationFailure).
 					Times(1)
+				mockTransactional.
+					EXPECT().
+					Rollback(propagatedContext).
+					Return(nil).
+					Times(1)
 			},
 			expectError: fosite.ErrInvalidRequest,
 		},