Added response cookies support

1 files changed, 134 insertions, 8 deletions

diff --git a/cookie.go b/cookie.go
index 3dc1f2c..57d964b 100644
--- a/cookie.go
+++ b/cookie.go
@@ -2,9 +2,132 @@ package fasthttp
 
 import (
 	"bytes"
+	"errors"
+	"time"
 )
 
-func appendCookieBytes(dst []byte, cookies []argsKV) []byte {
+// CookieExpireDelete may be set on Cookie.Expire for expiring the given cookie.
+var CookieExpireDelete = time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
+
+// Cookie represents HTTP response cookie.
+type Cookie struct {
+	// Key is cookie name.
+	Key []byte
+
+	// Value is cookie value.
+	Value []byte
+
+	// Expiration time for the cookie.
+	//
+	// Set expiration time to CookieExpireDelete for expiring (deleting)
+	// the cookie on the client.
+	//
+	// By default cookie lifetime is limited by browser session.
+	Expire time.Time
+
+	// Domain for the cookie.
+	//
+	// By default cookie is set to the current domain.
+	Domain []byte
+
+	// Path for the cookie.
+	//
+	// By default cookie is set to the current page only.
+	Path []byte
+
+	bufKV argsKV
+}
+
+var zeroTime time.Time
+
+// Clear clears the cookie.
+func (c *Cookie) Clear() {
+	c.Key = c.Key[:0]
+	c.Value = c.Value[:0]
+	c.Expire = zeroTime
+	c.Domain = c.Domain[:0]
+	c.Path = c.Path[:0]
+}
+
+// AppendBytes appends cookie representation to dst and returns dst
+// (maybe newly allocated).
+func (c *Cookie) AppendBytes(dst []byte) []byte {
+	if len(c.Key) > 0 {
+		dst = appendQuotedArg(dst, c.Key)
+		dst = append(dst, '=')
+	}
+	dst = appendQuotedArg(dst, c.Value)
+
+	if !c.Expire.IsZero() {
+		c.bufKV.value = c.Expire.In(gmtLocation).AppendFormat(c.bufKV.value[:0], time.RFC1123)
+		dst = append(dst, ';', ' ')
+		dst = append(dst, strCookieExpires...)
+		dst = append(dst, '=')
+		dst = append(dst, c.bufKV.value...)
+	}
+	if len(c.Domain) > 0 {
+		dst = appendCookiePart(dst, strCookieDomain, c.Domain)
+	}
+	if len(c.Path) > 0 {
+		dst = appendCookiePart(dst, strCookiePath, c.Path)
+	}
+	return dst
+}
+
+var errNoCookies = errors.New("no cookies found")
+
+// Parse parses Set-Cookie header.
+func (c *Cookie) Parse(src []byte) error {
+	c.Clear()
+
+	var s cookieScanner
+	s.b = src
+
+	kv := &c.bufKV
+	if !s.next(kv, true) {
+		return errNoCookies
+	}
+
+	c.Key = append(c.Key[:0], kv.key...)
+	c.Value = append(c.Value[:0], kv.value...)
+
+	for s.next(kv, false) {
+		if len(kv.key) == 0 && len(kv.value) == 0 {
+			continue
+		}
+		switch {
+		case bytes.Equal(strCookieExpires, kv.key):
+			v := unsafeBytesToStr(kv.value)
+			exptime, err := time.Parse(time.RFC1123, v)
+			if err != nil {
+				return err
+			}
+			c.Expire = exptime
+		case bytes.Equal(strCookieDomain, kv.key):
+			c.Domain = append(c.Domain[:0], kv.value...)
+		case bytes.Equal(strCookiePath, kv.key):
+			c.Path = append(c.Path[:0], kv.value...)
+		}
+	}
+	return nil
+}
+
+func appendCookiePart(dst, key, value []byte) []byte {
+	dst = append(dst, ';', ' ')
+	dst = append(dst, key...)
+	dst = append(dst, '=')
+	return append(dst, value...)
+}
+
+func getCookieKey(dst, src []byte) []byte {
+	n := bytes.IndexByte(src, '=')
+	if n >= 0 {
+		src = src[:n]
+	}
+	return decodeCookieArg(dst[:0], src, true)
+}
+
+func appendRequestCookieBytes(dst []byte, cookies []argsKV) []byte {
 	for i, n := 0, len(cookies); i < n; i++ {
 		kv := &cookies[i]
 		if len(kv.key) > 0 {
@@ -19,10 +142,10 @@ func appendCookieBytes(dst []byte, cookies []argsKV) []byte {
 	return dst
 }
 
-func parseCookies(cookies []argsKV, src []byte, kv *argsKV) []argsKV {
+func parseRequestCookies(cookies []argsKV, src []byte, kv *argsKV) []argsKV {
 	var s cookieScanner
 	s.b = src
-	for s.next(kv) {
+	for s.next(kv, true) {
 		if len(kv.key) > 0 || len(kv.value) > 0 {
 			cookies = setArg(cookies, kv.key, kv.value)
 		}
@@ -34,7 +157,7 @@ type cookieScanner struct {
 	b []byte
 }
 
-func (s *cookieScanner) next(kv *argsKV) bool {
+func (s *cookieScanner) next(kv *argsKV, decode bool) bool {
 	if len(s.b) == 0 {
 		return false
 	}
@@ -52,21 +175,24 @@ func (s *cookieScanner) next(kv *argsKV) bool {
 	n = bytes.IndexByte(b, '=')
 	if n < 0 {
 		kv.key = kv.key[:0]
-		kv.value = decodeCookieArg(kv.value[:0], b)
+		kv.value = decodeCookieArg(kv.value[:0], b, decode)
 		return true
 	}
 
-	kv.key = decodeCookieArg(kv.key[:0], b[:n])
-	kv.value = decodeCookieArg(kv.value[:0], b[n+1:])
+	kv.key = decodeCookieArg(kv.key[:0], b[:n], decode)
+	kv.value = decodeCookieArg(kv.value[:0], b[n+1:], decode)
 	return true
 }
 
-func decodeCookieArg(dst, src []byte) []byte {
+func decodeCookieArg(dst, src []byte, decode bool) []byte {
 	for len(src) > 0 && src[0] == ' ' {
 		src = src[1:]
 	}
 	for len(src) > 0 && src[len(src)-1] == ' ' {
 		src = src[:len(src)-1]
 	}
+	if !decode {
+		return append(dst, src...)
+	}
 	return decodeArg(dst, src, true)
 }