more test code

common/buf/reader.go

@@ -27,26 +27,6 @@ func (r *BytesToBufferReader) Read() (MultiBuffer, error) {
 	return mb, nil
 }
 
-func (r *BytesToBufferReader) WriteTo(writer io.Writer) (int64, error) {
-	totalBytes := int64(0)
-	eof := false
-	for !eof {
-		if err := r.buffer.Reset(ReadFrom(r.reader)); err != nil {
-			if errors.Cause(err) == io.EOF {
-				eof = true
-			} else {
-				return totalBytes, err
-			}
-		}
-		nBytes, err := writer.Write(r.buffer.Bytes())
-		totalBytes += int64(nBytes)
-		if err != nil {
-			return totalBytes, err
-		}
-	}
-	return totalBytes, nil
-}
-
 type readerAdpater struct {
 	MultiBufferReader
 }
@@ -119,7 +99,7 @@ func (r *bufferToBytesReader) writeToInternal(writer io.Writer) (int64, error) {
 				return totalBytes, r.err
 			}
 		}
-		totalBytes := int64(r.current.Len())
+		totalBytes += int64(r.current.Len())
 		if err := mbWriter.Write(r.current); err != nil {
 			return totalBytes, err
 		}

common/buf/reader_test.go

@@ -2,10 +2,13 @@ package buf_test
 
 import (
 	"bytes"
+	"context"
+	"io"
 	"testing"
 
 	. "v2ray.com/core/common/buf"
 	"v2ray.com/core/testing/assert"
+	"v2ray.com/core/transport/ray"
 )
 
 func TestAdaptiveReader(t *testing.T) {
@@ -19,3 +22,49 @@ func TestAdaptiveReader(t *testing.T) {
 	assert.Error(err).IsNil()
 	assert.Int(b.Len()).Equals(32 * 1024)
 }
+
+func TestBytesReaderWriteTo(t *testing.T) {
+	assert := assert.On(t)
+
+	stream := ray.NewStream(context.Background())
+	reader := ToBytesReader(stream)
+	b1 := New()
+	b1.AppendBytes('a', 'b', 'c')
+	b2 := New()
+	b2.AppendBytes('e', 'f', 'g')
+	assert.Error(stream.Write(NewMultiBufferValue(b1, b2))).IsNil()
+	stream.Close()
+
+	stream2 := ray.NewStream(context.Background())
+	writer := ToBytesWriter(stream2)
+
+	nBytes, err := io.Copy(writer, reader)
+	assert.Error(err).IsNil()
+	assert.Int64(nBytes).Equals(6)
+
+	mb, err := stream2.Read()
+	assert.Error(err).IsNil()
+	assert.Int(len(mb)).Equals(2)
+	assert.String(mb[0].String()).Equals("abc")
+	assert.String(mb[1].String()).Equals("efg")
+}
+
+func TestBytesReaderMultiBuffer(t *testing.T) {
+	assert := assert.On(t)
+
+	stream := ray.NewStream(context.Background())
+	reader := ToBytesReader(stream)
+	b1 := New()
+	b1.AppendBytes('a', 'b', 'c')
+	b2 := New()
+	b2.AppendBytes('e', 'f', 'g')
+	assert.Error(stream.Write(NewMultiBufferValue(b1, b2))).IsNil()
+	stream.Close()
+
+	mbReader := NewReader(reader)
+	mb, err := mbReader.Read()
+	assert.Error(err).IsNil()
+	assert.Int(len(mb)).Equals(2)
+	assert.String(mb[0].String()).Equals("abc")
+	assert.String(mb[1].String()).Equals("efg")
+}