github.com/qiniu/go-sdk/v7@v7.27.0/sms/bytes/seekable/seekable.go (about)

     1  // This package provide a method to read and replace http.Request's body.
     2  package seekable
     3  
     4  import (
     5  	"errors"
     6  	"io"
     7  	"net/http"
     8  
     9  	"github.com/qiniu/go-sdk/v7/sms/bytes"
    10  )
    11  
    12  // ---------------------------------------------------
    13  
    14  type Seekable interface {
    15  	Bytes() []byte
    16  	Read(val []byte) (n int, err error)
    17  	SeekToBegin() error
    18  }
    19  
    20  type SeekableCloser interface {
    21  	Seekable
    22  	io.Closer
    23  }
    24  
    25  // ---------------------------------------------------
    26  
    27  type readCloser struct {
    28  	Seekable
    29  	io.Closer
    30  }
    31  
    32  var (
    33  	ErrNoBody       = errors.New("no body")
    34  	ErrTooLargeBody = errors.New("too large body")
    35  )
    36  
    37  const MaxBodyLength int64 = 16 * 1024 * 1024
    38  
    39  func New(req *http.Request) (r SeekableCloser, err error) {
    40  	if req.Body == nil {
    41  		return nil, ErrNoBody
    42  	}
    43  	var ok bool
    44  	if r, ok = req.Body.(SeekableCloser); ok {
    45  		return
    46  	}
    47  	b, err2 := ReadAll(req)
    48  	if err2 != nil {
    49  		return nil, err2
    50  	}
    51  	r = bytes.NewReader(b)
    52  	req.Body = readCloser{r, req.Body}
    53  	return
    54  }
    55  
    56  type readCloser2 struct {
    57  	io.Reader
    58  	io.Closer
    59  }
    60  
    61  func ReadAll(req *http.Request) (b []byte, err error) {
    62  	if req.ContentLength > MaxBodyLength {
    63  		return nil, ErrTooLargeBody
    64  	} else if req.ContentLength > 0 {
    65  		b = make([]byte, int(req.ContentLength))
    66  		_, err = io.ReadFull(req.Body, b)
    67  		return
    68  	} else if req.ContentLength == 0 {
    69  		return nil, ErrNoBody
    70  	}
    71  	b, err = io.ReadAll(io.LimitReader(req.Body, MaxBodyLength+1))
    72  	if int64(len(b)) > MaxBodyLength {
    73  		r := io.MultiReader(bytes.NewReader(b), req.Body)
    74  		req.Body = readCloser2{r, req.Body}
    75  		return nil, ErrTooLargeBody
    76  	}
    77  	return
    78  }
    79  
    80  // ---------------------------------------------------