Merge remote-tracking branch 'og/main'

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