forked from dgrr/cookiejar
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cookiejar.go
173 lines (149 loc) · 3.78 KB
/
cookiejar.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
package cookiejar
import (
"io"
"sync"
"unsafe"
"github.com/valyala/fasthttp"
)
var cookiePool = sync.Pool{
New: func() interface{} {
return &CookieJar{}
},
}
// AcquireCookieJar returns an empty CookieJar object from pool
func AcquireCookieJar() *CookieJar {
return cookiePool.Get().(*CookieJar)
}
// ReleaseCookieJar returns CookieJar to the pool
func ReleaseCookieJar(c *CookieJar) {
c.Release()
cookiePool.Put(c)
}
// CookieJar is container of cookies
//
// This object is used to handle multiple cookies
type CookieJar map[string]*fasthttp.Cookie
// Set sets cookie using key-value
//
// This function can replace an existent cookie
func (cj *CookieJar) Set(key, value string) {
setCookie(cj, key, value)
}
// Get returns and delete a value from cookiejar.
func (cj *CookieJar) Get() *fasthttp.Cookie {
for k, v := range *cj {
delete(*cj, k)
return v
}
return nil
}
func (cj *CookieJar) Del(k string) {
delete(*cj, k)
}
// SetBytesK sets cookie using key=value
//
// This function can replace an existent cookie.
func (cj *CookieJar) SetBytesK(key []byte, value string) {
setCookie(cj, b2s(key), value)
}
// SetBytesV sets cookie using key=value
//
// This function can replace an existent cookie.
func (cj *CookieJar) SetBytesV(key string, value []byte) {
setCookie(cj, key, b2s(value))
}
// SetBytesKV sets cookie using key=value
//
// This function can replace an existent cookie.
func (cj *CookieJar) SetBytesKV(key, value []byte) {
setCookie(cj, b2s(key), b2s(value))
}
func setCookie(cj *CookieJar, key, value string) {
c, ok := (*cj)[key]
if !ok {
c = fasthttp.AcquireCookie()
}
c.SetKey(key)
c.SetValue(value)
(*cj)[key] = c
}
// SetCookie sets cookie using its key.
//
// After that you can use Peek or Get function to get cookie value.
func (cj *CookieJar) Put(cookie *fasthttp.Cookie) {
c, ok := (*cj)[b2s(cookie.Key())]
if ok {
fasthttp.ReleaseCookie(c)
}
(*cj)[b2s(cookie.Key())] = cookie
}
// Peek peeks cookie value using key.
//
// This function does not delete cookie
func (cj *CookieJar) Peek(key string) *fasthttp.Cookie {
return (*cj)[key]
}
// Release releases all cookie values.
func (cj *CookieJar) Release() {
for k := range *cj {
cj.ReleaseCookie(k)
}
}
// ReleaseCookie releases a cookie specified by parsed key.
func (cj *CookieJar) ReleaseCookie(key string) {
c, ok := (*cj)[key]
if ok {
fasthttp.ReleaseCookie(c)
delete(*cj, key)
}
}
// PeekValue returns value of specified cookie-key.
func (cj *CookieJar) PeekValue(key string) []byte {
c, ok := (*cj)[key]
if ok {
return c.Value()
}
return nil
}
// ReadResponse gets all Response cookies reading Set-Cookie header.
func (cj *CookieJar) ReadResponse(r *fasthttp.Response) {
r.Header.VisitAllCookie(func(key, value []byte) {
cookie := fasthttp.AcquireCookie()
cookie.ParseBytes(value)
cj.Put(cookie)
})
}
// ReadRequest gets all cookies from a Request reading Set-Cookie header.
func (cj *CookieJar) ReadRequest(r *fasthttp.Request) {
r.Header.VisitAllCookie(func(key, value []byte) {
cookie := fasthttp.AcquireCookie()
cookie.ParseBytes(value)
cj.Put(cookie)
})
}
// WriteTo writes all cookies representation to w.
func (cj *CookieJar) WriteTo(w io.Writer) (n int64, err error) {
for _, c := range *cj {
nn, err := c.WriteTo(w)
n += nn
if err != nil {
break
}
}
return
}
// FillRequest dumps all cookies stored in cj into Request adding this values to Cookie header.
func (cj *CookieJar) FillRequest(r *fasthttp.Request) {
for _, c := range *cj {
r.Header.SetCookieBytesKV(c.Key(), c.Value())
}
}
// FillResponse dumps all cookies stored in cj into Response adding this values to Cookie header.
func (cj *CookieJar) FillResponse(r *fasthttp.Response) {
for _, c := range *cj {
r.Header.SetCookie(c)
}
}
func b2s(b []byte) string {
return *(*string)(unsafe.Pointer(&b))
}