mirror of
https://github.com/v2fly/v2ray-core.git
synced 2024-12-21 09:36:34 -05:00
simplify ray stream
This commit is contained in:
parent
d809973621
commit
2f565bfd5e
@ -1,34 +0,0 @@
|
||||
package buf
|
||||
|
||||
type MergingReader struct {
|
||||
reader Reader
|
||||
timeoutReader TimeoutReader
|
||||
}
|
||||
|
||||
func NewMergingReader(reader Reader) Reader {
|
||||
return &MergingReader{
|
||||
reader: reader,
|
||||
timeoutReader: reader.(TimeoutReader),
|
||||
}
|
||||
}
|
||||
|
||||
func (r *MergingReader) Read() (MultiBuffer, error) {
|
||||
mb, err := r.reader.Read()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if r.timeoutReader == nil {
|
||||
return mb, nil
|
||||
}
|
||||
|
||||
for {
|
||||
mb2, err := r.timeoutReader.ReadTimeout(0)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
mb.AppendMulti(mb2)
|
||||
}
|
||||
|
||||
return mb, nil
|
||||
}
|
@ -1,33 +0,0 @@
|
||||
package buf_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"context"
|
||||
|
||||
. "v2ray.com/core/common/buf"
|
||||
"v2ray.com/core/testing/assert"
|
||||
"v2ray.com/core/transport/ray"
|
||||
)
|
||||
|
||||
func TestMergingReader(t *testing.T) {
|
||||
assert := assert.On(t)
|
||||
|
||||
stream := ray.NewStream(context.Background())
|
||||
b1 := New()
|
||||
b1.AppendBytes('a', 'b', 'c')
|
||||
stream.Write(NewMultiBufferValue(b1))
|
||||
|
||||
b2 := New()
|
||||
b2.AppendBytes('e', 'f', 'g')
|
||||
stream.Write(NewMultiBufferValue(b2))
|
||||
|
||||
b3 := New()
|
||||
b3.AppendBytes('h', 'i', 'j')
|
||||
stream.Write(NewMultiBufferValue(b3))
|
||||
|
||||
reader := NewMergingReader(stream)
|
||||
b, err := reader.Read()
|
||||
assert.Error(err).IsNil()
|
||||
assert.Int(b.Len()).Equals(9)
|
||||
}
|
@ -13,7 +13,7 @@ type MultiBufferReader interface {
|
||||
type MultiBuffer []*Buffer
|
||||
|
||||
func NewMultiBuffer() MultiBuffer {
|
||||
return MultiBuffer(make([]*Buffer, 0, 8))
|
||||
return MultiBuffer(make([]*Buffer, 0, 32))
|
||||
}
|
||||
|
||||
func NewMultiBufferValue(b ...*Buffer) MultiBuffer {
|
||||
|
@ -4,6 +4,7 @@ package freedom
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
@ -112,8 +113,13 @@ func (v *Handler) Process(ctx context.Context, outboundRay ray.OutboundRay, dial
|
||||
ctx, timer := signal.CancelAfterInactivity(ctx, timeout)
|
||||
|
||||
requestDone := signal.ExecuteAsync(func() error {
|
||||
v2writer := buf.NewWriter(conn)
|
||||
if err := buf.PipeUntilEOF(timer, input, v2writer); err != nil {
|
||||
var writer buf.Writer
|
||||
if destination.Network == net.Network_TCP {
|
||||
writer = buf.NewWriter(conn)
|
||||
} else {
|
||||
writer = &seqWriter{writer: conn}
|
||||
}
|
||||
if err := buf.PipeUntilEOF(timer, input, writer); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
@ -145,3 +151,19 @@ func init() {
|
||||
return New(ctx, config.(*Config))
|
||||
}))
|
||||
}
|
||||
|
||||
type seqWriter struct {
|
||||
writer io.Writer
|
||||
}
|
||||
|
||||
func (w *seqWriter) Write(mb buf.MultiBuffer) error {
|
||||
defer mb.Release()
|
||||
|
||||
for _, b := range mb {
|
||||
if _, err := w.writer.Write(b.Bytes()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
@ -105,8 +105,7 @@ func (v *Client) Process(ctx context.Context, outboundRay ray.OutboundRay, diale
|
||||
}
|
||||
|
||||
requestDone := signal.ExecuteAsync(func() error {
|
||||
mergedInput := buf.NewMergingReader(outboundRay.OutboundInput())
|
||||
if err := buf.PipeUntilEOF(timer, mergedInput, bodyWriter); err != nil {
|
||||
if err := buf.PipeUntilEOF(timer, outboundRay.OutboundInput(), bodyWriter); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
|
@ -160,8 +160,7 @@ func (s *Server) handleConnection(ctx context.Context, conn internet.Connection,
|
||||
return newError("failed to write response").Base(err)
|
||||
}
|
||||
|
||||
mergeReader := buf.NewMergingReader(ray.InboundOutput())
|
||||
payload, err := mergeReader.Read()
|
||||
payload, err := ray.InboundOutput().Read()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@ -174,7 +173,7 @@ func (s *Server) handleConnection(ctx context.Context, conn internet.Connection,
|
||||
return err
|
||||
}
|
||||
|
||||
if err := buf.PipeUntilEOF(timer, mergeReader, responseWriter); err != nil {
|
||||
if err := buf.PipeUntilEOF(timer, ray.InboundOutput(), responseWriter); err != nil {
|
||||
return newError("failed to transport all TCP response").Base(err)
|
||||
}
|
||||
|
||||
|
@ -140,12 +140,8 @@ func transferResponse(timer signal.ActivityTimer, session *encoding.ServerSessio
|
||||
|
||||
bodyWriter := session.EncodeResponseBody(request, output)
|
||||
|
||||
var reader buf.Reader = input
|
||||
if request.Command == protocol.RequestCommandTCP {
|
||||
reader = buf.NewMergingReader(input)
|
||||
}
|
||||
// Optimize for small response packet
|
||||
data, err := reader.Read()
|
||||
data, err := input.Read()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@ -161,7 +157,7 @@ func transferResponse(timer signal.ActivityTimer, session *encoding.ServerSessio
|
||||
}
|
||||
}
|
||||
|
||||
if err := buf.PipeUntilEOF(timer, reader, bodyWriter); err != nil {
|
||||
if err := buf.PipeUntilEOF(timer, input, bodyWriter); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
|
@ -123,12 +123,7 @@ func (v *Handler) Process(ctx context.Context, outboundRay ray.OutboundRay, dial
|
||||
return err
|
||||
}
|
||||
|
||||
var inputReader buf.Reader = input
|
||||
if request.Command == protocol.RequestCommandTCP {
|
||||
inputReader = buf.NewMergingReader(input)
|
||||
}
|
||||
|
||||
if err := buf.PipeUntilEOF(timer, inputReader, bodyWriter); err != nil {
|
||||
if err := buf.PipeUntilEOF(timer, input, bodyWriter); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
|
@ -3,15 +3,12 @@ package ray
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"v2ray.com/core/common/buf"
|
||||
)
|
||||
|
||||
const (
|
||||
bufferSize = 512
|
||||
)
|
||||
|
||||
// NewRay creates a new Ray for direct traffic transport.
|
||||
func NewRay(ctx context.Context) Ray {
|
||||
return &directRay{
|
||||
@ -42,121 +39,132 @@ func (v *directRay) InboundOutput() InputStream {
|
||||
}
|
||||
|
||||
type Stream struct {
|
||||
buffer chan buf.MultiBuffer
|
||||
access sync.Mutex
|
||||
data buf.MultiBuffer
|
||||
ctx context.Context
|
||||
close chan bool
|
||||
err chan bool
|
||||
wakeup chan bool
|
||||
close bool
|
||||
err bool
|
||||
}
|
||||
|
||||
func NewStream(ctx context.Context) *Stream {
|
||||
return &Stream{
|
||||
ctx: ctx,
|
||||
buffer: make(chan buf.MultiBuffer, bufferSize),
|
||||
close: make(chan bool),
|
||||
err: make(chan bool),
|
||||
wakeup: make(chan bool, 1),
|
||||
}
|
||||
}
|
||||
|
||||
func (v *Stream) Read() (buf.MultiBuffer, error) {
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
func (s *Stream) getData() (buf.MultiBuffer, error) {
|
||||
s.access.Lock()
|
||||
defer s.access.Unlock()
|
||||
|
||||
if s.data != nil {
|
||||
mb := s.data
|
||||
s.data = nil
|
||||
return mb, nil
|
||||
}
|
||||
|
||||
if s.close {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
if s.err {
|
||||
return nil, io.ErrClosedPipe
|
||||
case <-v.err:
|
||||
return nil, io.ErrClosedPipe
|
||||
case b := <-v.buffer:
|
||||
return b, nil
|
||||
default:
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *Stream) Read() (buf.MultiBuffer, error) {
|
||||
for {
|
||||
mb, err := s.getData()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if mb != nil {
|
||||
return mb, nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
return nil, io.ErrClosedPipe
|
||||
case b := <-v.buffer:
|
||||
return b, nil
|
||||
case <-v.close:
|
||||
return nil, io.EOF
|
||||
case <-v.err:
|
||||
case <-s.ctx.Done():
|
||||
return nil, io.ErrClosedPipe
|
||||
case <-s.wakeup:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (v *Stream) ReadTimeout(timeout time.Duration) (buf.MultiBuffer, error) {
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
return nil, io.ErrClosedPipe
|
||||
case <-v.err:
|
||||
return nil, io.ErrClosedPipe
|
||||
case b := <-v.buffer:
|
||||
return b, nil
|
||||
default:
|
||||
if timeout == 0 {
|
||||
return nil, buf.ErrReadTimeout
|
||||
func (s *Stream) ReadTimeout(timeout time.Duration) (buf.MultiBuffer, error) {
|
||||
for {
|
||||
mb, err := s.getData()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if mb != nil {
|
||||
return mb, nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
return nil, io.ErrClosedPipe
|
||||
case b := <-v.buffer:
|
||||
return b, nil
|
||||
case <-v.close:
|
||||
return nil, io.EOF
|
||||
case <-v.err:
|
||||
case <-s.ctx.Done():
|
||||
return nil, io.ErrClosedPipe
|
||||
case <-time.After(timeout):
|
||||
return nil, buf.ErrReadTimeout
|
||||
case <-s.wakeup:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (v *Stream) Write(data buf.MultiBuffer) (err error) {
|
||||
func (s *Stream) Write(data buf.MultiBuffer) (err error) {
|
||||
if data.IsEmpty() {
|
||||
return
|
||||
}
|
||||
|
||||
s.access.Lock()
|
||||
defer s.access.Unlock()
|
||||
|
||||
if s.err {
|
||||
data.Release()
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
if s.close {
|
||||
data.Release()
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
|
||||
if s.data == nil {
|
||||
s.data = data
|
||||
} else {
|
||||
s.data.AppendMulti(data)
|
||||
}
|
||||
s.wakeUp()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Stream) wakeUp() {
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
return io.ErrClosedPipe
|
||||
case <-v.err:
|
||||
return io.ErrClosedPipe
|
||||
case <-v.close:
|
||||
return io.ErrClosedPipe
|
||||
case s.wakeup <- true:
|
||||
default:
|
||||
select {
|
||||
case <-v.ctx.Done():
|
||||
return io.ErrClosedPipe
|
||||
case <-v.err:
|
||||
return io.ErrClosedPipe
|
||||
case <-v.close:
|
||||
return io.ErrClosedPipe
|
||||
case v.buffer <- data:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (v *Stream) Close() {
|
||||
defer swallowPanic()
|
||||
|
||||
close(v.close)
|
||||
func (s *Stream) Close() {
|
||||
s.access.Lock()
|
||||
s.close = true
|
||||
s.wakeUp()
|
||||
s.access.Unlock()
|
||||
}
|
||||
|
||||
func (v *Stream) CloseError() {
|
||||
defer swallowPanic()
|
||||
|
||||
close(v.err)
|
||||
|
||||
n := len(v.buffer)
|
||||
for i := 0; i < n; i++ {
|
||||
select {
|
||||
case b := <-v.buffer:
|
||||
b.Release()
|
||||
default:
|
||||
return
|
||||
}
|
||||
func (s *Stream) CloseError() {
|
||||
s.access.Lock()
|
||||
s.err = true
|
||||
if s.data != nil {
|
||||
s.data.Release()
|
||||
s.data = nil
|
||||
}
|
||||
s.wakeUp()
|
||||
s.access.Unlock()
|
||||
}
|
||||
|
||||
func (v *Stream) Release() {}
|
||||
|
||||
func swallowPanic() {
|
||||
recover()
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user