#141: Add load balancing support for proxy chain

chain.go

@@ -95,12 +95,17 @@ func (c *Chain) Dial(addr string) (net.Conn, error) {
 		return net.Dial("tcp", addr)
 	}
 
-	conn, nodes, err := c.getConn()
+	route, err := c.selectRoute()
 	if err != nil {
 		return nil, err
 	}
 
-	cc, err := nodes[len(nodes)-1].Client.Connect(conn, addr)
+	conn, err := c.getConn(route)
+	if err != nil {
+		return nil, err
+	}
+
+	cc, err := route.LastNode().Client.Connect(conn, addr)
 	if err != nil {
 		conn.Close()
 		return nil, err
@@ -111,26 +116,44 @@ func (c *Chain) Dial(addr string) (net.Conn, error) {
 // Conn obtains a handshaked connection to the last node of the chain.
 // If the chain is empty, it returns an ErrEmptyChain error.
 func (c *Chain) Conn() (conn net.Conn, err error) {
-	conn, _, err = c.getConn()
+	route, err := c.selectRoute()
+	if err != nil {
+		return nil, err
+	}
+	conn, err = c.getConn(route)
 	return
 }
 
-func (c *Chain) getConn() (conn net.Conn, nodes []Node, err error) {
-	if c.IsEmpty() {
-		err = ErrEmptyChain
-		return
-	}
-	groups := c.nodeGroups
-	selector := groups[0].Selector
+func (c *Chain) selectRoute() (route *Chain, err error) {
+	route = NewChain()
+	for _, group := range c.nodeGroups {
+		selector := group.Selector
 		if selector == nil {
 			selector = &defaultSelector{}
 		}
 		// select node from node group
-	node, err := selector.Select(groups[0].Nodes(), groups[0].Options...)
+		node, err := selector.Select(group.Nodes(), group.Options...)
 		if err != nil {
+			return nil, err
+		}
+		if node.Client.Transporter.Multiplex() {
+			node.DialOptions = append(node.DialOptions,
+				ChainDialOption(route),
+			)
+			route = NewChain() // cutoff the chain for multiplex
+		}
+		route.AddNode(node)
+	}
+	return
+}
+
+func (c *Chain) getConn(route *Chain) (conn net.Conn, err error) {
+	if route.IsEmpty() {
+		err = ErrEmptyChain
 		return
 	}
-	nodes = append(nodes, node)
+	nodes := route.Nodes()
+	node := nodes[0]
 
 	addr, err := selectIP(&node)
 	if err != nil {
@@ -147,21 +170,7 @@ func (c *Chain) getConn() (conn net.Conn, nodes []Node, err error) {
 	}
 
 	preNode := node
-	for i := range groups {
-		if i == len(groups)-1 {
-			break
-		}
-		selector = groups[i+1].Selector
-		if selector == nil {
-			selector = &defaultSelector{}
-		}
-		node, err = selector.Select(groups[i+1].Nodes(), groups[i+1].Options...)
-		if err != nil {
-			cn.Close()
-			return
-		}
-		nodes = append(nodes, node)
-
+	for _, node := range nodes[1:] {
 		addr, err = selectIP(&node)
 		if err != nil {
 			return
@@ -206,6 +215,7 @@ func selectIP(node *Node) (string, error) {
 			ip = ip + ":" + sport
 		}
 		addr = ip
+		// override the original address
 		node.HandshakeOptions = append(node.HandshakeOptions, AddrHandshakeOption(addr))
 	}
 	log.Log("select IP:", node.Addr, node.IPs, addr)

client.go

@@ -94,7 +94,6 @@ func (tr *tcpTransporter) Multiplex() bool {
 type DialOptions struct {
 	Timeout time.Duration
 	Chain   *Chain
-	// IPs     []string
 }
 
 // DialOption allows a common way to set dial options.

cmd/gost/main.go

@@ -62,6 +62,16 @@ func init() {
 }
 
 func main() {
+	// generate random self-signed certificate.
+	cert, err := gost.GenCertificate()
+	if err != nil {
+		log.Log(err)
+		os.Exit(1)
+	}
+	gost.DefaultTLSConfig = &tls.Config{
+		Certificates: []tls.Certificate{cert},
+	}
+
 	chain, err := initChain()
 	if err != nil {
 		log.Log(err)
@@ -71,23 +81,56 @@ func main() {
 		log.Log(err)
 		os.Exit(1)
 	}
+
 	select {}
 }
 
 func initChain() (*gost.Chain, error) {
 	chain := gost.NewChain()
 	for _, ns := range options.ChainNodes {
-		node, err := gost.ParseNode(ns)
+		// parse the base node
+		node, err := parseChainNode(ns)
 		if err != nil {
 			return nil, err
 		}
 
+		ngroup := gost.NewNodeGroup(node)
+
+		// parse node peers if exists
+		peerCfg, err := loadPeerConfig(node.Values.Get("peer"))
+		if err != nil {
+			log.Log(err)
+		}
+		ngroup.Options = append(ngroup.Options,
+			// gost.WithFilter(),
+			gost.WithStrategy(parseStrategy(peerCfg.Strategy)),
+		)
+		for _, s := range peerCfg.Nodes {
+			node, err = parseChainNode(s)
+			if err != nil {
+				return nil, err
+			}
+			ngroup.AddNode(node)
+		}
+
+		chain.AddNodeGroup(ngroup)
+	}
+
+	return chain, nil
+}
+
+func parseChainNode(ns string) (node gost.Node, err error) {
+	node, err = gost.ParseNode(ns)
+	if err != nil {
+		return
+	}
+
 	node.IPs = parseIP(node.Values.Get("ip"))
 	node.IPSelector = &gost.RoundRobinIPSelector{}
 
 	users, err := parseUsers(node.Values.Get("secrets"))
 	if err != nil {
-			return nil, err
+		return
 	}
 	if node.User == nil && len(users) > 0 {
 		node.User = users[0]
@@ -99,7 +142,7 @@ func initChain() (*gost.Chain, error) {
 
 	rootCAs, err := loadCA(node.Values.Get("ca"))
 	if err != nil {
-			return nil, err
+		return
 	}
 	tlsCfg := &tls.Config{
 		ServerName:         serverName,
@@ -127,12 +170,14 @@ func initChain() (*gost.Chain, error) {
 	case "mwss":
 		tr = gost.MWSSTransporter(wsOpts)
 	case "kcp":
+		/*
 			if !chain.IsEmpty() {
 				return nil, errors.New("KCP must be the first node in the proxy chain")
 			}
+		*/
 		config, err := parseKCPConfig(node.Values.Get("c"))
 		if err != nil {
-				return nil, err
+			return node, err
 		}
 		tr = gost.KCPTransporter(config)
 	case "ssh":
@@ -142,9 +187,11 @@ func initChain() (*gost.Chain, error) {
 			tr = gost.SSHTunnelTransporter()
 		}
 	case "quic":
+		/*
 			if !chain.IsEmpty() {
 				return nil, errors.New("QUIC must be the first node in the proxy chain")
 			}
+		*/
 		config := &gost.QUICConfig{
 			TLSConfig: tlsCfg,
 			KeepAlive: toBool(node.Values.Get("keepalive")),
@@ -159,7 +206,7 @@ func initChain() (*gost.Chain, error) {
 
 	case "obfs4":
 		if err := gost.Obfs4Init(node, false); err != nil {
-				return nil, err
+			return node, err
 		}
 		tr = gost.Obfs4Transporter()
 	case "ohttp":
@@ -168,13 +215,6 @@ func initChain() (*gost.Chain, error) {
 		tr = gost.TCPTransporter()
 	}
 
-		if tr.Multiplex() {
-			node.DialOptions = append(node.DialOptions,
-				gost.ChainDialOption(chain),
-			)
-			chain = gost.NewChain() // cutoff the chain for multiplex
-		}
-
 	var connector gost.Connector
 	switch node.Protocol {
 	case "http2":
@@ -221,10 +261,8 @@ func initChain() (*gost.Chain, error) {
 		Connector:   connector,
 		Transporter: tr,
 	}
-		chain.AddNode(node)
-	}
 
-	return chain, nil
+	return
 }
 
 func serve(chain *gost.Chain) error {
@@ -533,3 +571,32 @@ func parseIP(s string) (ips []string) {
 	}
 	return
 }
+
+type peerConfig struct {
+	Strategy string   `json:"strategy"`
+	Filters  []string `json:"filters"`
+	Nodes    []string `json:"nodes"`
+}
+
+func loadPeerConfig(peer string) (config peerConfig, err error) {
+	if peer == "" {
+		return
+	}
+	content, err := ioutil.ReadFile(peer)
+	if err != nil {
+		return
+	}
+	err = json.Unmarshal(content, &config)
+	return
+}
+
+func parseStrategy(s string) gost.Strategy {
+	switch s {
+	case "round":
+		return &gost.RoundStrategy{}
+	case "random":
+		fallthrough
+	default:
+		return &gost.RandomStrategy{}
+	}
+}

gost.go

@@ -38,7 +38,7 @@ var (
 	// PingTimeout is the timeout for pinging.
 	PingTimeout = 30 * time.Second
 	// PingRetries is the reties of ping.
-	PingRetries = 3
+	PingRetries = 1
 	// default udp node TTL in second for udp port forwarding.
 	defaultTTL = 60 * time.Second
 )
@@ -51,27 +51,19 @@ var (
 	DefaultUserAgent = "Chrome/60.0.3112.90"
 )
 
-func init() {
-	rawCert, rawKey, err := generateKeyPair()
-	if err != nil {
-		panic(err)
-	}
-	cert, err := tls.X509KeyPair(rawCert, rawKey)
-	if err != nil {
-		panic(err)
-	}
-	DefaultTLSConfig = &tls.Config{
-		Certificates: []tls.Certificate{cert},
-	}
-
-	// log.DefaultLogger = &LogLogger{}
-}
-
 // SetLogger sets a new logger for internal log system
 func SetLogger(logger log.Logger) {
 	log.DefaultLogger = logger
 }
 
+func GenCertificate() (cert tls.Certificate, err error) {
+	rawCert, rawKey, err := generateKeyPair()
+	if err != nil {
+		return
+	}
+	return tls.X509KeyPair(rawCert, rawKey)
+}
+
 func generateKeyPair() (rawCert, rawKey []byte, err error) {
 	// Create private key and self-signed certificate
 	// Adapted from https://golang.org/src/crypto/tls/generate_cert.go

quic.go

@@ -194,6 +194,7 @@ func (l *quicListener) sessionLoop(session quic.Session) {
 		stream, err := session.AcceptStream()
 		if err != nil {
 			log.Log("[quic] accept stream:", err)
+			session.Close(err)
 			return
 		}
 

selector.go

@@ -11,14 +11,10 @@ var (
 	ErrNoneAvailable = errors.New("none available")
 )
 
-// SelectOption used when making a select call
-type SelectOption func(*SelectOptions)
-
 // NodeSelector as a mechanism to pick nodes and mark their status.
 type NodeSelector interface {
 	Select(nodes []Node, opts ...SelectOption) (Node, error)
 	// Mark(node Node)
-	String() string
 }
 
 type defaultSelector struct {
@@ -26,35 +22,70 @@ type defaultSelector struct {
 
 func (s *defaultSelector) Select(nodes []Node, opts ...SelectOption) (Node, error) {
 	sopts := SelectOptions{
-		Strategy: defaultStrategy,
+		Strategy: &RoundStrategy{},
 	}
 	for _, opt := range opts {
 		opt(&sopts)
 	}
 
 	for _, filter := range sopts.Filters {
-		nodes = filter(nodes)
+		nodes = filter.Filter(nodes)
 	}
 	if len(nodes) == 0 {
 		return Node{}, ErrNoneAvailable
 	}
-	return sopts.Strategy(nodes), nil
-}
-
-func (s *defaultSelector) String() string {
-	return "default"
+	return sopts.Strategy.Apply(nodes), nil
 }
 
 // Filter is used to filter a node during the selection process
-type Filter func([]Node) []Node
+type Filter interface {
+	Filter([]Node) []Node
+}
 
 // Strategy is a selection strategy e.g random, round robin
-type Strategy func([]Node) Node
-
-func defaultStrategy(nodes []Node) Node {
-	return nodes[0]
+type Strategy interface {
+	Apply([]Node) Node
+	String() string
 }
 
+// RoundStrategy is a strategy for node selector
+type RoundStrategy struct {
+	count uint64
+}
+
+// Apply applies the round robin strategy for the nodes
+func (s *RoundStrategy) Apply(nodes []Node) Node {
+	if len(nodes) == 0 {
+		return Node{}
+	}
+	old := s.count
+	atomic.AddUint64(&s.count, 1)
+	return nodes[int(old%uint64(len(nodes)))]
+}
+
+func (s *RoundStrategy) String() string {
+	return "round"
+}
+
+// RandomStrategy is a strategy for node selector
+type RandomStrategy struct{}
+
+// Apply applies the random strategy for the nodes
+func (s *RandomStrategy) Apply(nodes []Node) Node {
+	if len(nodes) == 0 {
+		return Node{}
+	}
+
+	return nodes[time.Now().Nanosecond()%len(nodes)]
+}
+
+func (s *RandomStrategy) String() string {
+	return "random"
+}
+
+// SelectOption used when making a select call
+type SelectOption func(*SelectOptions)
+
 // SelectOptions is the options for node selection
 type SelectOptions struct {
 	Filters  []Filter
@@ -108,9 +139,9 @@ func (s *RoundRobinIPSelector) Select(ips []string) (string, error) {
 	if len(ips) == 0 {
 		return "", nil
 	}
-
-	count := atomic.AddUint64(&s.count, 1)
-	return ips[int(count%uint64(len(ips)))], nil
+	old := s.count
+	atomic.AddUint64(&s.count, 1)
+	return ips[int(old%uint64(len(ips)))], nil
 }
 
 func (s *RoundRobinIPSelector) String() string {