Cookiejar implementation.

client.go

@@ -163,6 +163,11 @@ type Client struct {
 	// User-Agent header to be excluded from the Request.
 	NoDefaultUserAgentHeader bool
 
+	// CookieJar stores cookies allowing user to handle cookies easily.
+	//
+	// If CookieJar is nil no cookie will be collected.
+	CookieJar *CookieJar
+
 	// Callback for establishing new connections to hosts.
 	//
 	// Default Dial is used if not set.
@@ -410,6 +415,7 @@ func (c *Client) Do(req *Request, resp *Response) error {
 			Name:                          c.Name,
 			NoDefaultUserAgentHeader:      c.NoDefaultUserAgentHeader,
 			Dial:                          c.Dial,
+			CookieJar:                     c.CookieJar,
 			DialDualStack:                 c.DialDualStack,
 			IsTLS:                         isTLS,
 			TLSConfig:                     c.TLSConfig,
@@ -525,6 +531,10 @@ type HostClient struct {
 	// Default Dial is used if not set.
 	Dial DialFunc
 
+	// CookieJar stores cookies. If CookieJar is nil
+	// no cookie will be collected.
+	CookieJar *CookieJar
+
 	// Attempt to connect to both ipv4 and ipv6 host addresses
 	// if set to true.
 	//
@@ -1101,8 +1111,17 @@ func (c *HostClient) do(req *Request, resp *Response) (bool, error) {
 		resp = AcquireResponse()
 	}
 
+	host := b2s(req.Host())
+	if c.CookieJar != nil {
+		c.CookieJar.dumpTo(host, req)
+	}
+
 	ok, err := c.doNonNilReqResp(req, resp)
 
+	if c.CookieJar != nil {
+		c.CookieJar.getFrom(host, resp)
+	}
+
 	if nilResp {
 		ReleaseResponse(resp)
 	}

client_test.go

@@ -50,6 +50,44 @@ func TestClientPostArgs(t *testing.T) {
 	}
 }
 
+func TestClientCookieJar(t *testing.T) {
+	ln := fasthttputil.NewInmemoryListener()
+	key, value := "cometo", "thefasthttpcon"
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			cookie := AcquireCookie()
+			cookie.SetKey(key)
+			cookie.SetValue(value)
+			ctx.Response.Header.SetCookie(cookie)
+		},
+	}
+	go s.Serve(ln)
+	c := &Client{
+		Dial: func(addr string) (net.Conn, error) {
+			return ln.Dial()
+		},
+		CookieJar: &CookieJar{},
+	}
+	req := AcquireRequest()
+	res := AcquireResponse()
+	req.SetRequestURI("http://fasthttp.con/hello/world")
+	err := c.Do(req, res)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cs := c.CookieJar.Get("fasthttp.con")
+	if len(cs) == 0 {
+		t.Fatalf("Unexpected len: %d", len(cs))
+	}
+
+	if k := string(cs[0].Key()); k != key {
+		t.Fatalf("Unexpected key: %s <> %s", k, key)
+	}
+	if v := string(cs[0].Value()); v != value {
+		t.Fatalf("Unexpected value: %s <> %s", v, value)
+	}
+}
+
 func TestClientRedirectSameSchema(t *testing.T) {
 
 	listenHTTPS1 := testClientRedirectListener(t, true)

cookie.go

@@ -21,6 +21,7 @@ var (
 // CookieSameSite is an enum for the mode in which the SameSite flag should be set for the given cookie.
 // See https://tools.ietf.org/html/draft-ietf-httpbis-cookie-same-site-00 for details.
 type CookieSameSite int
+
 const (
 	// CookieSameSiteDisabled removes the SameSite flag
 	CookieSameSiteDisabled CookieSameSite = iota

cookiejar.go

@@ -0,0 +1,80 @@
+package fasthttp
+
+import (
+	"sync"
+)
+
+// CookieJar manages cookie storage
+type CookieJar struct {
+	m           sync.RWMutex
+	hostCookies map[string][]*Cookie
+}
+
+// Get returns the cookies stored from a specific domain.
+//
+// If there were no cookies related with host returned slice will be nil.
+func (cj *CookieJar) Get(host string) (cookies []*Cookie) {
+	cj.m.RLock()
+	{
+		if cj.hostCookies == nil {
+			cj.hostCookies = make(map[string][]*Cookie)
+		}
+		cookies = cj.hostCookies[host]
+	}
+	cj.m.RUnlock()
+	return
+}
+
+// Set sets cookies for a specific host.
+func (cj *CookieJar) Set(host string, cookies ...*Cookie) {
+	cj.m.Lock()
+	{
+		if cj.hostCookies == nil {
+			cj.hostCookies = make(map[string][]*Cookie)
+		}
+		hc := cj.hostCookies[host]
+		hc = append(hc, cookies...)
+		cj.hostCookies[host] = hc
+	}
+	cj.m.Unlock()
+}
+
+func (cj *CookieJar) dumpTo(host string, req *Request) {
+	cj.m.Lock()
+	{
+		if cj.hostCookies == nil {
+			cj.hostCookies = make(map[string][]*Cookie)
+		}
+		cookies, ok := cj.hostCookies[host]
+		if ok {
+			for _, cookie := range cookies {
+				req.Header.SetCookieBytesKV(cookie.Key(), cookie.Value())
+			}
+		}
+	}
+	cj.m.Unlock()
+}
+
+func (cj *CookieJar) getFrom(host string, res *Response) {
+	cj.m.Lock()
+	{
+		if cj.hostCookies == nil {
+			cj.hostCookies = make(map[string][]*Cookie)
+		}
+		cookies, ok := cj.hostCookies[host]
+		if !ok {
+			cookies = make([]*Cookie, 0)
+		}
+		res.Header.VisitAllCookie(func(key, value []byte) {
+			cookie := &Cookie{}
+			cookie.SetKeyBytes(key)
+			if !res.Header.Cookie(cookie) {
+				// TODO: error?
+				cookie.SetValueBytes(value)
+			}
+			cookies = append(cookies, cookie)
+		})
+		cj.hostCookies[host] = cookies
+	}
+	cj.m.Unlock()
+}

cookiejar_test.go

@@ -0,0 +1,50 @@
+package fasthttp
+
+import "testing"
+
+func prepareJar() *CookieJar {
+	cj := &CookieJar{}
+	cj.hostCookies = make(map[string][]*Cookie)
+	return cj
+}
+
+func checkKeyValue(t *testing.T, cj *CookieJar, cookie *Cookie, host string, n int) {
+	cs := cj.Get(host)
+	if len(cs) < n {
+		t.Fatalf("Unexpected cookie length: %d. Expected %d", len(cs), n)
+	}
+	c := cs[n-1]
+	if c == nil {
+		t.Fatal("got a nil cookie")
+	}
+	if string(c.Key()) != string(cookie.Key()) {
+		t.Fatalf("key mismatch: %s <> %s", c.Key(), cookie.Key())
+	}
+	if string(c.Value()) != string(cookie.Value()) {
+		t.Fatalf("value mismatch: %s <> %s", c.Value(), cookie.Value())
+	}
+}
+
+func TestCookieJar_Get(t *testing.T) {
+	host := "fast.http"
+	cj := prepareJar()
+
+	cookie := &Cookie{}
+	cookie.SetKey("k")
+	cookie.SetValue("v")
+
+	cj.hostCookies[host] = append(cj.hostCookies[host], cookie)
+	checkKeyValue(t, cj, cookie, host, 1)
+}
+
+func TestCookieJar_Set(t *testing.T) {
+	host := "fast.http"
+	cj := &CookieJar{}
+
+	cookie := &Cookie{}
+	cookie.SetKey("k")
+	cookie.SetValue("v")
+
+	cj.Set(host, cookie)
+	checkKeyValue(t, cj, cookie, host, 1)
+}