1 Star 0 Fork 0

王南北丶/direwolf

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
download.go 3.70 KB
一键复制 编辑 原始数据 按行查看 历史
王南北丶 提交于 2020-01-21 02:52 . fix request context leak
package direwolf
import (
"bytes"
"context"
"errors"
"io"
"io/ioutil"
"net/http"
"strings"
"time"
)
// send is low level request method.
func send(session *Session, req *Request) (*Response, error) {
// Set timeout to request context.
// Default timeout is 30s.
timeout := time.Second * 30
if req.Timeout > 0 {
timeout = time.Second * time.Duration(req.Timeout)
} else if session.Timeout > 0 {
timeout = time.Second * time.Duration(session.Timeout)
}
ctx, timeoutCancel := context.WithTimeout(context.Background(), timeout)
// set proxy to request context.
if req.Proxy != nil {
ctx = context.WithValue(ctx, "http", req.Proxy.HTTP)
ctx = context.WithValue(ctx, "https", req.Proxy.HTTPS)
} else if session.Proxy != nil {
ctx = context.WithValue(ctx, "http", session.Proxy.HTTP)
ctx = context.WithValue(ctx, "https", session.Proxy.HTTPS)
}
// set RedirectNum to request context.
// default RedirectNum is 10.
if req.RedirectNum == 0 {
ctx = context.WithValue(ctx, "redirectNum", 10)
} else if req.RedirectNum > 0 {
ctx = context.WithValue(ctx, "redirectNum", req.RedirectNum)
} else {
ctx = context.WithValue(ctx, "redirectNum", 0)
}
// Make new http.Request with context
httpReq, err := http.NewRequestWithContext(ctx, req.Method, req.URL, nil)
if err != nil {
timeoutCancel()
return nil, WrapErr(err, "build Request error, please check request url or request method")
}
// Handle the Headers.
httpReq.Header = mergeHeaders(req.Headers, session.Headers)
// Handle the DataForm, Body or JsonBody.
// Set right Content-Type.
if req.PostForm != nil {
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
data := req.PostForm.URLEncode()
httpReq.Body = ioutil.NopCloser(strings.NewReader(data))
} else if req.Body != nil {
httpReq.Body = ioutil.NopCloser(bytes.NewReader(req.Body))
} else if req.JsonBody != nil {
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Body = ioutil.NopCloser(bytes.NewReader(req.JsonBody))
}
// Handle Cookies
if req.Cookies != nil {
for _, cookie := range req.Cookies {
httpReq.AddCookie(cookie)
}
}
resp, err := session.client.Do(httpReq) // do request
if err != nil {
if strings.Contains(err.Error(), "context deadline exceeded") { // check timeout error
timeoutCancel()
return nil, WrapErr(ErrTimeout, err.Error())
}
timeoutCancel()
return nil, WrapErr(err, "Request Error")
}
defer func() {
if err := resp.Body.Close(); err != nil {
panic(err)
}
}()
response, err := buildResponse(req, resp)
if err != nil {
timeoutCancel()
return nil, WrapErr(err, "build Response Error")
}
timeoutCancel() // cancel the timeout context after request successed.
return response, nil
}
// buildResponse build response with http.Response after do request.
func buildResponse(httpReq *Request, httpResp *http.Response) (*Response, error) {
content, err := ioutil.ReadAll(httpResp.Body)
if err != nil {
if !errors.Is(err, io.ErrUnexpectedEOF) { // Ignore Unexpected EOF error
return nil, WrapErr(err, "read Response.Body failed")
}
}
return &Response{
URL: httpReq.URL,
StatusCode: httpResp.StatusCode,
Proto: httpResp.Proto,
Headers: httpResp.Header,
Cookies: httpResp.Cookies(),
Request: httpReq,
ContentLength: httpResp.ContentLength,
Content: content,
encoding: "UTF-8",
}, nil
}
// mergeHeaders merge Request headers and Session Headers.
// Request has higher priority.
func mergeHeaders(h1, h2 http.Header) http.Header {
h := http.Header{}
for key, values := range h2 {
for _, value := range values {
h.Set(key, value)
}
}
for key, values := range h1 {
for _, value := range values {
h.Set(key, value)
}
}
return h
}
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Go
1
https://gitee.com/wnanbei/direwolf.git
git@gitee.com:wnanbei/direwolf.git
wnanbei
direwolf
direwolf
master

搜索帮助

0d507c66 1850385 C8b1a773 1850385