add bodyreader fn

master
兔子 3 years ago
parent 66c8abbcea
commit ce3ebbbf8a

@ -2,6 +2,7 @@ package starnet
import ( import (
"bytes" "bytes"
"context"
"crypto/rand" "crypto/rand"
"errors" "errors"
"fmt" "fmt"
@ -46,6 +47,7 @@ type Request struct {
CircleBuffer *stario.StarBuffer CircleBuffer *stario.StarBuffer
Proxy string Proxy string
Process func(float64) Process func(float64)
respReader io.ReadCloser
} }
func NewRequests(url string, postdata []byte, method string) Request { func NewRequests(url string, postdata []byte, method string) Request {
@ -61,7 +63,7 @@ func NewRequests(url string, postdata []byte, method string) Request {
if strings.ToUpper(method) == "POST" { if strings.ToUpper(method) == "POST" {
req.ReqHeader.Set("Content-Type", HEADER_FORM_URLENCODE) req.ReqHeader.Set("Content-Type", HEADER_FORM_URLENCODE)
} }
req.ReqHeader.Set("User-Agent", "B612 / 1.0.0") req.ReqHeader.Set("User-Agent", "B612 / 1.1.0")
return req return req
} }
@ -144,7 +146,7 @@ func Curl(curl Request) (resps Request, err error) {
if err != nil { if err != nil {
return Request{}, err return Request{}, err
} }
defer resp.Body.Close()
curl.PostBuffer = nil curl.PostBuffer = nil
curl.CircleBuffer = nil curl.CircleBuffer = nil
curl.RespHttpCode = resp.StatusCode curl.RespHttpCode = resp.StatusCode
@ -182,6 +184,8 @@ func Curl(curl Request) (resps Request, err error) {
return return
} }
curl.RecvData = buf.Bytes() curl.RecvData = buf.Bytes()
} else {
curl.respReader = resp.Body
} }
if curl.RecvIo != nil { if curl.RecvIo != nil {
if curl.WriteRecvData { if curl.WriteRecvData {
@ -196,6 +200,11 @@ func Curl(curl Request) (resps Request, err error) {
return curl, err return curl, err
} }
// RespBodyReader Only works when WriteRecvData set to false
func (curl *Request) RespBodyReader() io.ReadCloser {
return curl.respReader
}
func netcurl(curl Request) (*http.Response, error) { func netcurl(curl Request) (*http.Response, error) {
var req *http.Request var req *http.Request
var err error var err error
@ -219,14 +228,13 @@ func netcurl(curl Request) (*http.Response, error) {
} }
} }
transport := &http.Transport{ transport := &http.Transport{
Dial: func(netw, addr string) (net.Conn, error) { DialContext: func(ctx context.Context, netw, addr string) (net.Conn, error) {
deadline := time.Now().Add(time.Duration(curl.TimeOut) * time.Second)
c, err := net.DialTimeout(netw, addr, time.Second*time.Duration(curl.DialTimeOut)) c, err := net.DialTimeout(netw, addr, time.Second*time.Duration(curl.DialTimeOut))
if err != nil { if err != nil {
return nil, err return nil, err
} }
if curl.TimeOut != 0 { if curl.TimeOut != 0 {
c.SetDeadline(deadline) c.SetDeadline(time.Now().Add(time.Duration(curl.TimeOut) * time.Second))
} }
return c, nil return c, nil
}, },

Loading…
Cancel
Save