copy.go 2.6 KB

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