github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/rest/http.go (about)

     1  package rest
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"crypto/tls"
     7  	"encoding/base64"
     8  	"encoding/json"
     9  	"io/ioutil"
    10  	"mime/multipart"
    11  	"net/http"
    12  	"time"
    13  
    14  	"github.com/pkg/errors"
    15  )
    16  
    17  type Rest struct {
    18  	Body             []byte
    19  	Method, Addr     string
    20  	Headers          map[string]string
    21  	Result           interface{}
    22  	DisableKeepAlive bool
    23  	Context          context.Context
    24  	Timeout          time.Duration
    25  	BasicAuth        string
    26  }
    27  
    28  // Post execute HTTP POST request.
    29  func (r Rest) Post() (*Rsp, error) { return r.do("POST") }
    30  
    31  // Get execute HTTP GET request.
    32  func (r Rest) Get() (*Rsp, error) { return r.do("GET") }
    33  
    34  // Delete execute HTTP GET request.
    35  func (r Rest) Delete() (*Rsp, error) { return r.do("DELETE") }
    36  
    37  // Upload execute HTTP GET request.
    38  func (r Rest) Upload(filename string, fileData []byte) (*Rsp, error) {
    39  	body := &bytes.Buffer{}
    40  	writer := multipart.NewWriter(body)
    41  	part, _ := writer.CreateFormFile("file", filename)
    42  	_, _ = part.Write(fileData)
    43  	_ = writer.Close()
    44  
    45  	if r.Headers == nil {
    46  		r.Headers = make(map[string]string)
    47  	}
    48  
    49  	r.Body = body.Bytes()
    50  	r.Headers["Content-Type"] = writer.FormDataContentType()
    51  	r.Method = "POST"
    52  	return r.Do()
    53  }
    54  
    55  type Rsp struct {
    56  	Body   []byte
    57  	Status int
    58  	Header http.Header
    59  	Cost   time.Duration
    60  }
    61  
    62  var Client = &http.Client{
    63  	// Timeout: 10 * time.Second,
    64  	Transport: &http.Transport{
    65  		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
    66  		// DisableKeepAlives: true,
    67  	},
    68  }
    69  
    70  // Do execute HTTP method request.
    71  func (r Rest) Do() (*Rsp, error) {
    72  	return r.do(If(r.Method == "", "GET", r.Method))
    73  }
    74  
    75  // Do execute HTTP method request.
    76  func (r Rest) do(method string) (*Rsp, error) {
    77  	var ctx context.Context
    78  	if r.Context != nil {
    79  		ctx = r.Context
    80  	} else {
    81  		var cancel context.CancelFunc
    82  		ctx, cancel = context.WithTimeout(context.Background(), DurationOr(r.Timeout, 10*time.Second))
    83  		defer cancel()
    84  	}
    85  
    86  	req, err := http.NewRequestWithContext(ctx, method, r.Addr, bytes.NewReader(r.Body))
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  	req.Close = r.DisableKeepAlive
    91  	if req.Header.Get("Content-Type") == "" {
    92  		req.Header.Set("Content-Type", If(IsJSONBytes(r.Body),
    93  			"application/json; charset=utf-8", "text/plain; charset=utf-8"))
    94  	}
    95  
    96  	if r.BasicAuth != "" {
    97  		req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(r.BasicAuth)))
    98  	}
    99  
   100  	for k, v := range r.Headers {
   101  		req.Header.Set(k, v)
   102  	}
   103  
   104  	start := time.Now()
   105  	resp, err := Client.Do(req)
   106  	cost := time.Since(start)
   107  	if err != nil {
   108  		return nil, err
   109  	}
   110  
   111  	rsp := &Rsp{Status: resp.StatusCode, Header: resp.Header, Cost: cost}
   112  	bodyData, err := ioutil.ReadAll(resp.Body)
   113  	if err != nil {
   114  		return rsp, err
   115  	}
   116  	_ = resp.Body.Close()
   117  
   118  	rsp.Body = bodyData
   119  	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
   120  		if r.Result != nil {
   121  			_ = json.Unmarshal(bodyData, r.Result)
   122  		}
   123  		return rsp, nil
   124  	}
   125  
   126  	return rsp, errors.Wrapf(err, "status:%d", resp.StatusCode)
   127  }
   128  
   129  func DurationOr(a, b time.Duration) time.Duration {
   130  	if a == 0 {
   131  		return b
   132  	}
   133  
   134  	return a
   135  }
   136  
   137  // If tests condition to return a or b.
   138  func If(condition bool, a, b string) string {
   139  	if condition {
   140  		return a
   141  	}
   142  
   143  	return b
   144  }
   145  
   146  // IsJSONBytes tests bytes b is in JSON format.
   147  func IsJSONBytes(b []byte) bool {
   148  	if len(b) == 0 {
   149  		return false
   150  	}
   151  
   152  	var m interface{}
   153  	return json.Unmarshal(b, &m) == nil
   154  }