refine ctlcmd

common/buf/multi_buffer.go

@@ -3,7 +3,6 @@ package buf
 import (
 	"io"
 	"net"
-	"os"
 
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/errors"
@@ -14,22 +13,12 @@ import (
 func ReadAllToMultiBuffer(reader io.Reader) (MultiBuffer, error) {
 	mb := NewMultiBufferCap(128)
 
-	for {
-		b := New()
-		err := b.Reset(ReadFrom(reader))
-		if b.IsEmpty() {
-			b.Release()
-		} else {
-			mb.Append(b)
-		}
-		if err != nil {
-			if errors.Cause(err) == io.EOF || errors.Cause(err) == os.ErrClosed {
-				return mb, nil
-			}
-			mb.Release()
-			return nil, err
-		}
+	if _, err := mb.ReadFrom(reader); err != nil {
+		mb.Release()
+		return nil, err
 	}
+
+	return mb, nil
 }
 
 // ReadSizeToMultiBuffer reads specific number of bytes from reader into a MultiBuffer.
@@ -102,6 +91,28 @@ func (mb MultiBuffer) Copy(b []byte) int {
 	return total
 }
 
+// ReadFrom implements io.ReaderFrom.
+func (mb *MultiBuffer) ReadFrom(reader io.Reader) (int64, error) {
+	totalBytes := int64(0)
+
+	for {
+		b := New()
+		err := b.Reset(ReadFrom(reader))
+		if b.IsEmpty() {
+			b.Release()
+		} else {
+			mb.Append(b)
+		}
+		totalBytes += int64(b.Len())
+		if err != nil {
+			if errors.Cause(err) == io.EOF {
+				return totalBytes, nil
+			}
+			return totalBytes, err
+		}
+	}
+}
+
 // Read implements io.Reader.
 func (mb *MultiBuffer) Read(b []byte) (int, error) {
 	if mb.Len() == 0 {
@@ -125,6 +136,22 @@ func (mb *MultiBuffer) Read(b []byte) (int, error) {
 	return totalBytes, nil
 }
 
+// WriteTo implements io.WriterTo.
+func (mb *MultiBuffer) WriteTo(writer io.Writer) (int64, error) {
+	defer mb.Release()
+
+	totalBytes := int64(0)
+	for _, b := range *mb {
+		nBytes, err := writer.Write(b.Bytes())
+		totalBytes += int64(nBytes)
+		if err != nil {
+			return totalBytes, err
+		}
+	}
+
+	return totalBytes, nil
+}
+
 // Write implements io.Writer.
 func (mb *MultiBuffer) Write(b []byte) (int, error) {
 	totalBytes := len(b)