DiscardBytes

common/buf/writer.go

@@ -1,6 +1,10 @@
 package buf
 
-import "io"
+import (
+	"io"
+
+	"v2ray.com/core/common/errors"
+)
 
 // BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer.
 type BufferToBytesWriter struct {
@@ -109,6 +113,30 @@ func (noOpWriter) Write(b MultiBuffer) error {
 	return nil
 }
 
+type noOpBytesWriter struct{}
+
+func (noOpBytesWriter) Write(b []byte) (int, error) {
+	return len(b), nil
+}
+
+func (noOpBytesWriter) ReadFrom(reader io.Reader) (int64, error) {
+	b := New()
+	defer b.Release()
+
+	totalBytes := int64(0)
+	for {
+		err := b.Reset(ReadFrom(reader))
+		totalBytes += int64(b.Len())
+		if err != nil {
+			if errors.Cause(err) == io.EOF {
+				return totalBytes, nil
+			}
+			return totalBytes, err
+		}
+	}
+}
+
 var (
-	Discard Writer = noOpWriter{}
+	Discard      Writer    = noOpWriter{}
+	DiscardBytes io.Writer = noOpBytesWriter{}
 )

common/buf/writer_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"io"
 
+	"v2ray.com/core/common"
 	. "v2ray.com/core/common/buf"
 	"v2ray.com/core/transport/ray"
 	. "v2ray.com/ext/assert"
@@ -43,3 +44,27 @@ func TestBytesWriterReadFrom(t *testing.T) {
 	assert(mb.Len(), Equals, 8192)
 	assert(len(mb), Equals, 4)
 }
+
+func TestDiscardBytes(t *testing.T) {
+	assert := With(t)
+
+	b := New()
+	common.Must(b.Reset(ReadFrom(rand.Reader)))
+
+	nBytes, err := io.Copy(DiscardBytes, b)
+	assert(nBytes, Equals, int64(Size))
+	assert(err, IsNil)
+}
+
+func TestDiscardBytesMultiBuffer(t *testing.T) {
+	assert := With(t)
+
+	const size = 10240*1024 + 1
+	buffer := bytes.NewBuffer(make([]byte, 0, size))
+	common.Must2(buffer.ReadFrom(io.LimitReader(rand.Reader, size)))
+
+	r := NewReader(buffer)
+	nBytes, err := io.Copy(DiscardBytes, ToBytesReader(r))
+	assert(nBytes, Equals, int64(size))
+	assert(err, IsNil)
+}