From 6184521e56aca759a3e05d00b3ae2374b0b72350 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:16:46 +0100 Subject: [PATCH 01/11] pkg/http: Add `secureCookieFunc` field to CookieHandler. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index 1ebc9e2..922aec0 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 { From ac4c761df5b1e0d8e1ca4bd04fec164ed754a157 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:17:10 +0100 Subject: [PATCH 02/11] pkg/http: Add `IsRequestAware` method CookieHandler. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index 922aec0..cbd9655 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -117,3 +117,7 @@ func (c *CookieHandler) DeleteCookie(w http.ResponseWriter, name string) { SameSite: c.sameSite, }) } + +func (c *CookieHandler) IsRequestAware() bool { + return c.secureCookieFunc != nil +} From eae9da5e4c022003fc7439f6e3221381c34d8364 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:18:32 +0100 Subject: [PATCH 03/11] pkg/http: Use `secureCookieFunc` when checking a cookie (if set). Signed-off-by: Mark Laing --- pkg/http/cookie.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index cbd9655..661d81a 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -69,8 +69,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 From 05e528cff0d5145555434e8481daf739e1fe79dd Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:19:06 +0100 Subject: [PATCH 04/11] pkg/http: Error on `SetCookie` if cookie handler is request aware. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index 661d81a..b200300 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -97,6 +97,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 From 0a5b6fadbec73d614cd83f03c54e6bd79396aca6 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:19:40 +0100 Subject: [PATCH 05/11] pkg/http: Add method to set request aware cookies. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index b200300..9fe1531 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -118,6 +118,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, From a96fc0498cf46a830ebb39570340ef05da6053cc Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:19:58 +0100 Subject: [PATCH 06/11] pkg/http: Add function to create a new request aware cookie handler. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index 9fe1531..6b87d3d 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -31,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 { From 36fa09bbb69cc94b4f96e2978ace22687b40d8fa Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:22:12 +0100 Subject: [PATCH 07/11] pkg/client/rp: Update `trySetStateCookie` function signature. Use `SetRequestAwareCookie` if the cookie handle is request aware. This function signature can be updated because it is not exported. Signed-off-by: Mark Laing --- pkg/client/rp/relying_party.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/pkg/client/rp/relying_party.go b/pkg/client/rp/relying_party.go index e6fa078..9efc893 100644 --- a/pkg/client/rp/relying_party.go +++ b/pkg/client/rp/relying_party.go @@ -410,7 +410,7 @@ 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 } @@ -607,9 +607,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 } } From 5cad5e7c9d1db63e57570a8387b14d295965e055 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:24:21 +0100 Subject: [PATCH 08/11] pkg/client/rp: Add `GenerateAndStoreCodeChallengeWithRequest` function. It's not possible to add a `http.Request` argument to `GenerateAndStoreCodeChallenge` as this would be a breaking change. Instead, add a new function that accepts a request argument and call `SetRequestAwareCookie` here. Signed-off-by: Mark Laing --- pkg/client/rp/relying_party.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/pkg/client/rp/relying_party.go b/pkg/client/rp/relying_party.go index 9efc893..3cd8db1 100644 --- a/pkg/client/rp/relying_party.go +++ b/pkg/client/rp/relying_party.go @@ -436,6 +436,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") From 8e86a8e01e09aeb714ee86583880b7a7c4276be0 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Wed, 4 Jun 2025 10:26:17 +0100 Subject: [PATCH 09/11] pkg/client/rp: Update PKCE logic to pass request if required by cookie handler. Signed-off-by: Mark Laing --- pkg/client/rp/relying_party.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pkg/client/rp/relying_party.go b/pkg/client/rp/relying_party.go index 3cd8db1..93c5dd0 100644 --- a/pkg/client/rp/relying_party.go +++ b/pkg/client/rp/relying_party.go @@ -415,7 +415,14 @@ func AuthURLHandler(stateFn func() string, rp RelyingParty, urlParam ...URLParam 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 From 4c7a5c570002c7076e0176847072d57751bdcb12 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Fri, 6 Jun 2025 09:44:02 +0100 Subject: [PATCH 10/11] pkg/http: Don't set MaxAge if cookie handler is request aware. The securecookie field can be nil. Expect the caller to set max age on the securecookie returned by the secureCookieFunc. Signed-off-by: Mark Laing --- pkg/http/cookie.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkg/http/cookie.go b/pkg/http/cookie.go index 6b87d3d..7724aeb 100644 --- a/pkg/http/cookie.go +++ b/pkg/http/cookie.go @@ -63,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) } } From 6e7ee79a688ec79f420a13ed15aee10f15dc9bb4 Mon Sep 17 00:00:00 2001 From: Mark Laing Date: Fri, 6 Jun 2025 09:45:46 +0100 Subject: [PATCH 11/11] pkg/client: Add integration tests for request aware cookie handling. Adds a new type `cookieSpec` which is accepted as an argument to `RunAuthorizationCodeFlow`. `TestRelyingPartySession` now runs with `wrapServer` true/false and with two cookie handlers, one static and one request aware. The request aware handler extracts encryption keys from a secret using a salt from a "login_id" cookie. Signed-off-by: Mark Laing --- pkg/client/integration_test.go | 82 +++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 11 deletions(-) diff --git a/pkg/client/integration_test.go b/pkg/client/integration_test.go index 98a9d3a..fe20492 100644 --- a/pkg/client/integration_test.go +++ b/pkg/client/integration_test.go @@ -3,6 +3,9 @@ package client_test import ( "bytes" "context" + "crypto/hkdf" + cryptoRand "crypto/rand" + "crypto/sha256" "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,57 @@ 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, 512) + _, _ = cryptoRand.Read(secret) + hashFunc := sha256.New + requestAwareCookieHandler := httphelper.NewRequestAwareCookieHandler(func(r *http.Request) (*securecookie.SecureCookie, error) { + loginIDCookie, err := r.Cookie("login_id") + require.NoError(t, err) + + loginID, err := uuid.Parse(loginIDCookie.Value) + require.NoError(t, err) + + prk, err := hkdf.Extract(hashFunc, secret, loginID[:]) + require.NoError(t, err) + + hash, err := hkdf.Expand(hashFunc, prk, "INTEGRITY", 64) + require.NoError(t, err) + + block, err := hkdf.Expand(hashFunc, prk, "ENCRYPTION", 32) + require.NoError(t, err) + + 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 +131,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 +277,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 +332,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 +351,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 +358,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 +372,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"),