diff --git a/myhttp/httpc/const.go b/myhttp/httpc/const.go index 197dad0..49ece2f 100644 --- a/myhttp/httpc/const.go +++ b/myhttp/httpc/const.go @@ -1,5 +1,11 @@ package httpc +import ( + "fmt" + "net/http" + "strings" +) + const ( HeaderUserAgent = "User-Agent" HeaderContentType = "Content-Type" @@ -17,3 +23,66 @@ const ( AcceptHtml = "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7" AcceptCNLanguage = "zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7" ) + +func RedirectAllCookies(req *http.Request, via []*http.Request) error { + _, cookieHeader := GetRedirectCookieHeaders(req, via) + if cookieHeader != "" { + req.Header.Set("Cookie", cookieHeader) + } + return nil +} + +func GetRedirectCookieHeaders(req *http.Request, via []*http.Request) (map[string]string, string) { + if len(via) == 0 { + return map[string]string{}, "" + } + + lastReq := via[len(via)-1] + + cookieMap := parseCookieMap(lastReq.Header.Get("Cookie")) + + for _, c := range lastReq.Cookies() { + cookieMap[c.Name] = c.Value + } + + if lastReq.Response != nil { + for _, c := range lastReq.Response.Cookies() { + cookieMap[c.Name] = c.Value + } + } + + var pairs []string + for name, value := range cookieMap { + pairs = append(pairs, fmt.Sprintf("%s=%s", name, value)) + } + + if len(pairs) > 0 { + // 直接 Set 最终合并后的完整字符串 + return cookieMap, strings.Join(pairs, "; ") + } + + return map[string]string{}, "" +} + +func parseCookieMap(str string) map[string]string { + res := make(map[string]string) + if len(str) == 0 { + return res + } + + sp := strings.Split(str, ";") + + for _, v := range sp { + split := strings.SplitN(v, "=", 2) + if len(split) == 2 { + name := strings.TrimSpace(split[0]) + value := strings.TrimSpace(split[1]) + + if name != "" && value != "" { + res[name] = value + } + } + } + + return res +} diff --git a/myhttp/httpc/httpclient.go b/myhttp/httpc/httpclient.go index 10eaf0d..4998422 100644 --- a/myhttp/httpc/httpclient.go +++ b/myhttp/httpc/httpclient.go @@ -5,14 +5,16 @@ import ( "context" "errors" "fmt" - "git.makemake.in/kzkzzzz/mycommon/mylog" - jsoniter "github.com/json-iterator/go" - "golang.org/x/time/rate" "io" "net/http" + "net/http/cookiejar" "net/url" "strings" "time" + + "git.makemake.in/kzkzzzz/mycommon/mylog" + jsoniter "github.com/json-iterator/go" + "golang.org/x/time/rate" ) var ( @@ -64,6 +66,7 @@ type Request struct { httpClient *HttpClient contentType string noWaitQps bool + cookies []*http.Cookie } type HttpClient struct { @@ -79,17 +82,24 @@ func New(opts ...ConfigOpt) *HttpClient { } if config.timeout <= 0 { - config.timeout = time.Second * 3 + config.timeout = time.Second * 10 } if config.transport == nil { config.transport = NewTransport(3072, time.Second*90) } + //if config.redirectFn == nil { + // config.redirectFn = RedirectAllCookies + //} + if config.client == nil { + cookieJar, _ := cookiejar.New(nil) + client := &http.Client{ Transport: config.transport, Timeout: config.timeout, + Jar: cookieJar, } if config.redirectFn != nil { @@ -188,6 +198,18 @@ func (r *Request) SetHeader(k, v string) *Request { return r } +func (r *Request) SetCookies(v ...*http.Cookie) *Request { + r.cookies = append(r.cookies, v...) + return r +} + +func (r *Request) SetCookieMap(cookiesMap map[string]*http.Cookie) *Request { + for _, v := range cookiesMap { + r.cookies = append(r.cookies, v) + } + return r +} + func (r *Request) SetHeaders(headers map[string]string) *Request { for k, v := range headers { r.SetHeader(k, v) @@ -252,6 +274,10 @@ func (r *Request) Do(method, rawUrl string) (*Response, error) { return nil, err } + for _, cookie := range r.cookies { + req.AddCookie(cookie) + } + query := req.URL.Query() for k, v := range r.mapQuery { diff --git a/myhttp/httpc/option.go b/myhttp/httpc/option.go index f09b4f2..d111280 100644 --- a/myhttp/httpc/option.go +++ b/myhttp/httpc/option.go @@ -1,9 +1,10 @@ package httpc import ( - "golang.org/x/time/rate" "net/http" "time" + + "golang.org/x/time/rate" ) type (