diff --git a/config/parsing/selector/parse.go b/config/parsing/selector/parse.go index b017d6ff..6f3e63bf 100644 --- a/config/parsing/selector/parse.go +++ b/config/parsing/selector/parse.go @@ -47,6 +47,8 @@ func ParseNodeSelector(cfg *config.SelectorConfig) selector.Selector[*chain.Node strategy = xs.FIFOStrategy[*chain.Node]() case "hash": strategy = xs.HashStrategy[*chain.Node]() + case "parallel": + strategy = xs.ParallelStrategy[*chain.Node]() default: strategy = xs.RoundRobinStrategy[*chain.Node]() } diff --git a/selector/parallel.go b/selector/parallel.go new file mode 100644 index 00000000..2f418926 --- /dev/null +++ b/selector/parallel.go @@ -0,0 +1,167 @@ +package selector + +import ( + "context" + "errors" + "net" + "sync" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/connector" + "github.com/go-gost/core/logger" + "github.com/go-gost/core/selector" + xnet "github.com/go-gost/x/internal/net" +) + +type parallelStrategy[T any] struct{} + +func ParallelStrategy[T any]() selector.Strategy[T] { + return ¶llelStrategy[T]{} +} + +func (s *parallelStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) { + if len(vs) == 0 { + return + } + if len(vs) == 1 { + return vs[0] + } + + nodes := make([]*chain.Node, 0, len(vs)) + for _, node := range vs { + nodes = append(nodes, any(node).(*chain.Node)) + } + + vn := chain.NewNode("parallel", "", chain.TransportNodeOption(¶llelTransporter{nodes: nodes})) + return any(vn).(T) +} + +type parallelConn struct { + net.Conn + node *chain.Node +} + +type parallelTransporter struct { + nodes []*chain.Node +} + +func (tr *parallelTransporter) Dial(ctx context.Context, addr string) (net.Conn, error) { + if len(tr.nodes) == 0 { + return nil, errors.New("no nodes available") + } + + ctx, cancel := context.WithCancel(ctx) + + type result struct { + conn net.Conn + node *chain.Node + err error + } + + resCh := make(chan result, len(tr.nodes)) + + var wg sync.WaitGroup + for _, node := range tr.nodes { + wg.Add(1) + go func(n *chain.Node) { + defer wg.Done() + + raddr, err := xnet.Resolve(ctx, "ip", n.Addr, n.Options().Resolver, n.Options().HostMapper, logger.Default()) + if err != nil { + resCh <- result{err: err, node: n} + return + } + + cc, err := n.Options().Transport.Dial(ctx, raddr) + resCh <- result{conn: cc, node: n, err: err} + }(node) + } + + var best result + var errs []error + + for i := 0; i < len(tr.nodes); i++ { + res := <-resCh + if res.err != nil { + errs = append(errs, res.err) + continue + } + + best = res + break + } + + cancel() + + if best.conn == nil { + go func() { + wg.Wait() + close(resCh) + }() + return nil, errors.Join(errs...) + } + + go func() { + // Drain the rest of the channel to clean up late successful connections + for i := len(errs) + 1; i < len(tr.nodes); i++ { + res := <-resCh + if res.err == nil && res.conn != nil { + res.conn.Close() + } + } + wg.Wait() + close(resCh) + }() + + return ¶llelConn{ + Conn: best.conn, + node: best.node, + }, nil +} + +func (tr *parallelTransporter) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) { + pc, ok := conn.(*parallelConn) + if !ok { + return conn, nil + } + + cc, err := pc.node.Options().Transport.Handshake(ctx, pc.Conn) + if err != nil { + return nil, err + } + pc.Conn = cc + + return pc, nil +} + +func (tr *parallelTransporter) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) { + pc, ok := conn.(*parallelConn) + if !ok { + return nil, errors.New("invalid connection type") + } + + return pc.node.Options().Transport.Connect(ctx, pc.Conn, network, address) +} + +func (tr *parallelTransporter) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) { + pc, ok := conn.(*parallelConn) + if !ok { + return nil, errors.New("invalid connection type") + } + + return pc.node.Options().Transport.Bind(ctx, pc.Conn, network, address, opts...) +} + +func (tr *parallelTransporter) Multiplex() bool { + return false +} + +func (tr *parallelTransporter) Options() *chain.TransportOptions { + return &chain.TransportOptions{} +} + +func (tr *parallelTransporter) Copy() chain.Transporter { + return ¶llelTransporter{ + nodes: append([]*chain.Node(nil), tr.nodes...), + } +}