From fce930e0990d83d2177b45ebe376b037b317f151 Mon Sep 17 00:00:00 2001
From: aeneasr <3372410+aeneasr@users.noreply.github.com>
Date: Fri, 27 Oct 2023 10:04:10 +0200
Subject: [PATCH] chore: improve csrf context

---
 consent/csrf.go             |  3 ++-
 consent/handler.go          |  1 -
 consent/helper_test.go      |  2 +-
 consent/strategy_default.go |  4 ++--
 flow/consent_types.go       | 10 ++++++++++
 flow/flow.go                |  2 ++
 6 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/consent/csrf.go b/consent/csrf.go
index 42588390a52..cd691febe9c 100644
--- a/consent/csrf.go
+++ b/consent/csrf.go
@@ -6,6 +6,7 @@
 package consent
 
 import (
+	"github.com/ory/hydra/v2/flow"
 	"net/http"
 	"strings"
 	"time"
@@ -45,7 +46,7 @@ func createCsrfSession(w http.ResponseWriter, r *http.Request, conf x.CookieConf
 	return nil
 }
 
-func validateCsrfSession(r *http.Request, conf x.CookieConfigProvider, store sessions.Store, name, expectedCSRF string) error {
+func ValidateCsrfSession(r *http.Request, conf x.CookieConfigProvider, store sessions.Store, name, expectedCSRF string, f *flow.Flow) error {
 	if cookie, err := getCsrfSession(r, store, conf, name); err != nil {
 		return errorsx.WithStack(fosite.ErrRequestForbidden.WithHint("CSRF session cookie could not be decoded."))
 	} else if csrf, err := mapx.GetString(cookie.Values, "csrf"); err != nil {
diff --git a/consent/handler.go b/consent/handler.go
index d0d3fd2aa2b..e313b7229e6 100644
--- a/consent/handler.go
+++ b/consent/handler.go
@@ -492,7 +492,6 @@ func (h *Handler) acceptOAuth2LoginRequest(w http.ResponseWriter, r *http.Reques
 	}
 
 	events.Trace(ctx, events.LoginAccepted, events.WithClientID(request.Client.GetID()), events.WithSubject(request.Subject))
-
 	h.r.Writer().Write(w, r, &flow.OAuth2RedirectTo{
 		RedirectTo: urlx.SetQuery(ru, url.Values{"login_verifier": {verifier}}).String(),
 	})
diff --git a/consent/helper_test.go b/consent/helper_test.go
index a5f09e81cdd..4a347575f3a 100644
--- a/consent/helper_test.go
+++ b/consent/helper_test.go
@@ -267,7 +267,7 @@ func TestValidateCsrfSession(t *testing.T) {
 				assert.NoError(t, err, "failed to save cookie %s", c.name)
 			}
 
-			err := validateCsrfSession(r, config, store, name, tc.csrfValue)
+			err := ValidateCsrfSession(r, config, store, name, tc.csrfValue, new(flow.Flow))
 			if tc.expectError {
 				assert.Error(t, err)
 			} else {
diff --git a/consent/strategy_default.go b/consent/strategy_default.go
index 94ed48df9d0..8b3d21b3e92 100644
--- a/consent/strategy_default.go
+++ b/consent/strategy_default.go
@@ -380,7 +380,7 @@ func (s *DefaultStrategy) verifyAuthentication(
 	}
 
 	clientSpecificCookieNameLoginCSRF := fmt.Sprintf("%s_%s", s.r.Config().CookieNameLoginCSRF(ctx), session.LoginRequest.Client.CookieSuffix())
-	if err := validateCsrfSession(r, s.r.Config(), store, clientSpecificCookieNameLoginCSRF, session.LoginRequest.CSRF); err != nil {
+	if err := ValidateCsrfSession(r, s.r.Config(), store, clientSpecificCookieNameLoginCSRF, session.LoginRequest.CSRF, f); err != nil {
 		return nil, err
 	}
 
@@ -689,7 +689,7 @@ func (s *DefaultStrategy) verifyConsent(ctx context.Context, _ http.ResponseWrit
 	}
 
 	clientSpecificCookieNameConsentCSRF := fmt.Sprintf("%s_%s", s.r.Config().CookieNameConsentCSRF(ctx), session.ConsentRequest.Client.CookieSuffix())
-	if err := validateCsrfSession(r, s.r.Config(), store, clientSpecificCookieNameConsentCSRF, session.ConsentRequest.CSRF); err != nil {
+	if err := ValidateCsrfSession(r, s.r.Config(), store, clientSpecificCookieNameConsentCSRF, session.ConsentRequest.CSRF, f); err != nil {
 		return nil, nil, err
 	}
 
diff --git a/flow/consent_types.go b/flow/consent_types.go
index 9a2666c7867..31759d27d40 100644
--- a/flow/consent_types.go
+++ b/flow/consent_types.go
@@ -174,6 +174,11 @@ type AcceptOAuth2ConsentRequest struct {
 	// the flow.
 	WasHandled bool `json:"-"`
 
+	// Context is an optional object which can hold arbitrary data. The data will be made available when fetching the
+	// consent request under the "context" field. This is useful in scenarios where login and consent endpoints share
+	// data.
+	Context sqlxx.JSONRawMessage `json:"context"`
+
 	ConsentRequest  *OAuth2ConsentRequest `json:"-"`
 	Error           *RequestDeniedError   `json:"-"`
 	RequestedAt     time.Time             `json:"-"`
@@ -240,6 +245,11 @@ type OAuth2ConsentSession struct {
 	// the flow.
 	WasHandled bool `json:"-" db:"was_used"`
 
+	// Context is an optional object which can hold arbitrary data. The data will be made available when fetching the
+	// consent request under the "context" field. This is useful in scenarios where login and consent endpoints share
+	// data.
+	Context sqlxx.JSONRawMessage `json:"context"`
+
 	// Consent Request
 	//
 	// The consent request that lead to this consent session.
diff --git a/flow/flow.go b/flow/flow.go
index 7e8eeb077c8..1fd5affc4f6 100644
--- a/flow/flow.go
+++ b/flow/flow.go
@@ -394,6 +394,7 @@ func (f *Flow) HandleConsentRequest(r *AcceptOAuth2ConsentRequest) error {
 	f.ConsentHandledAt = r.HandledAt
 	f.ConsentWasHandled = r.WasHandled
 	f.ConsentError = r.Error
+	f.Context = r.Context
 
 	if r.Session != nil {
 		f.SessionIDToken = r.Session.IDToken
@@ -458,6 +459,7 @@ func (f *Flow) GetHandledConsentRequest() *AcceptOAuth2ConsentRequest {
 		RememberFor:        crf,
 		HandledAt:          f.ConsentHandledAt,
 		WasHandled:         f.ConsentWasHandled,
+		Context:            f.Context,
 		ConsentRequest:     f.GetConsentRequest(),
 		Error:              f.ConsentError,
 		RequestedAt:        f.RequestedAt,