1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-10-29 23:37:54 -04:00
v2fly/common/buf/copy.go
2018-10-24 12:02:02 +02:00

145 lines
2.9 KiB
Go

package buf
import (
"io"
"time"
"v2ray.com/core/common/errors"
"v2ray.com/core/common/signal"
)
type errorHandler func(error) error
type dataHandler func(MultiBuffer)
type copyHandler struct {
onReadError []errorHandler
onData []dataHandler
onWriteError []errorHandler
}
func (h *copyHandler) readFrom(reader Reader) (MultiBuffer, error) {
mb, err := reader.ReadMultiBuffer()
if err != nil {
for _, handler := range h.onReadError {
err = handler(err)
}
}
return mb, err
}
func (h *copyHandler) writeTo(writer Writer, mb MultiBuffer) error {
err := writer.WriteMultiBuffer(mb)
if err != nil {
for _, handler := range h.onWriteError {
err = handler(err)
}
}
return err
}
// 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 := handler.readFrom(reader)
if !buffer.IsEmpty() {
for _, handler := range handler.onData {
handler(buffer)
}
if werr := handler.writeTo(writer, 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)
}