github.com/searKing/golang/go@v1.2.117/net/http/rewind.go (about)

     1  // Copyright 2022 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package http
     6  
     7  import (
     8  	"bytes"
     9  	"errors"
    10  	"io"
    11  	"io/ioutil"
    12  	"net/http"
    13  	"os"
    14  	"reflect"
    15  	"strings"
    16  
    17  	io_ "github.com/searKing/golang/go/io"
    18  )
    19  
    20  var (
    21  	ErrBodyNotRewindable = errors.New("body not rewindable")
    22  )
    23  var nopCloserType = reflect.TypeOf(io.NopCloser(nil))
    24  
    25  func RequestRewindableWithFileName(name string) (body io.ReadCloser, getBody func() (io.ReadCloser, error), err error) {
    26  	return BodyRewindableWithFilePosition(name, 0, io.SeekCurrent)
    27  }
    28  
    29  func BodyRewindableWithFilePosition(name string, offset int64, whence int) (body io.ReadCloser, getBody func() (io.ReadCloser, error), err error) {
    30  	file, err := os.Open(name)
    31  	if err != nil {
    32  		return nil, nil, err
    33  	}
    34  
    35  	if _, err = file.Seek(offset, whence); err != nil {
    36  		return nil, nil, err
    37  	}
    38  
    39  	return BodyRewindableWithFile(file)
    40  }
    41  
    42  // BodyRewindableWithFile returns a Request suitable for use with Redirect, like 307 redirect for PUT or POST.
    43  // Only a nil GetBody in Request may be replace with a rewindable GetBody, which is a Body *os.File.
    44  // See: https://github.com/golang/go/issues/7912
    45  // See also: https://go-review.googlesource.com/c/go/+/29852/13/src/net/http/client.go#391
    46  func BodyRewindableWithFile(file *os.File) (body io.ReadCloser, getBody func() (io.ReadCloser, error), err error) {
    47  	offset, err := file.Seek(0, io.SeekCurrent)
    48  	if err != nil {
    49  		return nil, nil, err
    50  	}
    51  
    52  	return file, func() (io.ReadCloser, error) {
    53  		file_, err_ := os.Open(file.Name())
    54  		if err_ != nil {
    55  			return nil, err
    56  		}
    57  
    58  		if _, err_ = file_.Seek(offset, io.SeekStart); err != nil {
    59  			return nil, err_
    60  		}
    61  		return file_, err_
    62  	}, nil
    63  }
    64  
    65  // RequestWithBodyRewindable returns a Request suitable for use with Redirect, like 307 redirect for PUT or POST.
    66  // Only a nil GetBody in Request may be replaced with a rewindable GetBody, which is a Body replayer.
    67  // A body with a type not ioutil.NopCloser(nil) may return error as the Body in Request will be closed before redirect automatically.
    68  // So you can close body by yourself to ensure rewindable always:
    69  // Examples:
    70  //
    71  //	body := req.Body
    72  //	defer body.Close() // body will not be closed inside
    73  //	req.Body = ioutil.NopCloser(body)
    74  //	_ = RequestWithBodyRewindable(req)
    75  //
    76  // // do http requests...
    77  //
    78  // See: https://github.com/golang/go/issues/7912
    79  // See also: https://go-review.googlesource.com/c/go/+/29852/13/src/net/http/client.go#391
    80  func RequestWithBodyRewindable(req *http.Request) error {
    81  	if req.Body == nil || req.Body == http.NoBody {
    82  		// No copying needed.
    83  		return nil
    84  	}
    85  
    86  	// If the request body can be reset back to its original
    87  	// state via the optional req.GetBody, do that.
    88  	if req.GetBody != nil {
    89  		return nil
    90  	}
    91  
    92  	body, neverClose := isNeverCloseReader(req.Body)
    93  	if !neverClose {
    94  		// Body in Request will be closed before redirect automatically, so io.Seeker can not be used.
    95  
    96  		// take care of *os.File, which can be reopen
    97  		switch body_ := body.(type) {
    98  		case *os.File:
    99  			_, getBody, err := BodyRewindableWithFile(body_)
   100  			if err != nil {
   101  				return err
   102  			}
   103  			req.GetBody = getBody
   104  			return nil
   105  		}
   106  		return ErrBodyNotRewindable
   107  	}
   108  
   109  	// handle never closed body
   110  
   111  	// NewRequest and NewRequestWithContext in net/http will handle
   112  	// See: https://github.com/golang/go/blob/2117ea9737bc9cb2e30cb087b76a283f68768819/src/net/http/request.go#L873
   113  	switch v := body.(type) {
   114  	case *bytes.Buffer:
   115  		req.ContentLength = int64(v.Len())
   116  		buf := v.Bytes()
   117  		req.GetBody = func() (io.ReadCloser, error) {
   118  			r := bytes.NewReader(buf)
   119  			return io.NopCloser(r), nil
   120  		}
   121  		return nil
   122  	case *bytes.Reader:
   123  		req.ContentLength = int64(v.Len())
   124  		snapshot := *v
   125  		req.GetBody = func() (io.ReadCloser, error) {
   126  			r := snapshot
   127  			return io.NopCloser(&r), nil
   128  		}
   129  		return nil
   130  	case *strings.Reader:
   131  		req.ContentLength = int64(v.Len())
   132  		snapshot := *v
   133  		req.GetBody = func() (io.ReadCloser, error) {
   134  			r := snapshot
   135  			return io.NopCloser(&r), nil
   136  		}
   137  		return nil
   138  	case *os.File:
   139  		_, getBody, err := BodyRewindableWithFile(v)
   140  		if err != nil {
   141  			return err
   142  		}
   143  		req.GetBody = getBody
   144  		return nil
   145  	}
   146  
   147  	// Handle unknown types
   148  
   149  	// Use io.Seeker to rewind if Seek succeed
   150  	{
   151  		if seeker, ok := body.(io.Seeker); ok {
   152  			offset, err := seeker.Seek(0, io.SeekCurrent)
   153  			if err == nil {
   154  				req.GetBody = func() (io.ReadCloser, error) {
   155  					_, err := seeker.Seek(offset, io.SeekStart)
   156  					if err != nil {
   157  						return nil, err
   158  					}
   159  					return req.Body, nil
   160  				}
   161  			}
   162  		}
   163  	}
   164  
   165  	// Use a replay reader to capture any body sent in case we have to replay it again
   166  	// All data will be buffered in memory in case of reread.
   167  	{
   168  		replayR := io_.ReplayReader(req.Body)
   169  		replayRC := replayReadCloser{Reader: replayR, Closer: req.Body}
   170  		req.Body = replayRC
   171  		req.GetBody = func() (io.ReadCloser, error) {
   172  			replayR.Replay()
   173  
   174  			// Refresh the body reader so the body can be sent again
   175  			// take care of req.Body set to nil by caller outside
   176  			if req.Body == nil {
   177  				return nil, nil
   178  			}
   179  			return ioutil.NopCloser(replayR), nil
   180  		}
   181  	}
   182  	return nil
   183  }
   184  
   185  type replayReadCloser struct {
   186  	io.Reader
   187  	io.Closer
   188  }
   189  
   190  // isNeverCloseReader reports whether r is a type known to not closed.
   191  // Its caller uses this as an optional optimization to
   192  // send fewer TCP packets.
   193  func isNeverCloseReader(r io.Reader) (rr io.Reader, nopClose bool) {
   194  	return _isNeverCloseReader(r, false)
   195  }
   196  
   197  func _isNeverCloseReader(r io.Reader, nopclose bool) (rr io.Reader, nopClose bool) {
   198  	switch r.(type) {
   199  	case *bytes.Reader, *bytes.Buffer, *strings.Reader:
   200  		return r, true
   201  	}
   202  	if reflect.TypeOf(r) == nopCloserType {
   203  		return _isNeverCloseReader(reflect.ValueOf(r).Field(0).Interface().(io.Reader), true)
   204  	}
   205  	return r, nopclose
   206  }