-
Notifications
You must be signed in to change notification settings - Fork 1
/
cookie.go
69 lines (59 loc) · 1.54 KB
/
cookie.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
package youcrawl
import (
"net/http"
"net/http/cookiejar"
"sync"
)
type CookieStore interface {
GetCookie(key string) *cookiejar.Jar
SetCookie(key string, jar *cookiejar.Jar)
GetOrCreate(key string) *cookiejar.Jar
}
type DefaultCookieStore struct {
sync.Map
}
func (s *DefaultCookieStore) GetCookie(key string) *cookiejar.Jar {
rawJar, _ := s.Load(key)
jar := rawJar.(*cookiejar.Jar)
return jar
}
func (s *DefaultCookieStore) GetOrCreate(key string) *cookiejar.Jar {
newJar, _ := cookiejar.New(&cookiejar.Options{})
rawJar, _ := s.Map.LoadOrStore(key, newJar)
jar := rawJar.(*cookiejar.Jar)
return jar
}
func (s *DefaultCookieStore) SetCookie(key string, jar *cookiejar.Jar) {
s.Store(key, jar)
}
type CookieMiddleware struct {
Store CookieStore
GetKey func(c *http.Client, r *http.Request, ctx *Context) string
}
func (m *CookieMiddleware) RequestCallback(c *http.Client, r *http.Request, ctx *Context) {
key := "Default"
if m.GetKey != nil {
key = m.GetKey(c, r, ctx)
}
jar := c.Jar.(*cookiejar.Jar)
m.Store.SetCookie(key, jar)
ctx.Cookie = jar
}
func (m *CookieMiddleware) Process(c *http.Client, r *http.Request, ctx *Context) {
key := "Default"
if m.GetKey != nil {
key = m.GetKey(c, r, ctx)
}
jar := m.Store.GetOrCreate(key)
c.Jar = jar
}
type CookieMiddlewareOption struct {
GetKey func(c *http.Client, r *http.Request, ctx *Context) string
}
func NewCookieMiddleware(option CookieMiddlewareOption) *CookieMiddleware {
store := DefaultCookieStore{}
return &CookieMiddleware{
Store: &store,
GetKey: option.GetKey,
}
}