1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-06-18 05:25:23 +00:00
v2fly/common/buf/reader.go
2019-01-06 00:34:38 +01:00

195 lines
4.1 KiB
Go

package buf
import (
"io"
"v2ray.com/core/common"
"v2ray.com/core/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, without allocating large buffer in advance.
func ReadBuffer(r io.Reader) (*Buffer, error) {
// Use an one-byte buffer to wait for incoming payload.
var firstByte [1]byte
nBytes, err := r.Read(firstByte[:])
if err != nil {
return nil, err
}
b := New()
if nBytes > 0 {
common.Must(b.WriteByte(firstByte[0]))
}
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.")
}
// 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)
if err != nil {
return nil, err
}
return MultiBuffer{b}, nil
}
// 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
}