github.com/ipfans/trojan-go@v0.11.0/proxy/stack.go (about)

     1  package proxy
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/ipfans/trojan-go/log"
     7  	"github.com/ipfans/trojan-go/tunnel"
     8  )
     9  
    10  type Node struct {
    11  	Name       string
    12  	Next       map[string]*Node
    13  	IsEndpoint bool
    14  	context.Context
    15  	tunnel.Server
    16  	tunnel.Client
    17  }
    18  
    19  func (n *Node) BuildNext(name string) *Node {
    20  	if next, found := n.Next[name]; found {
    21  		return next
    22  	}
    23  	t, err := tunnel.GetTunnel(name)
    24  	if err != nil {
    25  		log.Fatal(err)
    26  	}
    27  	s, err := t.NewServer(n.Context, n.Server)
    28  	if err != nil {
    29  		log.Fatal(err)
    30  	}
    31  	newNode := &Node{
    32  		Name:    name,
    33  		Next:    make(map[string]*Node),
    34  		Context: n.Context,
    35  		Server:  s,
    36  	}
    37  	n.Next[name] = newNode
    38  	return newNode
    39  }
    40  
    41  func (n *Node) LinkNextNode(next *Node) *Node {
    42  	if next, found := n.Next[next.Name]; found {
    43  		return next
    44  	}
    45  	n.Next[next.Name] = next
    46  	t, err := tunnel.GetTunnel(next.Name)
    47  	if err != nil {
    48  		log.Fatal(err)
    49  	}
    50  	s, err := t.NewServer(next.Context, n.Server) // context of the child nodes have been initialized
    51  	if err != nil {
    52  		log.Fatal(err)
    53  	}
    54  	next.Server = s
    55  	return next
    56  }
    57  
    58  func FindAllEndpoints(root *Node) []tunnel.Server {
    59  	list := make([]tunnel.Server, 0)
    60  	if root.IsEndpoint || len(root.Next) == 0 {
    61  		list = append(list, root.Server)
    62  	}
    63  	for _, next := range root.Next {
    64  		list = append(list, FindAllEndpoints(next)...)
    65  	}
    66  	return list
    67  }
    68  
    69  // CreateClientStack create client tunnel stacks from lists
    70  func CreateClientStack(ctx context.Context, clientStack []string) (tunnel.Client, error) {
    71  	var client tunnel.Client
    72  	for _, name := range clientStack {
    73  		t, err := tunnel.GetTunnel(name)
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  		client, err = t.NewClient(ctx, client)
    78  		if err != nil {
    79  			return nil, err
    80  		}
    81  	}
    82  	return client, nil
    83  }
    84  
    85  // CreateServerStack create server tunnel stack from list
    86  func CreateServerStack(ctx context.Context, serverStack []string) (tunnel.Server, error) {
    87  	var server tunnel.Server
    88  	for _, name := range serverStack {
    89  		t, err := tunnel.GetTunnel(name)
    90  		if err != nil {
    91  			return nil, err
    92  		}
    93  		server, err = t.NewServer(ctx, server)
    94  		if err != nil {
    95  			return nil, err
    96  		}
    97  	}
    98  	return server, nil
    99  }