|  | @@ -9,8 +9,10 @@ import (
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  	"v2ray.com/core"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/app/proxyman"
 | 
	
		
			
				|  |  | +	"v2ray.com/core/common"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/common/buf"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/common/net"
 | 
	
		
			
				|  |  | +	"v2ray.com/core/common/signal"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/proxy"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/transport/internet"
 | 
	
		
			
				|  |  |  	"v2ray.com/core/transport/internet/tcp"
 | 
	
	
		
			
				|  | @@ -19,7 +21,7 @@ import (
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  type worker interface {
 | 
	
		
			
				|  |  |  	Start() error
 | 
	
		
			
				|  |  | -	Close()
 | 
	
		
			
				|  |  | +	Close() error
 | 
	
		
			
				|  |  |  	Port() net.Port
 | 
	
		
			
				|  |  |  	Proxy() proxy.Inbound
 | 
	
		
			
				|  |  |  }
 | 
	
	
		
			
				|  | @@ -34,13 +36,11 @@ type tcpWorker struct {
 | 
	
		
			
				|  |  |  	dispatcher   core.Dispatcher
 | 
	
		
			
				|  |  |  	sniffers     []proxyman.KnownProtocols
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -	ctx    context.Context
 | 
	
		
			
				|  |  | -	cancel context.CancelFunc
 | 
	
		
			
				|  |  | -	hub    internet.Listener
 | 
	
		
			
				|  |  | +	hub internet.Listener
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (w *tcpWorker) callback(conn internet.Connection) {
 | 
	
		
			
				|  |  | -	ctx, cancel := context.WithCancel(w.ctx)
 | 
	
		
			
				|  |  | +	ctx, cancel := context.WithCancel(context.Background())
 | 
	
		
			
				|  |  |  	if w.recvOrigDest {
 | 
	
		
			
				|  |  |  		dest, err := tcp.GetOriginalDestination(conn)
 | 
	
		
			
				|  |  |  		if err != nil {
 | 
	
	
		
			
				|  | @@ -70,45 +70,24 @@ func (w *tcpWorker) Proxy() proxy.Inbound {
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (w *tcpWorker) Start() error {
 | 
	
		
			
				|  |  | -	ctx, cancel := context.WithCancel(context.Background())
 | 
	
		
			
				|  |  | -	w.ctx = ctx
 | 
	
		
			
				|  |  | -	w.cancel = cancel
 | 
	
		
			
				|  |  | -	ctx = internet.ContextWithStreamSettings(ctx, w.stream)
 | 
	
		
			
				|  |  | -	conns := make(chan internet.Connection, 16)
 | 
	
		
			
				|  |  | -	hub, err := internet.ListenTCP(ctx, w.address, w.port, conns)
 | 
	
		
			
				|  |  | +	ctx := internet.ContextWithStreamSettings(context.Background(), w.stream)
 | 
	
		
			
				|  |  | +	hub, err := internet.ListenTCP(ctx, w.address, w.port, func(conn internet.Connection) {
 | 
	
		
			
				|  |  | +		go w.callback(conn)
 | 
	
		
			
				|  |  | +	})
 | 
	
		
			
				|  |  |  	if err != nil {
 | 
	
		
			
				|  |  |  		return newError("failed to listen TCP on ", w.port).AtWarning().Base(err)
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  | -	go w.handleConnections(conns)
 | 
	
		
			
				|  |  |  	w.hub = hub
 | 
	
		
			
				|  |  |  	return nil
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -func (w *tcpWorker) handleConnections(conns <-chan internet.Connection) {
 | 
	
		
			
				|  |  | -	for {
 | 
	
		
			
				|  |  | -		select {
 | 
	
		
			
				|  |  | -		case <-w.ctx.Done():
 | 
	
		
			
				|  |  | -			w.hub.Close()
 | 
	
		
			
				|  |  | -		L:
 | 
	
		
			
				|  |  | -			for {
 | 
	
		
			
				|  |  | -				select {
 | 
	
		
			
				|  |  | -				case conn := <-conns:
 | 
	
		
			
				|  |  | -					conn.Close()
 | 
	
		
			
				|  |  | -				default:
 | 
	
		
			
				|  |  | -					break L
 | 
	
		
			
				|  |  | -				}
 | 
	
		
			
				|  |  | -			}
 | 
	
		
			
				|  |  | -			return
 | 
	
		
			
				|  |  | -		case conn := <-conns:
 | 
	
		
			
				|  |  | -			go w.callback(conn)
 | 
	
		
			
				|  |  | -		}
 | 
	
		
			
				|  |  | -	}
 | 
	
		
			
				|  |  | -}
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -func (w *tcpWorker) Close() {
 | 
	
		
			
				|  |  | +func (w *tcpWorker) Close() error {
 | 
	
		
			
				|  |  |  	if w.hub != nil {
 | 
	
		
			
				|  |  | -		w.cancel()
 | 
	
		
			
				|  |  | +		common.Close(w.hub)
 | 
	
		
			
				|  |  | +		common.Close(w.proxy)
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +	return nil
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (w *tcpWorker) Port() net.Port {
 | 
	
	
		
			
				|  | @@ -121,8 +100,7 @@ type udpConn struct {
 | 
	
		
			
				|  |  |  	output           func([]byte) (int, error)
 | 
	
		
			
				|  |  |  	remote           net.Addr
 | 
	
		
			
				|  |  |  	local            net.Addr
 | 
	
		
			
				|  |  | -	ctx              context.Context
 | 
	
		
			
				|  |  | -	cancel           context.CancelFunc
 | 
	
		
			
				|  |  | +	done             *signal.Done
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (c *udpConn) updateActivity() {
 | 
	
	
		
			
				|  | @@ -135,7 +113,7 @@ func (c *udpConn) Read(buf []byte) (int, error) {
 | 
	
		
			
				|  |  |  		defer in.Release()
 | 
	
		
			
				|  |  |  		c.updateActivity()
 | 
	
		
			
				|  |  |  		return copy(buf, in.Bytes()), nil
 | 
	
		
			
				|  |  | -	case <-c.ctx.Done():
 | 
	
		
			
				|  |  | +	case <-c.done.C():
 | 
	
		
			
				|  |  |  		return 0, io.EOF
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  |  }
 | 
	
	
		
			
				|  | @@ -150,6 +128,7 @@ func (c *udpConn) Write(buf []byte) (int, error) {
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (c *udpConn) Close() error {
 | 
	
		
			
				|  |  | +	common.Close(c.done)
 | 
	
		
			
				|  |  |  	return nil
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
	
		
			
				|  | @@ -189,8 +168,7 @@ type udpWorker struct {
 | 
	
		
			
				|  |  |  	tag          string
 | 
	
		
			
				|  |  |  	dispatcher   core.Dispatcher
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -	ctx        context.Context
 | 
	
		
			
				|  |  | -	cancel     context.CancelFunc
 | 
	
		
			
				|  |  | +	done       *signal.Done
 | 
	
		
			
				|  |  |  	activeConn map[connId]*udpConn
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
	
		
			
				|  | @@ -215,6 +193,7 @@ func (w *udpWorker) getConnection(id connId) (*udpConn, bool) {
 | 
	
		
			
				|  |  |  			IP:   w.address.IP(),
 | 
	
		
			
				|  |  |  			Port: int(w.port),
 | 
	
		
			
				|  |  |  		},
 | 
	
		
			
				|  |  | +		done: signal.NewDone(),
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  |  	w.activeConn[id] = conn
 | 
	
		
			
				|  |  |  
 | 
	
	
		
			
				|  | @@ -230,16 +209,15 @@ func (w *udpWorker) callback(b *buf.Buffer, source net.Destination, originalDest
 | 
	
		
			
				|  |  |  	conn, existing := w.getConnection(id)
 | 
	
		
			
				|  |  |  	select {
 | 
	
		
			
				|  |  |  	case conn.input <- b:
 | 
	
		
			
				|  |  | +	case <-conn.done.C():
 | 
	
		
			
				|  |  | +		b.Release()
 | 
	
		
			
				|  |  |  	default:
 | 
	
		
			
				|  |  |  		b.Release()
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  	if !existing {
 | 
	
		
			
				|  |  |  		go func() {
 | 
	
		
			
				|  |  | -			ctx := w.ctx
 | 
	
		
			
				|  |  | -			ctx, cancel := context.WithCancel(ctx)
 | 
	
		
			
				|  |  | -			conn.ctx = ctx
 | 
	
		
			
				|  |  | -			conn.cancel = cancel
 | 
	
		
			
				|  |  | +			ctx := context.Background()
 | 
	
		
			
				|  |  |  			if originalDest.IsValid() {
 | 
	
		
			
				|  |  |  				ctx = proxy.ContextWithOriginalTarget(ctx, originalDest)
 | 
	
		
			
				|  |  |  			}
 | 
	
	
		
			
				|  | @@ -251,8 +229,8 @@ func (w *udpWorker) callback(b *buf.Buffer, source net.Destination, originalDest
 | 
	
		
			
				|  |  |  			if err := w.proxy.Process(ctx, net.Network_UDP, conn, w.dispatcher); err != nil {
 | 
	
		
			
				|  |  |  				newError("connection ends").Base(err).WriteToLog()
 | 
	
		
			
				|  |  |  			}
 | 
	
		
			
				|  |  | +			conn.Close()
 | 
	
		
			
				|  |  |  			w.removeConn(id)
 | 
	
		
			
				|  |  | -			cancel()
 | 
	
		
			
				|  |  |  		}()
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  |  }
 | 
	
	
		
			
				|  | @@ -265,9 +243,7 @@ func (w *udpWorker) removeConn(id connId) {
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (w *udpWorker) Start() error {
 | 
	
		
			
				|  |  |  	w.activeConn = make(map[connId]*udpConn, 16)
 | 
	
		
			
				|  |  | -	ctx, cancel := context.WithCancel(context.Background())
 | 
	
		
			
				|  |  | -	w.ctx = ctx
 | 
	
		
			
				|  |  | -	w.cancel = cancel
 | 
	
		
			
				|  |  | +	w.done = signal.NewDone()
 | 
	
		
			
				|  |  |  	h, err := udp.ListenUDP(w.address, w.port, udp.ListenOption{
 | 
	
		
			
				|  |  |  		Callback:            w.callback,
 | 
	
		
			
				|  |  |  		ReceiveOriginalDest: w.recvOrigDest,
 | 
	
	
		
			
				|  | @@ -280,11 +256,13 @@ func (w *udpWorker) Start() error {
 | 
	
		
			
				|  |  |  	return nil
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -func (w *udpWorker) Close() {
 | 
	
		
			
				|  |  | +func (w *udpWorker) Close() error {
 | 
	
		
			
				|  |  |  	if w.hub != nil {
 | 
	
		
			
				|  |  |  		w.hub.Close()
 | 
	
		
			
				|  |  | -		w.cancel()
 | 
	
		
			
				|  |  | +		w.done.Close()
 | 
	
		
			
				|  |  | +		common.Close(w.proxy)
 | 
	
		
			
				|  |  |  	}
 | 
	
		
			
				|  |  | +	return nil
 | 
	
		
			
				|  |  |  }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  func (w *udpWorker) monitor() {
 | 
	
	
		
			
				|  | @@ -293,7 +271,7 @@ func (w *udpWorker) monitor() {
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  	for {
 | 
	
		
			
				|  |  |  		select {
 | 
	
		
			
				|  |  | -		case <-w.ctx.Done():
 | 
	
		
			
				|  |  | +		case <-w.done.C():
 | 
	
		
			
				|  |  |  			return
 | 
	
		
			
				|  |  |  		case <-timer.C:
 | 
	
		
			
				|  |  |  			nowSec := time.Now().Unix()
 | 
	
	
		
			
				|  | @@ -301,7 +279,7 @@ func (w *udpWorker) monitor() {
 | 
	
		
			
				|  |  |  			for addr, conn := range w.activeConn {
 | 
	
		
			
				|  |  |  				if nowSec-atomic.LoadInt64(&conn.lastActivityTime) > 8 {
 | 
	
		
			
				|  |  |  					delete(w.activeConn, addr)
 | 
	
		
			
				|  |  | -					conn.cancel()
 | 
	
		
			
				|  |  | +					conn.Close()
 | 
	
		
			
				|  |  |  				}
 | 
	
		
			
				|  |  |  			}
 | 
	
		
			
				|  |  |  			w.Unlock()
 |