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

     1  package reverse
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"google.golang.org/protobuf/proto"
     8  
     9  	"github.com/imannamdari/v2ray-core/v5/common/mux"
    10  	"github.com/imannamdari/v2ray-core/v5/common/net"
    11  	"github.com/imannamdari/v2ray-core/v5/common/session"
    12  	"github.com/imannamdari/v2ray-core/v5/common/task"
    13  	"github.com/imannamdari/v2ray-core/v5/features/routing"
    14  	"github.com/imannamdari/v2ray-core/v5/transport"
    15  	"github.com/imannamdari/v2ray-core/v5/transport/pipe"
    16  )
    17  
    18  // Bridge is a component in reverse proxy, that relays connections from Portal to local address.
    19  type Bridge struct {
    20  	ctx         context.Context
    21  	dispatcher  routing.Dispatcher
    22  	tag         string
    23  	domain      string
    24  	workers     []*BridgeWorker
    25  	monitorTask *task.Periodic
    26  }
    27  
    28  // NewBridge creates a new Bridge instance.
    29  func NewBridge(ctx context.Context, config *BridgeConfig, dispatcher routing.Dispatcher) (*Bridge, error) {
    30  	if config.Tag == "" {
    31  		return nil, newError("bridge tag is empty")
    32  	}
    33  	if config.Domain == "" {
    34  		return nil, newError("bridge domain is empty")
    35  	}
    36  
    37  	b := &Bridge{
    38  		ctx:        ctx,
    39  		dispatcher: dispatcher,
    40  		tag:        config.Tag,
    41  		domain:     config.Domain,
    42  	}
    43  	b.monitorTask = &task.Periodic{
    44  		Execute:  b.monitor,
    45  		Interval: time.Second * 2,
    46  	}
    47  	return b, nil
    48  }
    49  
    50  func (b *Bridge) cleanup() {
    51  	var activeWorkers []*BridgeWorker
    52  
    53  	for _, w := range b.workers {
    54  		if w.IsActive() {
    55  			activeWorkers = append(activeWorkers, w)
    56  		}
    57  	}
    58  
    59  	if len(activeWorkers) != len(b.workers) {
    60  		b.workers = activeWorkers
    61  	}
    62  }
    63  
    64  func (b *Bridge) monitor() error {
    65  	b.cleanup()
    66  
    67  	var numConnections uint32
    68  	var numWorker uint32
    69  
    70  	for _, w := range b.workers {
    71  		if w.IsActive() {
    72  			numConnections += w.Connections()
    73  			numWorker++
    74  		}
    75  	}
    76  
    77  	if numWorker == 0 || numConnections/numWorker > 16 {
    78  		worker, err := NewBridgeWorker(b.ctx, b.domain, b.tag, b.dispatcher)
    79  		if err != nil {
    80  			newError("failed to create bridge worker").Base(err).AtWarning().WriteToLog()
    81  			return nil
    82  		}
    83  		b.workers = append(b.workers, worker)
    84  	}
    85  
    86  	return nil
    87  }
    88  
    89  func (b *Bridge) Start() error {
    90  	return b.monitorTask.Start()
    91  }
    92  
    93  func (b *Bridge) Close() error {
    94  	return b.monitorTask.Close()
    95  }
    96  
    97  type BridgeWorker struct {
    98  	tag        string
    99  	worker     *mux.ServerWorker
   100  	dispatcher routing.Dispatcher
   101  	state      Control_State
   102  }
   103  
   104  func NewBridgeWorker(ctx context.Context, domain string, tag string, d routing.Dispatcher) (*BridgeWorker, error) {
   105  	bridgeCtx := session.ContextWithInbound(ctx, &session.Inbound{
   106  		Tag: tag,
   107  	})
   108  	link, err := d.Dispatch(bridgeCtx, net.Destination{
   109  		Network: net.Network_TCP,
   110  		Address: net.DomainAddress(domain),
   111  		Port:    0,
   112  	})
   113  	if err != nil {
   114  		return nil, err
   115  	}
   116  
   117  	w := &BridgeWorker{
   118  		dispatcher: d,
   119  		tag:        tag,
   120  	}
   121  
   122  	worker, err := mux.NewServerWorker(ctx, w, link)
   123  	if err != nil {
   124  		return nil, err
   125  	}
   126  	w.worker = worker
   127  
   128  	return w, nil
   129  }
   130  
   131  func (w *BridgeWorker) Type() interface{} {
   132  	return routing.DispatcherType()
   133  }
   134  
   135  func (w *BridgeWorker) Start() error {
   136  	return nil
   137  }
   138  
   139  func (w *BridgeWorker) Close() error {
   140  	return nil
   141  }
   142  
   143  func (w *BridgeWorker) IsActive() bool {
   144  	return w.state == Control_ACTIVE && !w.worker.Closed()
   145  }
   146  
   147  func (w *BridgeWorker) Connections() uint32 {
   148  	return w.worker.ActiveConnections()
   149  }
   150  
   151  func (w *BridgeWorker) handleInternalConn(link transport.Link) {
   152  	go func() {
   153  		reader := link.Reader
   154  		for {
   155  			mb, err := reader.ReadMultiBuffer()
   156  			if err != nil {
   157  				break
   158  			}
   159  			for _, b := range mb {
   160  				var ctl Control
   161  				if err := proto.Unmarshal(b.Bytes(), &ctl); err != nil {
   162  					newError("failed to parse proto message").Base(err).WriteToLog()
   163  					break
   164  				}
   165  				if ctl.State != w.state {
   166  					w.state = ctl.State
   167  				}
   168  			}
   169  		}
   170  	}()
   171  }
   172  
   173  func (w *BridgeWorker) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
   174  	if !isInternalDomain(dest) {
   175  		ctx = session.ContextWithInbound(ctx, &session.Inbound{
   176  			Tag: w.tag,
   177  		})
   178  		return w.dispatcher.Dispatch(ctx, dest)
   179  	}
   180  
   181  	opt := []pipe.Option{pipe.WithSizeLimit(16 * 1024)}
   182  	uplinkReader, uplinkWriter := pipe.New(opt...)
   183  	downlinkReader, downlinkWriter := pipe.New(opt...)
   184  
   185  	w.handleInternalConn(transport.Link{
   186  		Reader: downlinkReader,
   187  		Writer: uplinkWriter,
   188  	})
   189  
   190  	return &transport.Link{
   191  		Reader: uplinkReader,
   192  		Writer: downlinkWriter,
   193  	}, nil
   194  }