Implement cookies in context, handle set/get headers

Author: 2hmad

context.go

@@ -2,7 +2,9 @@ package routing
 
 import (
 	"bytes"
+	"github.com/gopulse/pulse-router/utils"
 	"github.com/valyala/fasthttp"
+	"time"
 )
 
 type handlerFunc func(ctx *Context) error
@@ -13,6 +15,20 @@ type Context struct {
 	paramValues []string
 	handlers    []handlerFunc
 	handlerIdx  int
+	Cookies     *fasthttp.Cookie
+}
+
+type Cookie struct {
+	Name        string    `json:"name"`
+	Value       string    `json:"value"`
+	Path        string    `json:"path"`
+	Domain      string    `json:"domain"`
+	MaxAge      int       `json:"max_age"`
+	Expires     time.Time `json:"expires"`
+	Secure      bool      `json:"secure"`
+	HTTPOnly    bool      `json:"http_only"`
+	SameSite    string    `json:"same_site"`
+	SessionOnly bool      `json:"session_only"`
 }
 
 // NewContext returns a new Context.
@@ -81,3 +97,63 @@ func (c *Context) Reset() {
 func (c *Context) Abort() {
 	c.handlerIdx = len(c.handlers)
 }
+
+// SetCookie sets a cookie with the given name, value, and options.
+func (c *Context) SetCookie(cookie *Cookie) {
+	acCookie := fasthttp.AcquireCookie()
+
+	acCookie.SetKey(cookie.Name)
+	acCookie.SetValue(cookie.Value)
+	acCookie.SetPath(cookie.Path)
+	acCookie.SetDomain(cookie.Domain)
+	acCookie.SetSecure(cookie.Secure)
+	acCookie.SetHTTPOnly(cookie.HTTPOnly)
+	acCookie.SetSecure(cookie.Secure)
+	if !cookie.SessionOnly {
+		acCookie.SetMaxAge(cookie.MaxAge)
+		acCookie.SetExpire(cookie.Expires)
+	}
+
+	switch utils.ToLower(cookie.SameSite) {
+	case string(rune(fasthttp.CookieSameSiteStrictMode)):
+		acCookie.SetSameSite(fasthttp.CookieSameSiteStrictMode)
+	case string(rune(fasthttp.CookieSameSiteNoneMode)):
+		acCookie.SetSameSite(fasthttp.CookieSameSiteNoneMode)
+	case string(rune(fasthttp.CookieSameSiteDisabled)):
+		acCookie.SetSameSite(fasthttp.CookieSameSiteDisabled)
+	default:
+		acCookie.SetSameSite(fasthttp.CookieSameSiteDefaultMode)
+	}
+
+	c.Ctx.Response.Header.SetCookie(acCookie)
+	fasthttp.ReleaseCookie(acCookie)
+}
+
+// GetCookie returns the value of the cookie with the given name.
+func (c *Context) GetCookie(name string) string {
+	cookie := c.Ctx.Request.Header.Cookie(name)
+	if cookie == nil {
+		return ""
+	}
+	return string(cookie)
+}
+
+// ClearCookie deletes the cookie with the given name.
+func (c *Context) ClearCookie(name string) {
+	c.SetCookie(&Cookie{
+		Name:     name,
+		Value:    "",
+		Path:     "/",
+		MaxAge:   -1,
+		Secure:   false,
+		HTTPOnly: true,
+	})
+}
+
+func (c *Context) SetHeader(key, value string) {
+	c.Ctx.Response.Header.Set(key, value)
+}
+
+func (c *Context) GetHeader(key string) string {
+	return string(c.Ctx.Request.Header.Peek(key))
+}