Xray-core/common/buf/copy.go

135 lines
2.9 KiB
Go
Raw Normal View History

2020-11-25 13:01:53 +02:00
package buf
import (
"io"
"time"
2020-12-04 03:36:16 +02:00
"github.com/xtls/xray-core/common/errors"
"github.com/xtls/xray-core/common/signal"
"github.com/xtls/xray-core/features/stats"
2020-11-25 13:01:53 +02:00
)
type dataHandler func(MultiBuffer)
2024-09-22 05:02:08 +03:00
type CopyHandler struct {
OnData []dataHandler
2020-11-25 13:01:53 +02:00
}
// SizeCounter is for counting bytes copied by Copy().
type SizeCounter struct {
Size int64
}
// CopyOption is an option for copying data.
2024-09-22 05:02:08 +03:00
type CopyOption func(*CopyHandler)
2020-11-25 13:01:53 +02:00
// UpdateActivity is a CopyOption to update activity on each data copy operation.
func UpdateActivity(timer signal.ActivityUpdater) CopyOption {
2024-09-22 05:02:08 +03:00
return func(handler *CopyHandler) {
handler.OnData = append(handler.OnData, func(MultiBuffer) {
2020-11-25 13:01:53 +02:00
timer.Update()
})
}
}
// CountSize is a CopyOption that sums the total size of data copied into the given SizeCounter.
func CountSize(sc *SizeCounter) CopyOption {
2024-09-22 05:02:08 +03:00
return func(handler *CopyHandler) {
handler.OnData = append(handler.OnData, func(b MultiBuffer) {
2020-11-25 13:01:53 +02:00
sc.Size += int64(b.Len())
})
}
}
// AddToStatCounter a CopyOption add to stat counter
func AddToStatCounter(sc stats.Counter) CopyOption {
2024-09-22 05:02:08 +03:00
return func(handler *CopyHandler) {
handler.OnData = append(handler.OnData, func(b MultiBuffer) {
if sc != nil {
sc.Add(int64(b.Len()))
}
})
}
}
2020-11-25 13:01:53 +02:00
type readError struct {
error
}
func (e readError) Error() string {
return e.error.Error()
}
2022-06-01 06:03:06 +03:00
func (e readError) Unwrap() error {
2020-11-25 13:01:53 +02:00
return e.error
}
// IsReadError returns true if the error in Copy() comes from reading.
func IsReadError(err error) bool {
_, ok := err.(readError)
return ok
}
type writeError struct {
error
}
func (e writeError) Error() string {
return e.error.Error()
}
2022-06-01 06:03:06 +03:00
func (e writeError) Unwrap() error {
2020-11-25 13:01:53 +02:00
return e.error
}
// IsWriteError returns true if the error in Copy() comes from writing.
func IsWriteError(err error) bool {
_, ok := err.(writeError)
return ok
}
2024-09-22 05:02:08 +03:00
func copyInternal(reader Reader, writer Writer, handler *CopyHandler) error {
2020-11-25 13:01:53 +02:00
for {
buffer, err := reader.ReadMultiBuffer()
if !buffer.IsEmpty() {
if werr := writer.WriteMultiBuffer(buffer); werr != nil {
return writeError{werr}
}
}
2024-09-22 05:02:08 +03:00
for _, handler := range handler.OnData {
handler(buffer)
}
2020-11-25 13:01:53 +02:00
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 {
2024-09-22 05:02:08 +03:00
var handler CopyHandler
2020-11-25 13:01:53 +02:00
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 = errors.New("not a TimeoutReader")
2020-11-25 13:01:53 +02:00
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)
}