github.com/imannamdari/v2ray-core/v5@v5.0.5/app/commander/outbound.go (about)

     1  package commander
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/imannamdari/v2ray-core/v5/common"
     8  	"github.com/imannamdari/v2ray-core/v5/common/net"
     9  	"github.com/imannamdari/v2ray-core/v5/common/signal/done"
    10  	"github.com/imannamdari/v2ray-core/v5/transport"
    11  )
    12  
    13  // OutboundListener is a net.Listener for listening gRPC connections.
    14  type OutboundListener struct {
    15  	buffer chan net.Conn
    16  	done   *done.Instance
    17  }
    18  
    19  func (l *OutboundListener) add(conn net.Conn) {
    20  	select {
    21  	case l.buffer <- conn:
    22  	case <-l.done.Wait():
    23  		conn.Close()
    24  	default:
    25  		conn.Close()
    26  	}
    27  }
    28  
    29  // Accept implements net.Listener.
    30  func (l *OutboundListener) Accept() (net.Conn, error) {
    31  	select {
    32  	case <-l.done.Wait():
    33  		return nil, newError("listen closed")
    34  	case c := <-l.buffer:
    35  		return c, nil
    36  	}
    37  }
    38  
    39  // Close implement net.Listener.
    40  func (l *OutboundListener) Close() error {
    41  	common.Must(l.done.Close())
    42  L:
    43  	for {
    44  		select {
    45  		case c := <-l.buffer:
    46  			c.Close()
    47  		default:
    48  			break L
    49  		}
    50  	}
    51  	return nil
    52  }
    53  
    54  // Addr implements net.Listener.
    55  func (l *OutboundListener) Addr() net.Addr {
    56  	return &net.TCPAddr{
    57  		IP:   net.IP{0, 0, 0, 0},
    58  		Port: 0,
    59  	}
    60  }
    61  
    62  // Outbound is a outbound.Handler that handles gRPC connections.
    63  type Outbound struct {
    64  	tag      string
    65  	listener *OutboundListener
    66  	access   sync.RWMutex
    67  	closed   bool
    68  }
    69  
    70  // Dispatch implements outbound.Handler.
    71  func (co *Outbound) Dispatch(ctx context.Context, link *transport.Link) {
    72  	co.access.RLock()
    73  
    74  	if co.closed {
    75  		common.Interrupt(link.Reader)
    76  		common.Interrupt(link.Writer)
    77  		co.access.RUnlock()
    78  		return
    79  	}
    80  
    81  	closeSignal := done.New()
    82  	c := net.NewConnection(net.ConnectionInputMulti(link.Writer), net.ConnectionOutputMulti(link.Reader), net.ConnectionOnClose(closeSignal))
    83  	co.listener.add(c)
    84  	co.access.RUnlock()
    85  	<-closeSignal.Wait()
    86  }
    87  
    88  // Tag implements outbound.Handler.
    89  func (co *Outbound) Tag() string {
    90  	return co.tag
    91  }
    92  
    93  // Start implements common.Runnable.
    94  func (co *Outbound) Start() error {
    95  	co.access.Lock()
    96  	co.closed = false
    97  	co.access.Unlock()
    98  	return nil
    99  }
   100  
   101  // Close implements common.Closable.
   102  func (co *Outbound) Close() error {
   103  	co.access.Lock()
   104  	defer co.access.Unlock()
   105  
   106  	co.closed = true
   107  	return co.listener.Close()
   108  }