diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go index 5470a2937..fb44919a0 100644 --- a/app/proxyman/mux/mux.go +++ b/app/proxyman/mux/mux.go @@ -192,25 +192,25 @@ func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool return true } -func drain(reader io.Reader) error { +func drain(reader *buf.BufferedReader) error { return buf.Copy(NewStreamReader(reader), buf.Discard) } -func (m *Client) handleStatueKeepAlive(meta *FrameMetadata, reader io.Reader) error { +func (m *Client) handleStatueKeepAlive(meta *FrameMetadata, reader *buf.BufferedReader) error { if meta.Option.Has(OptionData) { return drain(reader) } return nil } -func (m *Client) handleStatusNew(meta *FrameMetadata, reader io.Reader) error { +func (m *Client) handleStatusNew(meta *FrameMetadata, reader *buf.BufferedReader) error { if meta.Option.Has(OptionData) { return drain(reader) } return nil } -func (m *Client) handleStatusKeep(meta *FrameMetadata, reader io.Reader) error { +func (m *Client) handleStatusKeep(meta *FrameMetadata, reader *buf.BufferedReader) error { if !meta.Option.Has(OptionData) { return nil } @@ -221,7 +221,7 @@ func (m *Client) handleStatusKeep(meta *FrameMetadata, reader io.Reader) error { return drain(reader) } -func (m *Client) handleStatusEnd(meta *FrameMetadata, reader io.Reader) error { +func (m *Client) handleStatusEnd(meta *FrameMetadata, reader *buf.BufferedReader) error { if s, found := m.sessionManager.Get(meta.SessionID); found { s.Close() } @@ -315,14 +315,14 @@ func handle(ctx context.Context, s *Session, output buf.Writer) { s.Close() } -func (w *ServerWorker) handleStatusKeepAlive(meta *FrameMetadata, reader io.Reader) error { +func (w *ServerWorker) handleStatusKeepAlive(meta *FrameMetadata, reader *buf.BufferedReader) error { if meta.Option.Has(OptionData) { return drain(reader) } return nil } -func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, reader io.Reader) error { +func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, reader *buf.BufferedReader) error { log.Trace(newError("received request for ", meta.Target)) inboundRay, err := w.dispatcher.Dispatch(ctx, meta.Target) if err != nil { @@ -349,7 +349,7 @@ func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, return nil } -func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader io.Reader) error { +func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader *buf.BufferedReader) error { if !meta.Option.Has(OptionData) { return nil } @@ -359,7 +359,7 @@ func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader io.Reader) e return drain(reader) } -func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader io.Reader) error { +func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader *buf.BufferedReader) error { if s, found := w.sessionManager.Get(meta.SessionID); found { s.Close() } @@ -369,7 +369,7 @@ func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader io.Reader) er return nil } -func (w *ServerWorker) handleFrame(ctx context.Context, reader io.Reader) error { +func (w *ServerWorker) handleFrame(ctx context.Context, reader *buf.BufferedReader) error { meta, err := ReadMetadata(reader) if err != nil { return newError("failed to read metadata").Base(err) diff --git a/app/proxyman/mux/reader.go b/app/proxyman/mux/reader.go index 80192dc9d..df13ed2bf 100644 --- a/app/proxyman/mux/reader.go +++ b/app/proxyman/mux/reader.go @@ -67,12 +67,12 @@ func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) { // StreamReader reads Mux frame as a stream. type StreamReader struct { - reader io.Reader + reader *buf.BufferedReader leftOver int } // NewStreamReader creates a new StreamReader. -func NewStreamReader(reader io.Reader) *StreamReader { +func NewStreamReader(reader *buf.BufferedReader) *StreamReader { return &StreamReader{ reader: reader, leftOver: -1, @@ -94,25 +94,7 @@ func (r *StreamReader) ReadMultiBuffer() (buf.MultiBuffer, error) { r.leftOver = int(size) } - mb := buf.NewMultiBufferCap(32) - for r.leftOver > 0 { - readLen := buf.Size - if r.leftOver < readLen { - readLen = r.leftOver - } - b := buf.New() - if err := b.AppendSupplier(func(bb []byte) (int, error) { - return r.reader.Read(bb[:readLen]) - }); err != nil { - b.Release() - mb.Release() - return nil, err - } - r.leftOver -= b.Len() - mb.Append(b) - if b.Len() < readLen { - break - } - } - return mb, nil + mb, err := r.reader.ReadAtMost(r.leftOver) + r.leftOver -= mb.Len() + return mb, err } diff --git a/app/proxyman/mux/session.go b/app/proxyman/mux/session.go index 8704a0d43..66e1b7290 100644 --- a/app/proxyman/mux/session.go +++ b/app/proxyman/mux/session.go @@ -1,7 +1,6 @@ package mux import ( - "io" "sync" "v2ray.com/core/common/buf" @@ -139,7 +138,7 @@ func (s *Session) Close() { } // NewReader creates a buf.Reader based on the transfer type of this Session. -func (s *Session) NewReader(reader io.Reader) buf.Reader { +func (s *Session) NewReader(reader *buf.BufferedReader) buf.Reader { if s.transferType == protocol.TransferTypeStream { return NewStreamReader(reader) } diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index b2814e5f8..2c8e95858 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -166,6 +166,11 @@ func (mb *MultiBuffer) SliceBySize(size int) MultiBuffer { (*mb)[i] = nil } *mb = (*mb)[endIndex:] + if endIndex == 0 && len(*mb) > 0 { + b := New() + common.Must(b.Reset(ReadFullFrom((*mb)[0], size))) + return NewMultiBufferValue(b) + } return slice } diff --git a/common/buf/reader.go b/common/buf/reader.go index 9b80e9bfe..124e3e3a3 100644 --- a/common/buf/reader.go +++ b/common/buf/reader.go @@ -124,6 +124,20 @@ func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) { return r.stream.ReadMultiBuffer() } +// ReadAtMost returns a MultiBuffer with at most size. +func (r *BufferedReader) ReadAtMost(size int) (mb MultiBuffer, err error) { + if r.leftOver == nil { + r.leftOver, err = r.stream.ReadMultiBuffer() + } + if r.leftOver != nil { + mb = r.leftOver.SliceBySize(size) + if r.leftOver.IsEmpty() { + r.leftOver = nil + } + } + return +} + func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) { mbWriter := NewWriter(writer) totalBytes := int64(0)