From 4b240eb68311f78b8700b104751630e79b73458c Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sat, 4 Nov 2017 01:33:35 +0100
Subject: [PATCH] DiscardBytes

---
 common/buf/writer.go      | 32 ++++++++++++++++++++++++++++++--
 common/buf/writer_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/common/buf/writer.go b/common/buf/writer.go
index 084598e70..3944382af 100644
--- a/common/buf/writer.go
+++ b/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{}
 )
diff --git a/common/buf/writer_test.go b/common/buf/writer_test.go
index 69198f310..0dfef9b03 100644
--- a/common/buf/writer_test.go
+++ b/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)
+}