github.com/v2fly/v2ray-core/v4@v4.45.2/common/buf/copy.go (about)

     1  package buf
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/v2fly/v2ray-core/v4/common/errors"
     8  	"github.com/v2fly/v2ray-core/v4/common/signal"
     9  )
    10  
    11  type dataHandler func(MultiBuffer)
    12  
    13  type copyHandler struct {
    14  	onData []dataHandler
    15  }
    16  
    17  // SizeCounter is for counting bytes copied by Copy().
    18  type SizeCounter struct {
    19  	Size int64
    20  }
    21  
    22  // CopyOption is an option for copying data.
    23  type CopyOption func(*copyHandler)
    24  
    25  // UpdateActivity is a CopyOption to update activity on each data copy operation.
    26  func UpdateActivity(timer signal.ActivityUpdater) CopyOption {
    27  	return func(handler *copyHandler) {
    28  		handler.onData = append(handler.onData, func(MultiBuffer) {
    29  			timer.Update()
    30  		})
    31  	}
    32  }
    33  
    34  // CountSize is a CopyOption that sums the total size of data copied into the given SizeCounter.
    35  func CountSize(sc *SizeCounter) CopyOption {
    36  	return func(handler *copyHandler) {
    37  		handler.onData = append(handler.onData, func(b MultiBuffer) {
    38  			sc.Size += int64(b.Len())
    39  		})
    40  	}
    41  }
    42  
    43  type readError struct {
    44  	error
    45  }
    46  
    47  func (e readError) Error() string {
    48  	return e.error.Error()
    49  }
    50  
    51  func (e readError) Inner() error {
    52  	return e.error
    53  }
    54  
    55  // IsReadError returns true if the error in Copy() comes from reading.
    56  func IsReadError(err error) bool {
    57  	_, ok := err.(readError)
    58  	return ok
    59  }
    60  
    61  type writeError struct {
    62  	error
    63  }
    64  
    65  func (e writeError) Error() string {
    66  	return e.error.Error()
    67  }
    68  
    69  func (e writeError) Inner() error {
    70  	return e.error
    71  }
    72  
    73  // IsWriteError returns true if the error in Copy() comes from writing.
    74  func IsWriteError(err error) bool {
    75  	_, ok := err.(writeError)
    76  	return ok
    77  }
    78  
    79  func copyInternal(reader Reader, writer Writer, handler *copyHandler) error {
    80  	for {
    81  		buffer, err := reader.ReadMultiBuffer()
    82  		if !buffer.IsEmpty() {
    83  			for _, handler := range handler.onData {
    84  				handler(buffer)
    85  			}
    86  
    87  			if werr := writer.WriteMultiBuffer(buffer); werr != nil {
    88  				return writeError{werr}
    89  			}
    90  		}
    91  
    92  		if err != nil {
    93  			return readError{err}
    94  		}
    95  	}
    96  }
    97  
    98  // Copy dumps all payload from reader to writer or stops when an error occurs. It returns nil when EOF.
    99  func Copy(reader Reader, writer Writer, options ...CopyOption) error {
   100  	var handler copyHandler
   101  	for _, option := range options {
   102  		option(&handler)
   103  	}
   104  	err := copyInternal(reader, writer, &handler)
   105  	if err != nil && errors.Cause(err) != io.EOF {
   106  		return err
   107  	}
   108  	return nil
   109  }
   110  
   111  var ErrNotTimeoutReader = newError("not a TimeoutReader")
   112  
   113  func CopyOnceTimeout(reader Reader, writer Writer, timeout time.Duration) error {
   114  	timeoutReader, ok := reader.(TimeoutReader)
   115  	if !ok {
   116  		return ErrNotTimeoutReader
   117  	}
   118  	mb, err := timeoutReader.ReadMultiBufferTimeout(timeout)
   119  	if err != nil {
   120  		return err
   121  	}
   122  	return writer.WriteMultiBuffer(mb)
   123  }