github.com/gospider007/requests@v0.0.0-20240506025355-c73d46169a23/client.go (about)

     1  package requests
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net/url"
     8  
     9  	"net/http"
    10  
    11  	"github.com/gospider007/gtls"
    12  )
    13  
    14  // Connection Management
    15  type Client struct {
    16  	option    ClientOption
    17  	transport *roundTripper
    18  	ctx       context.Context
    19  	cnl       context.CancelFunc
    20  	closed    bool
    21  }
    22  
    23  var defaultClient, _ = NewClient(context.TODO())
    24  
    25  // New Connection Management
    26  func NewClient(preCtx context.Context, options ...ClientOption) (*Client, error) {
    27  	if preCtx == nil {
    28  		preCtx = context.TODO()
    29  	}
    30  	var option ClientOption
    31  	if len(options) > 0 {
    32  		option = options[0]
    33  	}
    34  	result := new(Client)
    35  	result.ctx, result.cnl = context.WithCancel(preCtx)
    36  	result.transport = newRoundTripper(result.ctx, option)
    37  	result.option = option
    38  	//cookiesjar
    39  	if !result.option.DisCookie {
    40  		if result.option.Jar == nil {
    41  			result.option.Jar = NewJar()
    42  		}
    43  	}
    44  	var err error
    45  	if result.option.Proxy != "" {
    46  		_, err = gtls.VerifyProxy(result.option.Proxy)
    47  	}
    48  	return result, err
    49  }
    50  
    51  // Modifying the client's proxy
    52  func (obj *Client) SetProxy(proxyUrl string) (err error) {
    53  	_, err = gtls.VerifyProxy(proxyUrl)
    54  	if err == nil {
    55  		obj.option.Proxy = proxyUrl
    56  	}
    57  	return
    58  }
    59  
    60  // Modify the proxy method of the client
    61  func (obj *Client) SetGetProxy(getProxy func(ctx context.Context, url *url.URL) (string, error)) {
    62  	obj.transport.setGetProxy(getProxy)
    63  }
    64  
    65  // Close idle connections. If the connection is in use, wait until it ends before closing
    66  func (obj *Client) CloseConns() {
    67  	obj.transport.closeConns()
    68  }
    69  
    70  // Close the connection, even if it is in use, it will be closed
    71  func (obj *Client) ForceCloseConns() {
    72  	obj.transport.forceCloseConns()
    73  }
    74  
    75  // Close the client and cannot be used again after shutdown
    76  func (obj *Client) Close() {
    77  	obj.closed = true
    78  	obj.ForceCloseConns()
    79  	obj.cnl()
    80  }
    81  
    82  // func checkRedirect(req *http.Request, via []*http.Request) error {
    83  // 	ctxData := GetReqCtxData(req.Context())
    84  // 	if ctxData.maxRedirect == 0 || ctxData.maxRedirect >= len(via) {
    85  // 		return nil
    86  // 	}
    87  // 	return http.ErrUseLastResponse
    88  // }
    89  
    90  func (obj *Client) do(req *http.Request, option *RequestOption) (resp *http.Response, err error) {
    91  	var redirectNum int
    92  	for {
    93  		redirectNum++
    94  		resp, err = obj.send(req, option)
    95  		if req.Body != nil {
    96  			req.Body.Close()
    97  		}
    98  		if err != nil {
    99  			return
   100  		}
   101  		if option.MaxRedirect < 0 { //dis redirect
   102  			return
   103  		}
   104  		if option.MaxRedirect > 0 && redirectNum > option.MaxRedirect {
   105  			return
   106  		}
   107  		loc := resp.Header.Get("Location")
   108  		if loc == "" {
   109  			return resp, nil
   110  		}
   111  		u, err := req.URL.Parse(loc)
   112  		if err != nil {
   113  			return resp, fmt.Errorf("failed to parse Location header %q: %v", loc, err)
   114  		}
   115  		ireq, err := NewRequestWithContext(req.Context(), http.MethodGet, u, nil)
   116  		if err != nil {
   117  			return resp, err
   118  		}
   119  		var shouldRedirect bool
   120  		ireq.Method, shouldRedirect, _ = redirectBehavior(req.Method, resp, ireq)
   121  		if !shouldRedirect {
   122  			return resp, nil
   123  		}
   124  		ireq.Response = resp
   125  		ireq.Header = defaultHeaders()
   126  		ireq.Header.Set("Referer", req.URL.String())
   127  		for key := range ireq.Header {
   128  			if val := req.Header.Get(key); val != "" {
   129  				ireq.Header.Set(key, val)
   130  			}
   131  		}
   132  		if getDomain(u) == getDomain(req.URL) {
   133  			if Authorization := req.Header.Get("Authorization"); Authorization != "" {
   134  				ireq.Header.Set("Authorization", Authorization)
   135  			}
   136  			cookies := Cookies(req.Cookies()).String()
   137  			if cookies != "" {
   138  				ireq.Header.Set("Cookie", cookies)
   139  			}
   140  			addCookie(ireq, resp.Cookies())
   141  		}
   142  		io.Copy(io.Discard, resp.Body)
   143  		resp.Body.Close()
   144  		req = ireq
   145  	}
   146  }
   147  func (obj *Client) send(req *http.Request, option *RequestOption) (resp *http.Response, err error) {
   148  	if option.Jar != nil {
   149  		addCookie(req, option.Jar.GetCookies(req.URL))
   150  	}
   151  	resp, err = obj.transport.RoundTrip(req)
   152  	if err != nil {
   153  		return nil, err
   154  	}
   155  	if option.Jar != nil {
   156  		if rc := resp.Cookies(); len(rc) > 0 {
   157  			option.Jar.SetCookies(req.URL, rc)
   158  		}
   159  	}
   160  	return resp, nil
   161  }