github.com/xmplusdev/xray-core@v1.8.10/common/buf/copy.go (about)

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