1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2025-01-04 16:37:12 -05:00

simplify buf.BufferedReader

This commit is contained in:
Darien Raymond 2018-04-21 00:54:53 +02:00
parent bd6c6b5361
commit 148a7d064d
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
15 changed files with 56 additions and 71 deletions

View File

@ -258,7 +258,7 @@ func (m *Client) fetchOutput() {
common.Must(m.done.Close()) common.Must(m.done.Close())
}() }()
reader := buf.NewBufferedReader(m.link.Reader) reader := &buf.BufferedReader{Reader: m.link.Reader}
for { for {
meta, err := ReadMetadata(reader) meta, err := ReadMetadata(reader)
@ -456,7 +456,7 @@ func (w *ServerWorker) handleFrame(ctx context.Context, reader *buf.BufferedRead
func (w *ServerWorker) run(ctx context.Context) { func (w *ServerWorker) run(ctx context.Context) {
input := w.link.Reader input := w.link.Reader
reader := buf.NewBufferedReader(input) reader := &buf.BufferedReader{Reader: input}
defer w.sessionManager.Close() defer w.sessionManager.Close()

View File

@ -59,7 +59,7 @@ func TestReaderWriter(t *testing.T) {
assert(writePayload(writer2, 'y'), IsNil) assert(writePayload(writer2, 'y'), IsNil)
writer2.Close() writer2.Close()
bytesReader := buf.NewBufferedReader(pReader) bytesReader := &buf.BufferedReader{Reader: pReader}
meta, err := ReadMetadata(bytesReader) meta, err := ReadMetadata(bytesReader)
assert(err, IsNil) assert(err, IsNil)

View File

@ -75,32 +75,17 @@ func (r *BytesToBufferReader) ReadMultiBuffer() (MultiBuffer, error) {
// BufferedReader is a Reader that keeps its internal buffer. // BufferedReader is a Reader that keeps its internal buffer.
type BufferedReader struct { type BufferedReader struct {
stream Reader // Reader is the underlying reader to be read from
leftOver MultiBuffer Reader Reader
buffered bool // Buffer is the internal buffer to be read from first
} Buffer MultiBuffer
// Direct indicates whether or not to use the internal buffer
// NewBufferedReader returns a new BufferedReader. Direct bool
func NewBufferedReader(reader Reader) *BufferedReader {
return &BufferedReader{
stream: reader,
buffered: true,
}
}
// SetBuffered sets whether to keep the interal buffer.
func (r *BufferedReader) SetBuffered(f bool) {
r.buffered = f
}
// IsBuffered returns true if internal buffer is used.
func (r *BufferedReader) IsBuffered() bool {
return r.buffered
} }
// BufferedBytes returns the number of bytes that is cached in this reader. // BufferedBytes returns the number of bytes that is cached in this reader.
func (r *BufferedReader) BufferedBytes() int32 { func (r *BufferedReader) BufferedBytes() int32 {
return r.leftOver.Len() return r.Buffer.Len()
} }
// ReadByte implements io.ByteReader. // ReadByte implements io.ByteReader.
@ -112,26 +97,26 @@ func (r *BufferedReader) ReadByte() (byte, error) {
// Read implements io.Reader. It reads from internal buffer first (if available) and then reads from the underlying reader. // 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) { func (r *BufferedReader) Read(b []byte) (int, error) {
if r.leftOver != nil { if r.Buffer != nil {
nBytes, _ := r.leftOver.Read(b) nBytes, _ := r.Buffer.Read(b)
if r.leftOver.IsEmpty() { if r.Buffer.IsEmpty() {
r.leftOver.Release() r.Buffer.Release()
r.leftOver = nil r.Buffer = nil
} }
return nBytes, nil return nBytes, nil
} }
if !r.buffered { if r.Direct {
if reader, ok := r.stream.(io.Reader); ok { if reader, ok := r.Reader.(io.Reader); ok {
return reader.Read(b) return reader.Read(b)
} }
} }
mb, err := r.stream.ReadMultiBuffer() mb, err := r.Reader.ReadMultiBuffer()
if mb != nil { if mb != nil {
nBytes, _ := mb.Read(b) nBytes, _ := mb.Read(b)
if !mb.IsEmpty() { if !mb.IsEmpty() {
r.leftOver = mb r.Buffer = mb
} }
return nBytes, err return nBytes, err
} }
@ -140,28 +125,28 @@ func (r *BufferedReader) Read(b []byte) (int, error) {
// ReadMultiBuffer implements Reader. // ReadMultiBuffer implements Reader.
func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) { func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) {
if r.leftOver != nil { if r.Buffer != nil {
mb := r.leftOver mb := r.Buffer
r.leftOver = nil r.Buffer = nil
return mb, nil return mb, nil
} }
return r.stream.ReadMultiBuffer() return r.Reader.ReadMultiBuffer()
} }
// ReadAtMost returns a MultiBuffer with at most size. // ReadAtMost returns a MultiBuffer with at most size.
func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) { func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
if r.leftOver == nil { if r.Buffer == nil {
mb, err := r.stream.ReadMultiBuffer() mb, err := r.Reader.ReadMultiBuffer()
if mb.IsEmpty() && err != nil { if mb.IsEmpty() && err != nil {
return nil, err return nil, err
} }
r.leftOver = mb r.Buffer = mb
} }
mb := r.leftOver.SliceBySize(size) mb := r.Buffer.SliceBySize(size)
if r.leftOver.IsEmpty() { if r.Buffer.IsEmpty() {
r.leftOver = nil r.Buffer = nil
} }
return mb, nil return mb, nil
} }
@ -169,16 +154,16 @@ func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) { func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) {
mbWriter := NewWriter(writer) mbWriter := NewWriter(writer)
totalBytes := int64(0) totalBytes := int64(0)
if r.leftOver != nil { if r.Buffer != nil {
totalBytes += int64(r.leftOver.Len()) totalBytes += int64(r.Buffer.Len())
if err := mbWriter.WriteMultiBuffer(r.leftOver); err != nil { if err := mbWriter.WriteMultiBuffer(r.Buffer); err != nil {
return 0, err return 0, err
} }
r.leftOver = nil r.Buffer = nil
} }
for { for {
mb, err := r.stream.ReadMultiBuffer() mb, err := r.Reader.ReadMultiBuffer()
if mb != nil { if mb != nil {
totalBytes += int64(mb.Len()) totalBytes += int64(mb.Len())
if werr := mbWriter.WriteMultiBuffer(mb); werr != nil { if werr := mbWriter.WriteMultiBuffer(mb); werr != nil {
@ -202,8 +187,8 @@ func (r *BufferedReader) WriteTo(writer io.Writer) (int64, error) {
// Close implements io.Closer. // Close implements io.Closer.
func (r *BufferedReader) Close() error { func (r *BufferedReader) Close() error {
if !r.leftOver.IsEmpty() { if !r.Buffer.IsEmpty() {
r.leftOver.Release() r.Buffer.Release()
} }
return common.Close(r.stream) return common.Close(r.Reader)
} }

View File

@ -39,7 +39,7 @@ func TestBytesReaderWriteTo(t *testing.T) {
assert := With(t) assert := With(t)
pReader, pWriter := pipe.New() pReader, pWriter := pipe.New()
reader := NewBufferedReader(pReader) reader := &BufferedReader{Reader: pReader}
b1 := New() b1 := New()
b1.AppendBytes('a', 'b', 'c') b1.AppendBytes('a', 'b', 'c')
b2 := New() b2 := New()
@ -66,7 +66,7 @@ func TestBytesReaderMultiBuffer(t *testing.T) {
assert := With(t) assert := With(t)
pReader, pWriter := pipe.New() pReader, pWriter := pipe.New()
reader := NewBufferedReader(pReader) reader := &BufferedReader{Reader: pReader}
b1 := New() b1 := New()
b1.AppendBytes('a', 'b', 'c') b1.AppendBytes('a', 'b', 'c')
b2 := New() b2 := New()

View File

@ -67,7 +67,7 @@ func TestDiscardBytesMultiBuffer(t *testing.T) {
common.Must2(buffer.ReadFrom(io.LimitReader(rand.Reader, size))) common.Must2(buffer.ReadFrom(io.LimitReader(rand.Reader, size)))
r := NewReader(buffer) r := NewReader(buffer)
nBytes, err := io.Copy(DiscardBytes, NewBufferedReader(r)) nBytes, err := io.Copy(DiscardBytes, &BufferedReader{Reader: r})
assert(nBytes, Equals, int64(size)) assert(nBytes, Equals, int64(size))
assert(err, IsNil) assert(err, IsNil)
} }

View File

@ -91,7 +91,7 @@ type AuthenticationReader struct {
func NewAuthenticationReader(auth Authenticator, sizeParser ChunkSizeDecoder, reader io.Reader, transferType protocol.TransferType) *AuthenticationReader { func NewAuthenticationReader(auth Authenticator, sizeParser ChunkSizeDecoder, reader io.Reader, transferType protocol.TransferType) *AuthenticationReader {
return &AuthenticationReader{ return &AuthenticationReader{
auth: auth, auth: auth,
reader: buf.NewBufferedReader(buf.NewReader(reader)), reader: &buf.BufferedReader{Reader: buf.NewReader(reader)},
sizeParser: sizeParser, sizeParser: sizeParser,
transferType: transferType, transferType: transferType,
size: -1, size: -1,

View File

@ -68,7 +68,7 @@ type ChunkStreamReader struct {
func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader { func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader {
return &ChunkStreamReader{ return &ChunkStreamReader{
sizeDecoder: sizeDecoder, sizeDecoder: sizeDecoder,
reader: buf.NewBufferedReader(buf.NewReader(reader)), reader: &buf.BufferedReader{Reader: buf.NewReader(reader)},
buffer: make([]byte, sizeDecoder.SizeBytes()), buffer: make([]byte, sizeDecoder.SizeBytes()),
} }
} }

View File

@ -38,13 +38,13 @@ func ConnectionInputMulti(writer buf.Writer) ConnectionOption {
func ConnectionOutput(reader io.Reader) ConnectionOption { func ConnectionOutput(reader io.Reader) ConnectionOption {
return func(c *connection) { return func(c *connection) {
c.reader = buf.NewBufferedReader(buf.NewReader(reader)) c.reader = &buf.BufferedReader{Reader: buf.NewReader(reader)}
} }
} }
func ConnectionOutputMulti(reader buf.Reader) ConnectionOption { func ConnectionOutputMulti(reader buf.Reader) ConnectionOption {
return func(c *connection) { return func(c *connection) {
c.reader = buf.NewBufferedReader(reader) c.reader = &buf.BufferedReader{Reader: reader}
} }
} }

View File

@ -268,7 +268,7 @@ func (s *Server) handlePlainHTTP(ctx context.Context, request *http.Request, wri
} }
responseDone := func() error { responseDone := func() error {
responseReader := bufio.NewReaderSize(buf.NewBufferedReader(link.Reader), buf.Size) responseReader := bufio.NewReaderSize(&buf.BufferedReader{Reader: link.Reader}, buf.Size)
response, err := http.ReadResponse(responseReader, request) response, err := http.ReadResponse(responseReader, request)
if err == nil { if err == nil {
http_proto.RemoveHopByHopHeaders(response.Header) http_proto.RemoveHopByHopHeaders(response.Header)

View File

@ -52,7 +52,7 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
if err != nil { if err != nil {
return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError() return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
} }
br := buf.NewBufferedReader(r) br := &buf.BufferedReader{Reader: r}
reader = nil reader = nil
authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv)) authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
@ -109,7 +109,7 @@ func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHea
return nil, nil, newError("invalid remote address.") return nil, nil, newError("invalid remote address.")
} }
br.SetBuffered(false) br.Direct = true
var chunkReader buf.Reader var chunkReader buf.Reader
if request.Option.Has(RequestOptionOneTimeAuth) { if request.Option.Has(RequestOptionOneTimeAuth) {

View File

@ -140,8 +140,8 @@ func (s *Server) handlerUDPPayload(ctx context.Context, conn internet.Connection
func (s *Server) handleConnection(ctx context.Context, conn internet.Connection, dispatcher core.Dispatcher) error { func (s *Server) handleConnection(ctx context.Context, conn internet.Connection, dispatcher core.Dispatcher) error {
sessionPolicy := s.v.PolicyManager().ForLevel(s.user.Level) sessionPolicy := s.v.PolicyManager().ForLevel(s.user.Level)
conn.SetReadDeadline(time.Now().Add(sessionPolicy.Timeouts.Handshake)) conn.SetReadDeadline(time.Now().Add(sessionPolicy.Timeouts.Handshake))
bufferedReader := buf.NewBufferedReader(buf.NewReader(conn)) bufferedReader := buf.BufferedReader{Reader: buf.NewReader(conn)}
request, bodyReader, err := ReadTCPSession(s.user, bufferedReader) request, bodyReader, err := ReadTCPSession(s.user, &bufferedReader)
if err != nil { if err != nil {
log.Record(&log.AccessMessage{ log.Record(&log.AccessMessage{
From: conn.RemoteAddr(), From: conn.RemoteAddr(),
@ -153,7 +153,7 @@ func (s *Server) handleConnection(ctx context.Context, conn internet.Connection,
} }
conn.SetReadDeadline(time.Time{}) conn.SetReadDeadline(time.Time{})
bufferedReader.SetBuffered(false) bufferedReader.Direct = true
dest := request.Destination() dest := request.Destination()
log.Record(&log.AccessMessage{ log.Record(&log.AccessMessage{

View File

@ -70,7 +70,7 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
newError("failed to set deadline").Base(err).WithContext(ctx).WriteToLog() newError("failed to set deadline").Base(err).WithContext(ctx).WriteToLog()
} }
reader := buf.NewBufferedReader(buf.NewReader(conn)) reader := &buf.BufferedReader{Reader: buf.NewReader(conn)}
inboundDest, ok := proxy.InboundEntryPointFromContext(ctx) inboundDest, ok := proxy.InboundEntryPointFromContext(ctx)
if !ok { if !ok {

View File

@ -224,7 +224,7 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection i
return newError("unable to set read deadline").Base(err).AtWarning() return newError("unable to set read deadline").Base(err).AtWarning()
} }
reader := buf.NewBufferedReader(buf.NewReader(connection)) reader := &buf.BufferedReader{Reader: buf.NewReader(connection)}
session := encoding.NewServerSession(h.clients, h.sessionHistory) session := encoding.NewServerSession(h.clients, h.sessionHistory)
request, err := session.DecodeRequestHeader(reader) request, err := session.DecodeRequestHeader(reader)

View File

@ -146,14 +146,14 @@ func (v *Handler) Process(ctx context.Context, link *core.Link, dialer proxy.Dia
responseDone := func() error { responseDone := func() error {
defer timer.SetTimeout(sessionPolicy.Timeouts.UplinkOnly) defer timer.SetTimeout(sessionPolicy.Timeouts.UplinkOnly)
reader := buf.NewBufferedReader(buf.NewReader(conn)) reader := &buf.BufferedReader{Reader: buf.NewReader(conn)}
header, err := session.DecodeResponseHeader(reader) header, err := session.DecodeResponseHeader(reader)
if err != nil { if err != nil {
return newError("failed to read header").Base(err) return newError("failed to read header").Base(err)
} }
v.handleCommand(rec.Destination(), header.Command) v.handleCommand(rec.Destination(), header.Command)
reader.SetBuffered(false) reader.Direct = true
bodyReader := session.DecodeResponseBody(request, reader) bodyReader := session.DecodeResponseBody(request, reader)
return buf.Copy(bodyReader, output, buf.UpdateActivity(timer)) return buf.Copy(bodyReader, output, buf.UpdateActivity(timer))

View File

@ -84,11 +84,11 @@ func Dial(ctx context.Context, dest net.Destination) (internet.Connection, error
} }
preader, pwriter := pipe.New(pipe.WithSizeLimit(20 * 1024)) preader, pwriter := pipe.New(pipe.WithSizeLimit(20 * 1024))
breader := buf.NewBufferedReader(preader) breader := &buf.BufferedReader{Reader: preader}
request := &http.Request{ request := &http.Request{
Method: "PUT", Method: "PUT",
Host: httpSettings.getRandomHost(), Host: httpSettings.getRandomHost(),
Body: buf.NewBufferedReader(preader), Body: breader,
URL: &url.URL{ URL: &url.URL{
Scheme: "https", Scheme: "https",
Host: dest.NetAddr(), Host: dest.NetAddr(),