feat: Request aware cookie handling (#753)

* pkg/http: Add `secureCookieFunc` field to CookieHandler.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* pkg/http: Add `IsRequestAware` method CookieHandler.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* pkg/http: Use `secureCookieFunc` when checking a cookie (if set).

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* pkg/http: Error on `SetCookie` if cookie handler is request aware.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* pkg/http: Add method to set request aware cookies.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* pkg/http: Add function to create a new request aware cookie handler.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* 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 <mark.laing@canonical.com>

* 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 <mark.laing@canonical.com>

* pkg/client/rp: Update PKCE logic to pass request if required by cookie handler.

Signed-off-by: Mark Laing <mark.laing@canonical.com>

* 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 <mark.laing@canonical.com>

* 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 <mark.laing@canonical.com>

---------

Signed-off-by: Mark Laing <mark.laing@canonical.com>

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
+}