1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 18:17:52 -05:00
v2fly/common/buf/reader.go

144 lines
3.0 KiB
Go
Raw Normal View History

2016-12-09 07:17:34 -05:00
package buf
2017-04-17 16:35:20 -04:00
import (
"io"
2018-04-08 17:22:55 -04:00
"v2ray.com/core/common"
2017-04-17 16:35:20 -04:00
"v2ray.com/core/common/errors"
)
2016-12-09 07:17:34 -05:00
func readOne(r io.Reader) (*Buffer, error) {
b := New()
for i := 0; i < 64; i++ {
2018-11-02 10:01:33 -04:00
_, err := b.ReadFrom(r)
if !b.IsEmpty() {
return b, nil
}
if err != nil {
b.Release()
return nil, err
}
}
return nil, newError("Reader returns too many empty payloads.")
}
2018-02-13 05:15:04 -05:00
// BufferedReader is a Reader that keeps its internal buffer.
2017-11-09 16:33:15 -05:00
type BufferedReader struct {
2018-04-20 18:54:53 -04:00
// Reader is the underlying reader to be read from
Reader Reader
// Buffer is the internal buffer to be read from first
Buffer MultiBuffer
2016-12-09 07:17:34 -05:00
}
2018-04-01 17:31:53 -04:00
// BufferedBytes returns the number of bytes that is cached in this reader.
func (r *BufferedReader) BufferedBytes() int32 {
2018-04-20 18:54:53 -04:00
return r.Buffer.Len()
2018-04-01 17:31:53 -04:00
}
2018-02-13 05:15:04 -05:00
// ReadByte implements io.ByteReader.
func (r *BufferedReader) ReadByte() (byte, error) {
var b [1]byte
_, err := r.Read(b[:])
return b[0], err
}
2018-02-13 05:15:04 -05:00
// Read implements io.Reader. It reads from internal buffer first (if available) and then reads from the underlying reader.
2017-11-09 16:33:15 -05:00
func (r *BufferedReader) Read(b []byte) (int, error) {
2018-07-29 07:19:03 -04:00
if !r.Buffer.IsEmpty() {
buffer, nBytes := SplitBytes(r.Buffer, b)
2018-11-18 13:36:36 -05:00
r.Buffer = buffer
2018-04-20 18:54:53 -04:00
if r.Buffer.IsEmpty() {
r.Buffer = nil
2017-04-19 05:52:27 -04:00
}
2017-04-19 15:27:21 -04:00
return nBytes, nil
2016-12-09 07:17:34 -05:00
}
2018-04-20 18:54:53 -04:00
mb, err := r.Reader.ReadMultiBuffer()
2018-07-29 07:19:03 -04:00
if err != nil {
return 0, err
}
mb, nBytes := SplitBytes(mb, b)
2018-07-29 07:19:03 -04:00
if !mb.IsEmpty() {
r.Buffer = mb
2016-12-09 07:17:34 -05:00
}
2018-11-18 13:36:36 -05:00
return nBytes, nil
2016-12-09 07:17:34 -05:00
}
2017-04-15 16:22:29 -04:00
2018-02-13 05:15:04 -05:00
// ReadMultiBuffer implements Reader.
2017-11-09 16:33:15 -05:00
func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) {
2018-07-29 07:19:03 -04:00
if !r.Buffer.IsEmpty() {
2018-04-20 18:54:53 -04:00
mb := r.Buffer
r.Buffer = nil
2017-04-19 05:52:27 -04:00
return mb, nil
2017-04-15 16:22:29 -04:00
}
2017-04-19 05:52:27 -04:00
2018-04-20 18:54:53 -04:00
return r.Reader.ReadMultiBuffer()
2017-04-15 16:22:29 -04:00
}
2017-04-17 16:35:20 -04:00
// ReadAtMost returns a MultiBuffer with at most size.
2018-04-02 14:00:50 -04:00
func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
2018-07-29 07:19:03 -04:00
if r.Buffer.IsEmpty() {
2018-04-20 18:54:53 -04:00
mb, err := r.Reader.ReadMultiBuffer()
2017-11-26 09:08:37 -05:00
if mb.IsEmpty() && err != nil {
return nil, err
}
2018-04-20 18:54:53 -04:00
r.Buffer = mb
2017-11-26 09:08:37 -05:00
}
2018-11-18 14:16:14 -05:00
rb, mb := SplitSize(r.Buffer, size)
r.Buffer = rb
2018-04-20 18:54:53 -04:00
if r.Buffer.IsEmpty() {
r.Buffer = nil
}
2017-11-26 09:08:37 -05:00
return mb, nil
}
2017-11-09 16:33:15 -05:00
func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) {
2017-04-17 16:35:20 -04:00
mbWriter := NewWriter(writer)
2018-10-24 07:51:04 -04:00
var sc SizeCounter
2018-04-20 18:54:53 -04:00
if r.Buffer != nil {
2018-10-24 07:51:04 -04:00
sc.Size = int64(r.Buffer.Len())
2018-04-20 18:54:53 -04:00
if err := mbWriter.WriteMultiBuffer(r.Buffer); err != nil {
2017-04-19 10:29:36 -04:00
return 0, err
}
2018-04-20 18:54:53 -04:00
r.Buffer = nil
2017-04-19 05:52:27 -04:00
}
2018-10-24 07:51:04 -04:00
err := Copy(r.Reader, mbWriter, CountSize(&sc))
return sc.Size, err
2017-04-17 16:35:20 -04:00
}
2018-02-13 05:15:04 -05:00
// WriteTo implements io.WriterTo.
2017-11-09 16:33:15 -05:00
func (r *BufferedReader) WriteTo(writer io.Writer) (int64, error) {
2017-04-17 16:35:20 -04:00
nBytes, err := r.writeToInternal(writer)
if errors.Cause(err) == io.EOF {
return nBytes, nil
}
return nBytes, err
}
2018-04-16 08:57:13 -04:00
// Close implements io.Closer.
func (r *BufferedReader) Close() error {
2018-04-20 18:54:53 -04:00
if !r.Buffer.IsEmpty() {
2018-11-17 16:45:07 -05:00
ReleaseMulti(r.Buffer)
r.Buffer = nil
2018-04-16 08:57:13 -04:00
}
2018-04-20 18:54:53 -04:00
return common.Close(r.Reader)
2018-04-16 08:57:13 -04:00
}
2018-10-24 09:02:19 -04:00
// SingleReader is a Reader that read one Buffer every time.
type SingleReader struct {
io.Reader
}
2018-10-24 09:02:19 -04:00
// ReadMultiBuffer implements Reader.
func (r *SingleReader) ReadMultiBuffer() (MultiBuffer, error) {
b, err := readOne(r.Reader)
if err != nil {
return nil, err
}
2018-11-16 05:29:16 -05:00
return MultiBuffer{b}, nil
}