mux io

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)
 		}
 	}

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 {
-	m.access.RLock()
-	defer m.access.RUnlock()
+var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))
 
-	return len(m.sessions) >= maxParallel
+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) IsFullyUsed() bool {
+func (m *Client) isFullyOccupied() bool {
 	m.access.RLock()
 	defer m.access.RUnlock()
 
-	return m.count >= maxTotal
+	return len(m.sessions) >= maxParallel
 }
 
 func (m *Client) remove(id uint16) {
 	m.access.Lock()
+	defer m.access.Unlock()
+
 	delete(m.sessions, id)
-	m.access.Unlock()
+
+	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, session *clientSession) {
+func (m *Client) fetchInput(ctx context.Context, s *session) {
 	dest, _ := proxy.TargetFromContext(ctx)
-	writer := &muxWriter{
+	writer := &MuxWriter{
 		dest:   dest,
-		id:     session.id,
+		id:     s.id,
 		writer: m.inboundRay.InboundInput(),
 	}
 	_, timer := signal.CancelAfterInactivity(ctx, time.Minute*5)
-	buf.PipeUntilEOF(timer, session.outboundRay.OutboundInput(), writer)
+	buf.PipeUntilEOF(timer, s.input, writer)
 	writer.Close()
-	session.closeUplink()
+	s.closeUplink()
 }
 
-func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) {
+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
+				}
+			}
+		}
+	}
+}

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)
+}

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 {

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,

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)
 }