mirror of https://github.com/v2ray/v2ray-core
refine copy handler
Darien Raymond
parent
b86b998a20
commit
07847576b5
|
|
@ -0,0 +1,100 @@
|
|||
package buf
|
||||
|
||||
import (
|
||||
"io"
|
||||
|
||||
"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.Read()
|
||||
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.Write(mb)
|
||||
if err != nil {
|
||||
for _, handler := range h.onWriteError {
|
||||
err = handler(err)
|
||||
}
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
type CopyOption func(*copyHandler)
|
||||
|
||||
func IgnoreReaderError() CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onReadError = append(handler.onReadError, func(err error) error {
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func IgnoreWriterError() CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onWriteError = append(handler.onWriteError, func(err error) error {
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func UpdateActivity(timer signal.ActivityTimer) CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onData = append(handler.onData, func(MultiBuffer) {
|
||||
timer.Update()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func copyInternal(reader Reader, writer Writer, handler *copyHandler) error {
|
||||
for {
|
||||
buffer, err := handler.readFrom(reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if buffer.IsEmpty() {
|
||||
buffer.Release()
|
||||
continue
|
||||
}
|
||||
|
||||
for _, handler := range handler.onData {
|
||||
handler(buffer)
|
||||
}
|
||||
|
||||
if err := handler.writeTo(writer, buffer); err != nil {
|
||||
buffer.Release()
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Copy dumps all payload from reader to writer or stops when an error occurs.
|
||||
// ActivityTimer gets updated as soon as there is a payload.
|
||||
func Copy(reader Reader, writer Writer, options ...CopyOption) error {
|
||||
handler := new(copyHandler)
|
||||
for _, option := range options {
|
||||
option(handler)
|
||||
}
|
||||
err := copyInternal(reader, writer, handler)
|
||||
if err != nil && errors.Cause(err) != io.EOF {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
@ -3,9 +3,6 @@ package buf
|
|||
import (
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"v2ray.com/core/common/errors"
|
||||
"v2ray.com/core/common/signal"
|
||||
)
|
||||
|
||||
// Reader extends io.Reader with alloc.Buffer.
|
||||
|
|
@ -47,91 +44,6 @@ func ReadAtLeastFrom(reader io.Reader, size int) Supplier {
|
|||
}
|
||||
}
|
||||
|
||||
type copyHandler struct {
|
||||
onReadError func(error) error
|
||||
onData func()
|
||||
onWriteError func(error) error
|
||||
}
|
||||
|
||||
func (h *copyHandler) readFrom(reader Reader) (MultiBuffer, error) {
|
||||
mb, err := reader.Read()
|
||||
if err != nil && h.onReadError != nil {
|
||||
err = h.onReadError(err)
|
||||
}
|
||||
return mb, err
|
||||
}
|
||||
|
||||
func (h *copyHandler) writeTo(writer Writer, mb MultiBuffer) error {
|
||||
err := writer.Write(mb)
|
||||
if err != nil && h.onWriteError != nil {
|
||||
err = h.onWriteError(err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
type CopyOption func(*copyHandler)
|
||||
|
||||
func IgnoreReaderError() CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onReadError = func(err error) error {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func IgnoreWriterError() CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onWriteError = func(err error) error {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func UpdateActivity(timer signal.ActivityTimer) CopyOption {
|
||||
return func(handler *copyHandler) {
|
||||
handler.onData = func() {
|
||||
timer.Update()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func copyInternal(reader Reader, writer Writer, handler *copyHandler) error {
|
||||
for {
|
||||
buffer, err := handler.readFrom(reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if buffer.IsEmpty() {
|
||||
buffer.Release()
|
||||
continue
|
||||
}
|
||||
|
||||
if handler.onData != nil {
|
||||
handler.onData()
|
||||
}
|
||||
|
||||
if err := handler.writeTo(writer, buffer); err != nil {
|
||||
buffer.Release()
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Copy dumps all payload from reader to writer or stops when an error occurs.
|
||||
// ActivityTimer gets updated as soon as there is a payload.
|
||||
func Copy(reader Reader, writer Writer, options ...CopyOption) error {
|
||||
handler := new(copyHandler)
|
||||
for _, option := range options {
|
||||
option(handler)
|
||||
}
|
||||
err := copyInternal(reader, writer, handler)
|
||||
if err != nil && errors.Cause(err) != io.EOF {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewReader creates a new Reader.
|
||||
// The Reader instance doesn't take the ownership of reader.
|
||||
func NewReader(reader io.Reader) Reader {
|
||||
|
|
|
|||
Loading…
Reference in New Issue