v1.0.0

common/buf/reader.go

@@ -0,0 +1,174 @@
+package buf
+
+import (
+	"io"
+
+	"github.com/xtls/xray-core/v1/common"
+	"github.com/xtls/xray-core/v1/common/errors"
+)
+
+func readOneUDP(r io.Reader) (*Buffer, error) {
+	b := New()
+	for i := 0; i < 64; i++ {
+		_, err := b.ReadFrom(r)
+		if !b.IsEmpty() {
+			return b, nil
+		}
+		if err != nil {
+			b.Release()
+			return nil, err
+		}
+	}
+
+	b.Release()
+	return nil, newError("Reader returns too many empty payloads.")
+}
+
+// ReadBuffer reads a Buffer from the given reader.
+func ReadBuffer(r io.Reader) (*Buffer, error) {
+	b := New()
+	n, err := b.ReadFrom(r)
+	if n > 0 {
+		return b, err
+	}
+	b.Release()
+	return nil, err
+}
+
+// BufferedReader is a Reader that keeps its internal buffer.
+type BufferedReader struct {
+	// Reader is the underlying reader to be read from
+	Reader Reader
+	// Buffer is the internal buffer to be read from first
+	Buffer MultiBuffer
+	// Spliter is a function to read bytes from MultiBuffer
+	Spliter func(MultiBuffer, []byte) (MultiBuffer, int)
+}
+
+// BufferedBytes returns the number of bytes that is cached in this reader.
+func (r *BufferedReader) BufferedBytes() int32 {
+	return r.Buffer.Len()
+}
+
+// ReadByte implements io.ByteReader.
+func (r *BufferedReader) ReadByte() (byte, error) {
+	var b [1]byte
+	_, err := r.Read(b[:])
+	return b[0], err
+}
+
+// Read implements io.Reader. It reads from internal buffer first (if available) and then reads from the underlying reader.
+func (r *BufferedReader) Read(b []byte) (int, error) {
+	spliter := r.Spliter
+	if spliter == nil {
+		spliter = SplitBytes
+	}
+
+	if !r.Buffer.IsEmpty() {
+		buffer, nBytes := spliter(r.Buffer, b)
+		r.Buffer = buffer
+		if r.Buffer.IsEmpty() {
+			r.Buffer = nil
+		}
+		return nBytes, nil
+	}
+
+	mb, err := r.Reader.ReadMultiBuffer()
+	if err != nil {
+		return 0, err
+	}
+
+	mb, nBytes := spliter(mb, b)
+	if !mb.IsEmpty() {
+		r.Buffer = mb
+	}
+	return nBytes, nil
+}
+
+// ReadMultiBuffer implements Reader.
+func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) {
+	if !r.Buffer.IsEmpty() {
+		mb := r.Buffer
+		r.Buffer = nil
+		return mb, nil
+	}
+
+	return r.Reader.ReadMultiBuffer()
+}
+
+// ReadAtMost returns a MultiBuffer with at most size.
+func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
+	if r.Buffer.IsEmpty() {
+		mb, err := r.Reader.ReadMultiBuffer()
+		if mb.IsEmpty() && err != nil {
+			return nil, err
+		}
+		r.Buffer = mb
+	}
+
+	rb, mb := SplitSize(r.Buffer, size)
+	r.Buffer = rb
+	if r.Buffer.IsEmpty() {
+		r.Buffer = nil
+	}
+	return mb, nil
+}
+
+func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) {
+	mbWriter := NewWriter(writer)
+	var sc SizeCounter
+	if r.Buffer != nil {
+		sc.Size = int64(r.Buffer.Len())
+		if err := mbWriter.WriteMultiBuffer(r.Buffer); err != nil {
+			return 0, err
+		}
+		r.Buffer = nil
+	}
+
+	err := Copy(r.Reader, mbWriter, CountSize(&sc))
+	return sc.Size, err
+}
+
+// WriteTo implements io.WriterTo.
+func (r *BufferedReader) WriteTo(writer io.Writer) (int64, error) {
+	nBytes, err := r.writeToInternal(writer)
+	if errors.Cause(err) == io.EOF {
+		return nBytes, nil
+	}
+	return nBytes, err
+}
+
+// Interrupt implements common.Interruptible.
+func (r *BufferedReader) Interrupt() {
+	common.Interrupt(r.Reader)
+}
+
+// Close implements io.Closer.
+func (r *BufferedReader) Close() error {
+	return common.Close(r.Reader)
+}
+
+// SingleReader is a Reader that read one Buffer every time.
+type SingleReader struct {
+	io.Reader
+}
+
+// ReadMultiBuffer implements Reader.
+func (r *SingleReader) ReadMultiBuffer() (MultiBuffer, error) {
+	b, err := ReadBuffer(r.Reader)
+	return MultiBuffer{b}, err
+}
+
+// PacketReader is a Reader that read one Buffer every time.
+type PacketReader struct {
+	io.Reader
+}
+
+// ReadMultiBuffer implements Reader.
+func (r *PacketReader) ReadMultiBuffer() (MultiBuffer, error) {
+	b, err := readOneUDP(r.Reader)
+	if err != nil {
+		return nil, err
+	}
+	return MultiBuffer{b}, nil
+}