Repository: zitadel/oidc
Pull request #753 (2025-06-04 11:32:56)

feat: Request aware cookie handling
Hi :wave: firstly, thank you for this great library!

LXD is a hypervisor for system containers and VMs. It can be clustered to give a private cloud like experience.

When LXD is clustered, we have an issue wherein the authorization code flow fails if the host that receives the callback (receiver of `/oidc/callback`) is different from the host that begins the flow (receiver of `/oidc/login`; see https://github.com/canonical/lxd/issues/13644). This is because the cookie handler in the relying party has different hash and encryption keys for each cluster member. Thus, when the callback is executed, the relying party cannot decode the state or PKCE cookies and the callback fails.

As discussed in the linked issue, this can be mitigated via sticky sessions on an external load balancer. However, we are increasingly seeing LXD deployments placed behind DNS SRV records, where sticky sessions are not possible. 

We would like to be able to derive cookie encryption keys from some shared secret, so that the hash and encryption keys are the same on all cluster members. To obfuscate the secret we would like to use a nonce value during derivation. Ideally this nonce would be unique per client.

The relying party accepts a `CookieHandler`, which in turn accepts static hash and block keys. In order to use unique keys per-client (or on rotation), a new relying party must be created for each cookie handler. This is not ideal, because creation of a new relying party re-triggers OIDC discovery and creates a new remote key set (which will retrieve JWKs when required). 

This PR adds functionality to the CookieHandler to allow more dynamic handling of cookies, based on the incoming request. I have taken care not to break any existing APIs.

Please let me know what you think and what my next steps should be (e.g. testing)

### Definition of Ready

- [x] I am happy with the code
- [x] Short description of the feature/issue is added in the pr description
- [x] PR is linked to the corresponding user story
- [ ] Acceptance criteria are met (not sure)
- [ ] All open todos and follow ups are defined in a new ticket and justified (not sure)
- [ ] Deviations from the acceptance criteria and design are agreed with the PO and documented. (not sure)
- [x] No debug or dead code
- [ ] My code has no repetitions (some slight repetition to avoid API breakage)
- [ ] Critical parts are tested automatically (not sure)
- [ ] Where possible E2E tests are implemented
- [ ] Documentation/examples are up-to-date
- [ ] All non-functional requirements are met (not sure)
- [ ] Functionality of the acceptance criteria is checked manually on the dev system. (not sure)



markylaing: > Code looks great. Perhaps you can get some testing in by using (or creating a variant of) [testRelyingPartySession](https://github.com/zitadel/oidc/blob/7d57aaa99983368527b0fc93f72ab5542b2eefd2/pkg/client/integration_test.go#L63)

Great stuff I'll look into doing that and ping you when I have something :)
markylaing: @muhlemmer hi :wave: I've added integration tests now. Tests all ok locally :)
markylaing: > Thanks @markylaing, Tim is on hols atm, he's back in a couple of days, then he'll have a look again!

Great stuff thank you :pray: 

Patch:
diff --git a/pkg/client/rp/relying_party.go b/pkg/client/rp/relying_party.go
index e6fa078..93c5dd0 100644
--- a/pkg/client/rp/relying_party.go
+++ b/pkg/client/rp/relying_party.go
@@ -410,12 +410,19 @@ func AuthURLHandler(stateFn func() string, rp RelyingParty, urlParam ...URLParam
 		}
 
 		state := stateFn()
-		if err := trySetStateCookie(w, state, rp); err != nil {
+		if err := trySetStateCookie(r, w, state, rp); err != nil {
 			unauthorizedError(w, r, "failed to create state cookie: "+err.Error(), state, rp)
 			return
 		}
 		if rp.IsPKCE() {
-			codeChallenge, err := GenerateAndStoreCodeChallenge(w, rp)
+			var codeChallenge string
+			var err error
+			if rp.CookieHandler().IsRequestAware() {
+				codeChallenge, err = GenerateAndStoreCodeChallengeWithRequest(r, w, rp)
+			} else {
+				codeChallenge, err = GenerateAndStoreCodeChallenge(w, rp)
+			}
+
 			if err != nil {
 				unauthorizedError(w, r, "failed to create code challenge: "+err.Error(), state, rp)
 				return
@@ -436,6 +443,15 @@ func GenerateAndStoreCodeChallenge(w http.ResponseWriter, rp RelyingParty) (stri
 	return oidc.NewSHACodeChallenge(codeVerifier), nil
 }
 
+// GenerateAndStoreCodeChallenge generates a PKCE code challenge and stores its verifier into a secure cookie
+func GenerateAndStoreCodeChallengeWithRequest(r *http.Request, w http.ResponseWriter, rp RelyingParty) (string, error) {
+	codeVerifier := base64.RawURLEncoding.EncodeToString([]byte(uuid.New().String()))
+	if err := rp.CookieHandler().SetRequestAwareCookie(r, w, pkceCode, codeVerifier); err != nil {
+		return "", err
+	}
+	return oidc.NewSHACodeChallenge(codeVerifier), nil
+}
+
 // ErrMissingIDToken is returned when an id_token was expected,
 // but not received in the token response.
 var ErrMissingIDToken = errors.New("id_token missing")
@@ -607,9 +623,16 @@ func Userinfo[U SubjectGetter](ctx context.Context, token, tokenType, subject st
 	return userinfo, nil
 }
 
-func trySetStateCookie(w http.ResponseWriter, state string, rp RelyingParty) error {
+func trySetStateCookie(r *http.Request, w http.ResponseWriter, state string, rp RelyingParty) error {
 	if rp.CookieHandler() != nil {
-		if err := rp.CookieHandler().SetCookie(w, stateParam, state); err != nil {
+		var err error
+		if rp.CookieHandler().IsRequestAware() {
+			err = rp.CookieHandler().SetRequestAwareCookie(r, w, stateParam, state)
+		} else {
+			err = rp.CookieHandler().SetCookie(w, stateParam, state)
+		}
+
+		if err != nil {
 			return err
 		}
 	}
diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go
index 1ebc9e2..7724aeb 100644
--- a/pkg/http/cookie.go
+++ b/pkg/http/cookie.go
@@ -8,12 +8,13 @@ import (
 )
 
 type CookieHandler struct {
-	securecookie *securecookie.SecureCookie
-	secureOnly   bool
-	sameSite     http.SameSite
-	maxAge       int
-	domain       string
-	path         string
+	securecookie     *securecookie.SecureCookie
+	secureCookieFunc func(r *http.Request) (*securecookie.SecureCookie, error)
+	secureOnly       bool
+	sameSite         http.SameSite
+	maxAge           int
+	domain           string
+	path             string
 }
 
 func NewCookieHandler(hashKey, encryptKey []byte, opts ...CookieHandlerOpt) *CookieHandler {
@@ -30,6 +31,21 @@ func NewCookieHandler(hashKey, encryptKey []byte, opts ...CookieHandlerOpt) *Coo
 	return c
 }
 
+func NewRequestAwareCookieHandler(secureCookieFunc func(r *http.Request) (*securecookie.SecureCookie, error), opts ...CookieHandlerOpt) *CookieHandler {
+	c := &CookieHandler{
+		secureCookieFunc: secureCookieFunc,
+		secureOnly:       true,
+		sameSite:         http.SameSiteLaxMode,
+		path:             "/",
+	}
+
+	for _, opt := range opts {
+		opt(c)
+	}
+
+	return c
+}
+
 type CookieHandlerOpt func(*CookieHandler)
 
 func WithUnsecure() CookieHandlerOpt {
@@ -47,6 +63,10 @@ func WithSameSite(sameSite http.SameSite) CookieHandlerOpt {
 func WithMaxAge(maxAge int) CookieHandlerOpt {
 	return func(c *CookieHandler) {
 		c.maxAge = maxAge
+		if c.IsRequestAware() {
+			return
+		}
+
 		c.securecookie.MaxAge(maxAge)
 	}
 }
@@ -68,8 +88,17 @@ func (c *CookieHandler) CheckCookie(r *http.Request, name string) (string, error
 	if err != nil {
 		return "", err
 	}
+
+	secureCookie := c.securecookie
+	if c.IsRequestAware() {
+		secureCookie, err = c.secureCookieFunc(r)
+		if err != nil {
+			return "", err
+		}
+	}
+
 	var value string
-	if err := c.securecookie.Decode(name, cookie.Value, &value); err != nil {
+	if err := secureCookie.Decode(name, cookie.Value, &value); err != nil {
 		return "", err
 	}
 	return value, nil
@@ -87,6 +116,10 @@ func (c *CookieHandler) CheckQueryCookie(r *http.Request, name string) (string,
 }
 
 func (c *CookieHandler) SetCookie(w http.ResponseWriter, name, value string) error {
+	if c.IsRequestAware() {
+		return errors.New("Cookie handler is request aware")
+	}
+
 	encoded, err := c.securecookie.Encode(name, value)
 	if err != nil {
 		return err
@@ -104,6 +137,35 @@ func (c *CookieHandler) SetCookie(w http.ResponseWriter, name, value string) err
 	return nil
 }
 
+func (c *CookieHandler) SetRequestAwareCookie(r *http.Request, w http.ResponseWriter, name string, value string) error {
+	if !c.IsRequestAware() {
+		return errors.New("Cookie handler is not request aware")
+	}
+
+	secureCookie, err := c.secureCookieFunc(r)
+	if err != nil {
+		return err
+	}
+
+	encoded, err := secureCookie.Encode(name, value)
+	if err != nil {
+		return err
+	}
+
+	http.SetCookie(w, &http.Cookie{
+		Name:     name,
+		Value:    encoded,
+		Domain:   c.domain,
+		Path:     c.path,
+		MaxAge:   c.maxAge,
+		HttpOnly: true,
+		Secure:   c.secureOnly,
+		SameSite: c.sameSite,
+	})
+
+	return nil
+}
+
 func (c *CookieHandler) DeleteCookie(w http.ResponseWriter, name string) {
 	http.SetCookie(w, &http.Cookie{
 		Name:     name,
@@ -116,3 +178,7 @@ func (c *CookieHandler) DeleteCookie(w http.ResponseWriter, name string) {
 		SameSite: c.sameSite,
 	})
 }
+
+func (c *CookieHandler) IsRequestAware() bool {
+	return c.secureCookieFunc != nil
+}


Test patch:
diff --git a/pkg/client/integration_test.go b/pkg/client/integration_test.go
index 98a9d3a..746f7c1 100644
--- a/pkg/client/integration_test.go
+++ b/pkg/client/integration_test.go
@@ -3,6 +3,9 @@ package client_test
 import (
 	"bytes"
 	"context"
+	"crypto/hmac"
+	cryptoRand "crypto/rand"
+	"crypto/sha512"
 	"fmt"
 	"io"
 	"log/slog"
@@ -18,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
+	"github.com/gorilla/securecookie"
 	"github.com/jeremija/gosubmit"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -42,6 +47,15 @@ var Logger = slog.New(
 
 var CTX context.Context
 
+type cookieSpec struct {
+	cookieHandler *httphelper.CookieHandler
+	extraCookies  []*http.Cookie
+}
+
+var defaultCookieSpec = cookieSpec{
+	cookieHandler: httphelper.NewCookieHandler([]byte("test1234test1234"), []byte("test1234test1234"), httphelper.WithUnsecure()),
+}
+
 func TestMain(m *testing.M) {
 	os.Exit(func() int {
 		ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGINT)
@@ -53,14 +67,65 @@ func TestMain(m *testing.M) {
 }
 
 func TestRelyingPartySession(t *testing.T) {
-	for _, wrapServer := range []bool{false, true} {
-		t.Run(fmt.Sprint("wrapServer ", wrapServer), func(t *testing.T) {
-			testRelyingPartySession(t, wrapServer)
-		})
+	secret := make([]byte, 64)
+	_, _ = cryptoRand.Read(secret)
+	hashFunc := sha512.New
+	requestAwareCookieHandler := httphelper.NewRequestAwareCookieHandler(func(r *http.Request) (*securecookie.SecureCookie, error) {
+		// Expect a login ID cookie to be set.
+		loginIDCookie, err := r.Cookie("login_id")
+		require.NoError(t, err)
+
+		// Login ID cookie will contain a UUID.
+		loginID, err := uuid.Parse(loginIDCookie.Value)
+		require.NoError(t, err)
+
+		// Use a HMAC hash of the secret and login ID (salt).
+		h := hmac.New(hashFunc, secret)
+		_, err = h.Write(loginID[:])
+		require.NoError(t, err)
+
+		// Write a separator for the hash key.
+		_, err = h.Write([]byte("INTEGRITY"))
+		require.NoError(t, err)
+		hash := h.Sum(nil) // 64 bytes
+
+		// Write a separator for the block key.
+		_, err = h.Write([]byte("ENCRYPTION"))
+		require.NoError(t, err)
+		block := h.Sum(nil)[:32] // Sum is 64 bytes but we only want 32 for AES-256.
+
+		return securecookie.New(hash, block), nil
+	}, httphelper.WithUnsecure())
+
+	loginID := uuid.New()
+	loginIDCookie := &http.Cookie{
+		Name:     "login_id",
+		Value:    loginID.String(),
+		Secure:   false,
+		SameSite: http.SameSiteLaxMode,
+		Path:     "/",
+	}
+
+	cookieCases := []cookieSpec{
+		defaultCookieSpec,
+		{
+			cookieHandler: requestAwareCookieHandler,
+			extraCookies: []*http.Cookie{
+				loginIDCookie,
+			},
+		},
+	}
+
+	for _, cookieCase := range cookieCases {
+		for _, wrapServer := range []bool{false, true} {
+			t.Run(fmt.Sprint("wrapServer ", wrapServer, " requestAwareCookieHandler ", cookieCase.cookieHandler.IsRequestAware()), func(t *testing.T) {
+				testRelyingPartySession(t, wrapServer, cookieCase)
+			})
+		}
 	}
 }
 
-func testRelyingPartySession(t *testing.T, wrapServer bool) {
+func testRelyingPartySession(t *testing.T, wrapServer bool, cookieSpec cookieSpec) {
 	t.Log("------- start example OP ------")
 	targetURL := "http://local-site"
 	exampleStorage := storage.NewStorage(storage.NewUserStore(targetURL))
@@ -74,7 +139,7 @@ func testRelyingPartySession(t *testing.T, wrapServer bool) {
 	clientID := t.Name() + "-" + strconv.FormatInt(seed.Int63(), 25)
 
 	t.Log("------- run authorization code flow ------")
-	provider, tokens := RunAuthorizationCodeFlow(t, opServer, clientID, "secret")
+	provider, tokens := RunAuthorizationCodeFlow(t, opServer, clientID, "secret", cookieSpec)
 
 	t.Log("------- refresh tokens  ------")
 
@@ -220,7 +285,7 @@ func testResourceServerTokenExchange(t *testing.T, wrapServer bool) {
 	clientSecret := "secret"
 
 	t.Log("------- run authorization code flow ------")
-	provider, tokens := RunAuthorizationCodeFlow(t, opServer, clientID, clientSecret)
+	provider, tokens := RunAuthorizationCodeFlow(t, opServer, clientID, clientSecret, defaultCookieSpec)
 
 	resourceServer, err := rs.NewResourceServerClientCredentials(CTX, opServer.URL, clientID, clientSecret)
 	require.NoError(t, err, "new resource server")
@@ -275,7 +340,7 @@ func testResourceServerTokenExchange(t *testing.T, wrapServer bool) {
 	require.Nil(t, tokenExchangeResponse, "token exchange response")
 }
 
-func RunAuthorizationCodeFlow(t *testing.T, opServer *httptest.Server, clientID, clientSecret string) (provider rp.RelyingParty, tokens *oidc.Tokens[*oidc.IDTokenClaims]) {
+func RunAuthorizationCodeFlow(t *testing.T, opServer *httptest.Server, clientID, clientSecret string, cookieSpec cookieSpec) (provider rp.RelyingParty, tokens *oidc.Tokens[*oidc.IDTokenClaims]) {
 	targetURL := "http://local-site"
 	localURL, err := url.Parse(targetURL + "/login?requestID=1234")
 	require.NoError(t, err, "local url")
@@ -294,8 +359,6 @@ func RunAuthorizationCodeFlow(t *testing.T, opServer *httptest.Server, clientID,
 	}
 
 	t.Log("------- create RP ------")
-	key := []byte("test1234test1234")
-	cookieHandler := httphelper.NewCookieHandler(key, key, httphelper.WithUnsecure())
 	provider, err = rp.NewRelyingPartyOIDC(
 		CTX,
 		opServer.URL,
@@ -303,7 +366,7 @@ func RunAuthorizationCodeFlow(t *testing.T, opServer *httptest.Server, clientID,
 		clientSecret,
 		targetURL,
 		[]string{"openid", "email", "profile", "offline_access"},
-		rp.WithPKCE(cookieHandler),
+		rp.WithPKCE(cookieSpec.cookieHandler),
 		rp.WithAuthStyle(oauth2.AuthStyleInHeader),
 		rp.WithVerifierOpts(
 			rp.WithIssuedAtOffset(5*time.Second),
@@ -317,6 +380,11 @@ func RunAuthorizationCodeFlow(t *testing.T, opServer *httptest.Server, clientID,
 	state := "state-" + strconv.FormatInt(seed.Int63(), 25)
 	capturedW := httptest.NewRecorder()
 	get := httptest.NewRequest("GET", localURL.String(), nil)
+	for _, cookie := range cookieSpec.extraCookies {
+		get.AddCookie(cookie)
+		http.SetCookie(capturedW, cookie)
+	}
+
 	rp.AuthURLHandler(func() string { return state }, provider,
 		rp.WithPromptURLParam("Hello, World!", "Goodbye, World!"),
 		rp.WithURLParam("custom", "param"),