v2ray-core/common/buf/reader.go

package buf

import (
	"io"

	"v2ray.com/core/common/errors"
)

// BytesToBufferReader is a Reader that adjusts its reading speed automatically.
type BytesToBufferReader struct {
	reader io.Reader
	buffer []byte
}

// Read implements Reader.Read().
func (r *BytesToBufferReader) Read() (MultiBuffer, error) {
	nBytes, err := r.reader.Read(r.buffer)
	if err != nil {
		return nil, err
	}

	mb := NewMultiBuffer()
	p := r.buffer[:nBytes]
	for len(p) > 0 {
		b := New()
		nBytes, _ := b.Write(p)
		mb.Append(b)
		p = p[nBytes:]
	}
	return mb, nil
}

type readerAdpater struct {
	MultiBufferReader
}

func (r *readerAdpater) Read() (MultiBuffer, error) {
	return r.ReadMultiBuffer()
}

type bufferToBytesReader struct {
	stream   Reader
	leftOver MultiBuffer
}

func (r *bufferToBytesReader) Read(b []byte) (int, error) {
	if r.leftOver != nil {
		nBytes, _ := r.leftOver.Read(b)
		if r.leftOver.IsEmpty() {
			r.leftOver.Release()
			r.leftOver = nil
		}
		return nBytes, nil
	}

	mb, err := r.stream.Read()
	if err != nil {
		return 0, err
	}

	nBytes, _ := mb.Read(b)
	if !mb.IsEmpty() {
		r.leftOver = mb
	}
	return nBytes, nil
}

func (r *bufferToBytesReader) ReadMultiBuffer() (MultiBuffer, error) {
	if r.leftOver != nil {
		mb := r.leftOver
		r.leftOver = nil
		return mb, nil
	}

	return r.stream.Read()
}

func (r *bufferToBytesReader) writeToInternal(writer io.Writer) (int64, error) {
	mbWriter := NewWriter(writer)
	totalBytes := int64(0)
	if r.leftOver != nil {
		if err := mbWriter.Write(r.leftOver); err != nil {
			return 0, err
		}
		totalBytes += int64(r.leftOver.Len())
	}

	for {
		mb, err := r.stream.Read()
		if err != nil {
			return totalBytes, err
		}
		totalBytes += int64(mb.Len())
		if err := mbWriter.Write(mb); err != nil {
			return totalBytes, err
		}
	}
}

func (r *bufferToBytesReader) WriteTo(writer io.Writer) (int64, error) {
	nBytes, err := r.writeToInternal(writer)
	if errors.Cause(err) == io.EOF {
		return nBytes, nil
	}
	return nBytes, err
}