mycommon/myhttp/httpc/httpclient.go

package httpc

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"git.makemake.in/kzkzzzz/mycommon/mylog"
	jsoniter "github.com/json-iterator/go"
	"golang.org/x/time/rate"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var (
	defaultClient *HttpClient
)

func init() {
	defaultClient = New()
}

func ReInitDefault(timeout time.Duration) {
	defaultClient = New(WithTimout(timeout))
}

func ReInitDefaultOpt(opts ...ConfigOpt) {
	defaultClient = New(opts...)
}

func Client() *HttpClient {
	return defaultClient
}

func NewTransport(maxConn int, idleTimeout time.Duration) *http.Transport {
	if maxConn <= 0 {
		panic("max connection <= 0")
	}

	if idleTimeout <= 0 {
		panic("idle timeout <= 0")
	}

	tr := http.DefaultTransport.(*http.Transport).Clone()
	tr.MaxIdleConns = 0
	tr.MaxConnsPerHost = 0
	tr.MaxIdleConnsPerHost = maxConn
	tr.IdleConnTimeout = idleTimeout
	tr.DisableKeepAlives = false
	//tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: false}
	return tr
}

type Request struct {
	ctx         context.Context
	header      http.Header
	body        any
	mapQuery    map[string]string
	urlQuery    url.Values
	httpClient  *HttpClient
	contentType string
	noWaitQps   bool
}

type HttpClient struct {
	config     *Config
	client     *http.Client
	qpsLimiter *rate.Limiter
}

func New(opts ...ConfigOpt) *HttpClient {
	config := &Config{}
	for _, opt := range opts {
		opt(config)
	}

	if config.timeout <= 0 {
		config.timeout = time.Second * 3
	}

	if config.transport == nil {
		config.transport = NewTransport(3072, time.Second*90)
	}

	if config.client == nil {
		client := &http.Client{
			Transport: config.transport,
			Timeout:   config.timeout,
		}

		if config.redirectFn != nil {
			client.CheckRedirect = config.redirectFn
		}

		config.client = client
	}

	hc := &HttpClient{
		config: config,
		client: config.client,
	}

	if config.qpsLimiter == nil {
		hc.qpsLimiter = config.qpsLimiter
	} else if config.qps > 0 {
		hc.qpsLimiter = rate.NewLimiter(rate.Every(time.Second/time.Duration(config.qps)), config.qps)
	}

	return hc
}

func (h *HttpClient) RawClient() *http.Client {
	return h.client
}

func (h *HttpClient) SetProxy(pr string) *HttpClient {
	u, err := url.Parse(pr)
	if err != nil {
		panic(err)
	}

	h.client.Transport.(*http.Transport).Proxy = http.ProxyURL(u)
	return h
}

func (h *HttpClient) NewRequest(ctx context.Context) *Request {
	r := &Request{
		ctx:        ctx,
		header:     nil,
		mapQuery:   nil,
		httpClient: h,
		urlQuery:   url.Values{},
	}
	return r
}

func (h *HttpClient) NewPcRequest(ctx context.Context) *Request {
	r := h.NewRequest(ctx)
	r.SetHeaderPcAgent()
	return r
}

func (h *HttpClient) NewMobileRequest(ctx context.Context) *Request {
	r := h.NewRequest(ctx)
	r.SetHeaderMobileAgent()
	return r
}

func (r *Request) SetContentType(contentType string) *Request {
	r.contentType = contentType
	return r
}

func (r *Request) SetBody(body any) *Request {
	r.body = body
	return r
}

func (r *Request) SetQueryParam(k string, v string) *Request {
	if r.mapQuery == nil {
		r.mapQuery = make(map[string]string)
	}
	r.mapQuery[k] = v
	return r
}

func (r *Request) SetQueryParams(params map[string]string) *Request {
	for k, v := range params {
		r.SetQueryParam(k, v)
	}
	return r
}

func (r *Request) SetUrlQueryParam(k string, v string) *Request {
	r.urlQuery.Add(k, v)
	return r
}

func (r *Request) SetHeader(k, v string) *Request {
	if r.header == nil {
		r.header = http.Header{}
	}
	r.header.Set(k, v)
	return r
}

func (r *Request) SetHeaders(headers map[string]string) *Request {
	for k, v := range headers {
		r.SetHeader(k, v)
	}
	return r
}

func (r *Request) NoWaitQps() *Request {
	r.noWaitQps = true
	return r
}

func (r *Request) Get(rawUrl string) (*Response, error) {
	return r.Do(http.MethodGet, rawUrl)
}

func (r *Request) Post(rawUrl string) (*Response, error) {
	return r.Do(http.MethodPost, rawUrl)
}

var QpsLimitError = fmt.Errorf("qps limit")

func (r *Request) Do(method, rawUrl string) (*Response, error) {
	if r.httpClient.qpsLimiter != nil {
		if r.noWaitQps {
			allow := r.httpClient.qpsLimiter.Allow()
			if !allow {
				return nil, QpsLimitError
			}
		} else {
			err := r.httpClient.qpsLimiter.Wait(r.ctx)
			if err != nil {
				return nil, err
			}
		}
	}

	reqUrl, err := url.Parse(rawUrl)
	if err != nil {
		return nil, err
	}

	var reqBody io.Reader = nil
	if r.body != nil {
		switch v := r.body.(type) {
		case string:
			reqBody = strings.NewReader(v)
		case []byte:
			reqBody = bytes.NewReader(v)

		default:
			marshal, err := jsoniter.Marshal(r.body)
			if err != nil {
				return nil, err
			}
			reqBody = bytes.NewReader(marshal)
		}
	}

	req, err := http.NewRequest(method, reqUrl.String(), reqBody)
	if err != nil {
		return nil, err
	}

	query := req.URL.Query()

	for k, v := range r.mapQuery {
		query.Add(k, v)
	}
	for k, v := range r.urlQuery {
		query[k] = v
	}

	req.URL.RawQuery = query.Encode()

	if r.contentType != "" {
		req.Header.Set("Content-Type", r.contentType)
	}

	for k := range r.header {
		req.Header.Set(k, r.header.Get(k))
	}

	res, err := r.httpClient.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	body, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	if r.httpClient.config.noCheckStatus == false {
		if res.StatusCode != http.StatusOK {
			return nil, fmt.Errorf("status code err: %d (%s)", res.StatusCode, body)
		}
	}

	resp := &Response{
		body:     body,
		Response: res,
	}

	return resp, nil
}

func (r *Request) SetHeaderPcAgent() *Request {
	r.SetHeader(HeaderUserAgent, PcUserAgent)
	return r
}

func (r *Request) SetHeaderMobileAgent() *Request {
	r.SetHeader(HeaderUserAgent, MobileUserAgent)
	return r
}

func (r *Request) SetHeaderAcceptHtml() *Request {
	r.SetHeader(HeaderAccept, AcceptHtml)
	return r
}

func (r *Request) GetRedirectUrl(reqUrl string) (
	*url.URL, *Response, error) {

	//r.SetDebug(true)
	r.SetHeaderAcceptHtml()
	resp, err := r.Get(reqUrl)

	if err != nil && !errors.Is(err, http.ErrUseLastResponse) {
		return nil, nil, err
	}

	location, err := resp.Response.Location()
	if err != nil {
		mylog.Warnf("location code: %v", resp.Response.StatusCode)
		return nil, nil, err
	}

	return location, resp, nil
}

type Response struct {
	body     []byte
	Response *http.Response
}

func (r *Response) GetBody() []byte {
	return r.body
}