package buf import ( "io" "time" "v2ray.com/core/common/errors" "v2ray.com/core/common/signal" ) type dataHandler func(MultiBuffer) type copyHandler struct { onData []dataHandler } // SizeCounter is for counting bytes copied by Copy(). type SizeCounter struct { Size int64 } // CopyOption is an option for copying data. type CopyOption func(*copyHandler) // UpdateActivity is a CopyOption to update activity on each data copy operation. func UpdateActivity(timer signal.ActivityUpdater) CopyOption { return func(handler *copyHandler) { handler.onData = append(handler.onData, func(MultiBuffer) { timer.Update() }) } } // CountSize is a CopyOption that sums the total size of data copied into the given SizeCounter. func CountSize(sc *SizeCounter) CopyOption { return func(handler *copyHandler) { handler.onData = append(handler.onData, func(b MultiBuffer) { sc.Size += int64(b.Len()) }) } } type readError struct { error } func (e readError) Error() string { return e.error.Error() } func (e readError) Inner() error { return e.error } func IsReadError(err error) bool { _, ok := err.(readError) return ok } type writeError struct { error } func (e writeError) Error() string { return e.error.Error() } func (e writeError) Inner() error { return e.error } func IsWriteError(err error) bool { _, ok := err.(writeError) return ok } func copyInternal(reader Reader, writer Writer, handler *copyHandler) error { for { buffer, err := reader.ReadMultiBuffer() if !buffer.IsEmpty() { for _, handler := range handler.onData { handler(buffer) } if werr := writer.WriteMultiBuffer(buffer); werr != nil { return writeError{werr} } } if err != nil { return readError{err} } } } // Copy dumps all payload from reader to writer or stops when an error occurs. It returns nil when EOF. func Copy(reader Reader, writer Writer, options ...CopyOption) error { var handler copyHandler for _, option := range options { option(&handler) } err := copyInternal(reader, writer, &handler) if err != nil && errors.Cause(err) != io.EOF { return err } return nil } var ErrNotTimeoutReader = newError("not a TimeoutReader") func CopyOnceTimeout(reader Reader, writer Writer, timeout time.Duration) error { timeoutReader, ok := reader.(TimeoutReader) if !ok { return ErrNotTimeoutReader } mb, err := timeoutReader.ReadMultiBufferTimeout(timeout) if err != nil { return err } return writer.WriteMultiBuffer(mb) }