From ce3ebbbf8aa15e2a41a1b48041872677ea38f72c Mon Sep 17 00:00:00 2001
From: starainrt <i@b612.me>
Date: Fri, 11 Mar 2022 09:29:09 +0800
Subject: [PATCH] add bodyreader fn

---
 curl.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/curl.go b/curl.go
index fcbc43b..315cca6 100644
--- a/curl.go
+++ b/curl.go
@@ -2,6 +2,7 @@ package starnet
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"errors"
 	"fmt"
@@ -46,6 +47,7 @@ type Request struct {
 	CircleBuffer *stario.StarBuffer
 	Proxy        string
 	Process      func(float64)
+	respReader   io.ReadCloser
 }
 
 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" {
 		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
 }
 
@@ -144,7 +146,7 @@ func Curl(curl Request) (resps Request, err error) {
 	if err != nil {
 		return Request{}, err
 	}
-	defer resp.Body.Close()
+
 	curl.PostBuffer = nil
 	curl.CircleBuffer = nil
 	curl.RespHttpCode = resp.StatusCode
@@ -182,6 +184,8 @@ func Curl(curl Request) (resps Request, err error) {
 			return
 		}
 		curl.RecvData = buf.Bytes()
+	} else {
+		curl.respReader = resp.Body
 	}
 	if curl.RecvIo != nil {
 		if curl.WriteRecvData {
@@ -196,6 +200,11 @@ func Curl(curl Request) (resps Request, err error) {
 	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) {
 	var req *http.Request
 	var err error
@@ -219,14 +228,13 @@ func netcurl(curl Request) (*http.Response, error) {
 		}
 	}
 	transport := &http.Transport{
-		Dial: func(netw, addr string) (net.Conn, error) {
-			deadline := time.Now().Add(time.Duration(curl.TimeOut) * time.Second)
+		DialContext: func(ctx context.Context, netw, addr string) (net.Conn, error) {
 			c, err := net.DialTimeout(netw, addr, time.Second*time.Duration(curl.DialTimeOut))
 			if err != nil {
 				return nil, err
 			}
 			if curl.TimeOut != 0 {
-				c.SetDeadline(deadline)
+				c.SetDeadline(time.Now().Add(time.Duration(curl.TimeOut) * time.Second))
 			}
 			return c, nil
 		},