From 73c4696e00bb8237747d8f15d89c904fca951a66 Mon Sep 17 00:00:00 2001 From: Darien Raymond Date: Sun, 2 Apr 2017 09:48:30 +0200 Subject: [PATCH] mux io --- app/proxyman/mux/frame.go | 20 +-- app/proxyman/mux/mux.go | 241 ++++++++++++++++++++++++++++------ app/proxyman/mux/mux_test.go | 44 +++++++ app/proxyman/mux/reader.go | 19 +-- app/proxyman/mux/writer.go | 16 ++- testing/assert/destination.go | 6 + 6 files changed, 283 insertions(+), 63 deletions(-) create mode 100644 app/proxyman/mux/mux_test.go diff --git a/app/proxyman/mux/frame.go b/app/proxyman/mux/frame.go index b86853797..89e02905a 100644 --- a/app/proxyman/mux/frame.go +++ b/app/proxyman/mux/frame.go @@ -1,9 +1,8 @@ package mux import ( - "errors" - "v2ray.com/core/common/buf" + "v2ray.com/core/common/errors" "v2ray.com/core/common/net" "v2ray.com/core/common/serial" ) @@ -38,7 +37,7 @@ type TargetNetwork byte const ( TargetNetworkTCP TargetNetwork = 0x01 - TargetnetworkUDP TargetNetwork = 0x02 + TargetNetworkUDP TargetNetwork = 0x02 ) type AddressType byte @@ -71,7 +70,8 @@ type FrameMetadata struct { func (f FrameMetadata) AsSupplier() buf.Supplier { return func(b []byte) (int, error) { - b = serial.Uint16ToBytes(uint16(0), b) // place holder for length + lengthBytes := b + b = serial.Uint16ToBytes(uint16(0), b[:0]) // place holder for length b = serial.Uint16ToBytes(f.SessionID, b) b = append(b, byte(f.SessionStatus), byte(f.Option)) @@ -82,7 +82,7 @@ func (f FrameMetadata) AsSupplier() buf.Supplier { case net.Network_TCP: b = append(b, byte(TargetNetworkTCP)) case net.Network_UDP: - b = append(b, byte(TargetnetworkUDP)) + b = append(b, byte(TargetNetworkUDP)) } length++ @@ -101,11 +101,12 @@ func (f FrameMetadata) AsSupplier() buf.Supplier { length += 17 case net.AddressFamilyDomain: nDomain := len(addr.Domain()) - b = append(b, byte(nDomain)) + b = append(b, byte(AddressTypeDomain), byte(nDomain)) b = append(b, addr.Domain()...) - length += nDomain + 1 + length += nDomain + 2 } } + serial.Uint16ToBytes(uint16(length), lengthBytes[:0]) return length + 2, nil } } @@ -118,6 +119,7 @@ func ReadFrameFrom(b []byte) (*FrameMetadata, error) { f := &FrameMetadata{ SessionID: serial.BytesToUint16(b[:2]), SessionStatus: SessionStatus(b[2]), + Option: Option(b[3]), } b = b[4:] @@ -140,11 +142,13 @@ func ReadFrameFrom(b []byte) (*FrameMetadata, error) { nDomain := int(b[0]) addr = net.DomainAddress(string(b[1 : 1+nDomain])) b = b[nDomain+1:] + default: + return nil, errors.New("Proxyman|Mux: Unknown address type: ", addrType) } switch network { case TargetNetworkTCP: f.Target = net.TCPDestination(addr, port) - case TargetnetworkUDP: + case TargetNetworkUDP: f.Target = net.UDPDestination(addr, port) } } diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go index 7302a515d..32fb94e44 100644 --- a/app/proxyman/mux/mux.go +++ b/app/proxyman/mux/mux.go @@ -5,7 +5,10 @@ import ( "sync" "time" + "v2ray.com/core/app/dispatcher" + "v2ray.com/core/app/log" "v2ray.com/core/common/buf" + "v2ray.com/core/common/net" "v2ray.com/core/common/signal" "v2ray.com/core/proxy" "v2ray.com/core/transport/ray" @@ -16,16 +19,21 @@ const ( maxTotal = 128 ) -type clientSession struct { +type manager interface { + remove(id uint16) +} + +type session struct { sync.Mutex - outboundRay ray.OutboundRay - parent *Client + input ray.InputStream + output ray.OutputStream + parent manager id uint16 uplinkClosed bool downlinkClosed bool } -func (s *clientSession) checkAndRemove() { +func (s *session) checkAndRemove() { s.Lock() if s.uplinkClosed && s.downlinkClosed { s.parent.remove(s.id) @@ -33,14 +41,14 @@ func (s *clientSession) checkAndRemove() { s.Unlock() } -func (s *clientSession) closeUplink() { +func (s *session) closeUplink() { s.Lock() s.uplinkClosed = true s.Unlock() s.checkAndRemove() } -func (s *clientSession) closeDownlink() { +func (s *session) closeDownlink() { s.Lock() s.downlinkClosed = true s.Unlock() @@ -50,56 +58,101 @@ func (s *clientSession) closeDownlink() { type Client struct { access sync.RWMutex count uint16 - sessions map[uint16]*clientSession + sessions map[uint16]*session inboundRay ray.InboundRay + ctx context.Context + cancel context.CancelFunc } -func (m *Client) IsFullyOccupied() bool { +var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527)) + +func NewClient(p proxy.Outbound, dialer proxy.Dialer) (*Client, error) { + ctx, cancel := context.WithCancel(context.Background()) + ctx = proxy.ContextWithTarget(ctx, muxCoolDestination) + pipe := ray.NewRay(ctx) + err := p.Process(ctx, pipe, dialer) + if err != nil { + cancel() + return nil, err + } + return &Client{ + sessions: make(map[uint16]*session, 256), + inboundRay: pipe, + ctx: ctx, + cancel: cancel, + }, nil +} + +func (m *Client) isFullyOccupied() bool { m.access.RLock() defer m.access.RUnlock() return len(m.sessions) >= maxParallel } -func (m *Client) IsFullyUsed() bool { - m.access.RLock() - defer m.access.RUnlock() - - return m.count >= maxTotal -} - func (m *Client) remove(id uint16) { - m.access.Lock() - delete(m.sessions, id) - m.access.Unlock() -} - -func (m *Client) fetchInput(ctx context.Context, session *clientSession) { - dest, _ := proxy.TargetFromContext(ctx) - writer := &muxWriter{ - dest: dest, - id: session.id, - writer: m.inboundRay.InboundInput(), - } - _, timer := signal.CancelAfterInactivity(ctx, time.Minute*5) - buf.PipeUntilEOF(timer, session.outboundRay.OutboundInput(), writer) - writer.Close() - session.closeUplink() -} - -func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) { m.access.Lock() defer m.access.Unlock() + delete(m.sessions, id) + + if len(m.sessions) == 0 { + m.cancel() + m.inboundRay.InboundInput().Close() + } +} + +func (m *Client) Closed() bool { + select { + case <-m.ctx.Done(): + return true + default: + return false + } +} + +func (m *Client) fetchInput(ctx context.Context, s *session) { + dest, _ := proxy.TargetFromContext(ctx) + writer := &MuxWriter{ + dest: dest, + id: s.id, + writer: m.inboundRay.InboundInput(), + } + _, timer := signal.CancelAfterInactivity(ctx, time.Minute*5) + buf.PipeUntilEOF(timer, s.input, writer) + writer.Close() + s.closeUplink() +} + +func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool { + m.access.Lock() + defer m.access.Unlock() + + if len(m.sessions) >= maxParallel { + return false + } + + if m.count >= maxTotal { + return false + } + + select { + case <-m.ctx.Done(): + return false + default: + } + m.count++ id := m.count - session := &clientSession{ - outboundRay: outboundRay, - parent: m, - id: id, + s := &session{ + input: outboundRay.OutboundInput(), + output: outboundRay.OutboundOutput(), + parent: m, + id: id, } - m.sessions[id] = session - go m.fetchInput(ctx, session) + m.sessions[id] = s + go m.fetchInput(ctx, s) + return true } func (m *Client) fetchOutput() { @@ -110,10 +163,11 @@ func (m *Client) fetchOutput() { break } m.access.RLock() - session, found := m.sessions[meta.SessionID] + s, found := m.sessions[meta.SessionID] m.access.RUnlock() if found && meta.SessionStatus == SessionStatusEnd { - session.closeDownlink() + s.closeDownlink() + s.output.Close() } if !meta.Option.Has(OptionData) { continue @@ -125,7 +179,7 @@ func (m *Client) fetchOutput() { break } if found { - if err := session.outboundRay.OutboundOutput().Write(data); err != nil { + if err := s.output.Write(data); err != nil { break } } @@ -135,3 +189,104 @@ func (m *Client) fetchOutput() { } } } + +type Server struct { + dispatcher dispatcher.Interface +} + +func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (ray.InboundRay, error) { + if dest != muxCoolDestination { + return s.dispatcher.Dispatch(ctx, dest) + } + + ray := ray.NewRay(ctx) + + return ray, nil +} + +type ServerWorker struct { + dispatcher dispatcher.Interface + outboundRay ray.OutboundRay + sessions map[uint16]*session + access sync.RWMutex +} + +func (w *ServerWorker) remove(id uint16) { + w.access.Lock() + delete(w.sessions, id) + w.access.Unlock() +} + +func (w *ServerWorker) handle(ctx context.Context, s *session) { + for { + select { + case <-ctx.Done(): + return + default: + data, err := s.input.Read() + if err != nil { + return + } + w.outboundRay.OutboundOutput().Write(data) + } + } +} + +func (w *ServerWorker) run(ctx context.Context) { + input := w.outboundRay.OutboundInput() + reader := NewReader(input) + for { + select { + case <-ctx.Done(): + return + default: + } + + meta, err := reader.ReadMetadata() + if err != nil { + return + } + + w.access.RLock() + s, found := w.sessions[meta.SessionID] + w.access.RUnlock() + + if found && meta.SessionStatus == SessionStatusEnd { + s.closeUplink() + s.output.Close() + } + + if meta.SessionStatus == SessionStatusNew { + inboundRay, err := w.dispatcher.Dispatch(ctx, meta.Target) + if err != nil { + log.Info("Proxyman|Mux: Failed to dispatch request: ", err) + continue + } + s = &session{ + input: inboundRay.InboundOutput(), + output: inboundRay.InboundInput(), + parent: w, + id: meta.SessionID, + } + w.access.Lock() + w.sessions[meta.SessionID] = s + w.access.Unlock() + go w.handle(ctx, s) + } + + if meta.Option.Has(OptionData) { + for { + data, more, err := reader.Read() + if err != nil { + break + } + if s != nil { + s.output.Write(data) + } + if !more { + break + } + } + } + } +} diff --git a/app/proxyman/mux/mux_test.go b/app/proxyman/mux/mux_test.go new file mode 100644 index 000000000..05a013525 --- /dev/null +++ b/app/proxyman/mux/mux_test.go @@ -0,0 +1,44 @@ +package mux_test + +import ( + "context" + "testing" + + . "v2ray.com/core/app/proxyman/mux" + "v2ray.com/core/common/buf" + "v2ray.com/core/common/net" + "v2ray.com/core/testing/assert" + "v2ray.com/core/transport/ray" +) + +func TestReaderWriter(t *testing.T) { + assert := assert.On(t) + + dest := net.TCPDestination(net.DomainAddress("v2ray.com"), 80) + stream := ray.NewStream(context.Background()) + writer := NewWriter(1, dest, stream) + + payload := buf.New() + payload.AppendBytes('a', 'b', 'c', 'd') + writer.Write(payload) + + writer.Close() + + reader := NewReader(stream) + meta, err := reader.ReadMetadata() + assert.Error(err).IsNil() + assert.Uint16(meta.SessionID).Equals(1) + assert.Byte(byte(meta.SessionStatus)).Equals(byte(SessionStatusNew)) + assert.Destination(meta.Target).Equals(dest) + assert.Byte(byte(meta.Option)).Equals(byte(OptionData)) + + data, more, err := reader.Read() + assert.Error(err).IsNil() + assert.Bool(more).IsFalse() + assert.String(data.String()).Equals("abcd") + + meta, err = reader.ReadMetadata() + assert.Error(err).IsNil() + assert.Byte(byte(meta.SessionStatus)).Equals(byte(SessionStatusEnd)) + assert.Uint16(meta.SessionID).Equals(1) +} diff --git a/app/proxyman/mux/reader.go b/app/proxyman/mux/reader.go index 9a552f115..7c6ff24dc 100644 --- a/app/proxyman/mux/reader.go +++ b/app/proxyman/mux/reader.go @@ -1,23 +1,26 @@ package mux -import "io" -import "v2ray.com/core/common/buf" -import "v2ray.com/core/common/serial" +import ( + "io" -type muxReader struct { + "v2ray.com/core/common/buf" + "v2ray.com/core/common/serial" +) + +type MuxReader struct { reader io.Reader remainingLength int buffer *buf.Buffer } -func NewReader(reader buf.Reader) *muxReader { - return &muxReader{ +func NewReader(reader buf.Reader) *MuxReader { + return &MuxReader{ reader: buf.ToBytesReader(reader), buffer: buf.NewLocal(1024), } } -func (r *muxReader) ReadMetadata() (*FrameMetadata, error) { +func (r *MuxReader) ReadMetadata() (*FrameMetadata, error) { b := r.buffer b.Clear() @@ -32,7 +35,7 @@ func (r *muxReader) ReadMetadata() (*FrameMetadata, error) { return ReadFrameFrom(b.Bytes()) } -func (r *muxReader) Read() (*buf.Buffer, bool, error) { +func (r *MuxReader) Read() (*buf.Buffer, bool, error) { b := buf.New() var dataLen int if r.remainingLength > 0 { diff --git a/app/proxyman/mux/writer.go b/app/proxyman/mux/writer.go index 04ed8037a..62d6980fc 100644 --- a/app/proxyman/mux/writer.go +++ b/app/proxyman/mux/writer.go @@ -6,14 +6,22 @@ import ( "v2ray.com/core/common/serial" ) -type muxWriter struct { +type MuxWriter struct { id uint16 dest net.Destination writer buf.Writer followup bool } -func (w *muxWriter) writeInternal(b *buf.Buffer) error { +func NewWriter(id uint16, dest net.Destination, writer buf.Writer) *MuxWriter { + return &MuxWriter{ + id: id, + dest: dest, + writer: writer, + } +} + +func (w *MuxWriter) writeInternal(b *buf.Buffer) error { meta := FrameMetadata{ SessionID: w.id, Target: w.dest, @@ -49,7 +57,7 @@ func (w *muxWriter) writeInternal(b *buf.Buffer) error { return w.writer.Write(frame) } -func (w *muxWriter) Write(b *buf.Buffer) error { +func (w *MuxWriter) Write(b *buf.Buffer) error { defer b.Release() if err := w.writeInternal(b); err != nil { @@ -63,7 +71,7 @@ func (w *muxWriter) Write(b *buf.Buffer) error { return nil } -func (w *muxWriter) Close() { +func (w *MuxWriter) Close() { meta := FrameMetadata{ SessionID: w.id, Target: w.dest, diff --git a/testing/assert/destination.go b/testing/assert/destination.go index 581c08e07..a4b2e39e7 100644 --- a/testing/assert/destination.go +++ b/testing/assert/destination.go @@ -49,6 +49,12 @@ func (v *DestinationSubject) EqualsString(another string) { } } +func (v *DestinationSubject) Equals(another net.Destination) { + if v.value != another { + v.Fail("not equals to", another.String()) + } +} + func (v *DestinationSubject) HasAddress() *AddressSubject { return v.a.Address(v.value.Address) }