Merge pull request Ehco1996#34 from Ehco1996/ref

refactor package

cmd/ehco/main.go

@@ -7,7 +7,10 @@ import (
 
 	cli "github.com/urfave/cli/v2"
 
-	relay "github.com/Ehco1996/ehco/internal/relay"
+	"github.com/Ehco1996/ehco/internal/constant"
+	"github.com/Ehco1996/ehco/internal/logger"
+	"github.com/Ehco1996/ehco/internal/relay"
+	"github.com/Ehco1996/ehco/internal/tls"
 )
 
 var LocalAddr string
@@ -77,7 +80,7 @@ func main() {
 	app.Action = start
 	err := app.Run(os.Args)
 	if err != nil {
-		relay.Logger.Fatal(err)
+		logger.Logger.Fatal(err)
 	}
 
 }
@@ -89,7 +92,7 @@ func start(ctx *cli.Context) error {
 	if ConfigPath != "" {
 		config = relay.NewConfigByPath(ConfigPath)
 		if err := config.LoadConfig(); err != nil {
-			relay.Logger.Fatal(err)
+			logger.Logger.Fatal(err)
 		}
 	} else {
 		config = &relay.Config{
@@ -98,21 +101,21 @@ func start(ctx *cli.Context) error {
 				{
 					Listen:        LocalAddr,
 					ListenType:    ListenType,
-					Remote:        RemoteAddr,
-					UDPRemote:     UDPRemoteAddr,
+					TCPRemotes:    []string{RemoteAddr},
+					UDPRemotes:    []string{UDPRemoteAddr},
 					TransportType: TransportType,
 				},
 			},
 		}
 	}
 	initTls := false
 	for _, cfg := range config.Configs {
-		if !initTls && (cfg.ListenType == relay.Listen_WSS ||
-			cfg.ListenType == relay.Listen_MWSS ||
-			cfg.TransportType == relay.Transport_WSS ||
-			cfg.TransportType == relay.Transport_MWSS) {
+		if !initTls && (cfg.ListenType == constant.Listen_WSS ||
+			cfg.ListenType == constant.Listen_MWSS ||
+			cfg.TransportType == constant.Transport_WSS ||
+			cfg.TransportType == constant.Transport_MWSS) {
 			initTls = true
-			relay.InitTlsCfg()
+			tls.InitTlsCfg()
 		}
 		go serveRelay(cfg, ch)
 	}
@@ -121,8 +124,8 @@ func start(ctx *cli.Context) error {
 	if PprofPort != "" {
 		go func() {
 			pps := "0.0.0.0:" + PprofPort
-			relay.Logger.Infof("start pprof server at http://%s/debug/pprof/", pps)
-			relay.Logger.Fatal(http.ListenAndServe(pps, nil))
+			logger.Logger.Infof("start pprof server at http://%s/debug/pprof/", pps)
+			logger.Logger.Fatal(http.ListenAndServe(pps, nil))
 		}()
 	}
 
@@ -132,7 +135,7 @@ func start(ctx *cli.Context) error {
 func serveRelay(cfg relay.RelayConfig, ch chan error) {
 	r, err := relay.NewRelay(&cfg)
 	if err != nil {
-		relay.Logger.Fatal(err)
+		logger.Logger.Fatal(err)
 	}
 	ch <- r.ListenAndServe()
 }

config.json

@@ -1,63 +1,75 @@
 {
-    "configs": [
+    "relay_configs": [
         {
-            "listen": "0.0.0.0:1234",
+            "listen": "127.0.0.1:1234",
             "listen_type": "raw",
-            "remote": "0.0.0.0:5201",
-            "transport_type": "raw"
-        },
-        {
-            "listen": "0.0.0.0:1235",
-            "listen_type": "raw",
-            "remote": "ws://0.0.0.0:2000",
-            "transport_type": "ws"
+            "transport_type": "raw",
+            "tcp_remotes": [
+                "0.0.0.0:5201"
+            ],
+            "udp_remotes": [
+                "0.0.0.0:5201"
+            ]
         },
         {
-            "listen": "0.0.0.0:1236",
+            "listen": "127.0.0.1:1235",
             "listen_type": "raw",
-            "remote": "wss://0.0.0.0:2001",
-            "transport_type": "wss"
+            "transport_type": "ws",
+            "tcp_remotes": [
+                "ws://0.0.0.0:2443"
+            ],
+            "udp_remotes": [
+                "0.0.0.0:5201"
+            ]
         },
         {
-            "listen": "0.0.0.0:1237",
+            "listen": "127.0.0.1:1236",
             "listen_type": "raw",
-            "remote": "wss://0.0.0.0:2002",
-            "transport_type": "mwss"
+            "transport_type": "wss",
+            "tcp_remotes": [
+                "wss://0.0.0.0:3443"
+            ],
+            "udp_remotes": [
+                "0.0.0.0:5201"
+            ]
         },
         {
-            "listen": "0.0.0.0:1238",
+            "listen": "127.0.0.1:1237",
             "listen_type": "raw",
-            "remote": "",
-            "udp_remote": "0.0.0.0:5201",
             "transport_type": "mwss",
-            "lb_remotes": [
-                "wss://0.0.0.0:2002",
-                "wss://0.0.0.0:2003"
+            "tcp_remotes": [
+                "wss://0.0.0.0:4443"
+            ],
+            "udp_remotes": [
+                "0.0.0.0:5201"
             ]
         },
         {
-            "listen": "0.0.0.0:2000",
+            "listen": "127.0.0.1:2443",
             "listen_type": "ws",
-            "remote": "0.0.0.0:5201",
-            "transport_type": "raw"
+            "transport_type": "raw",
+            "tcp_remotes": [
+                "0.0.0.0:5201"
+            ],
+            "udp_remotes": []
         },
         {
-            "listen": "0.0.0.0:2001",
+            "listen": "127.0.0.1:3443",
             "listen_type": "wss",
-            "remote": "0.0.0.0:5201",
-            "transport_type": "raw"
-        },
-        {
-            "listen": "0.0.0.0:2002",
-            "listen_type": "mwss",
-            "remote": "0.0.0.0:5201",
-            "transport_type": "raw"
+            "transport_type": "raw",
+            "tcp_remotes": [
+                "0.0.0.0:5201"
+            ],
+            "udp_remotes": []
         },
         {
-            "listen": "0.0.0.0:2003",
+            "listen": "127.0.0.1:4443",
             "listen_type": "mwss",
-            "remote": "0.0.0.0:5201",
-            "transport_type": "raw"
+            "transport_type": "raw",
+            "tcp_remotes": [
+                "0.0.0.0:5201"
+            ],
+            "udp_remotes": []
         }
     ]
 }

go.mod

@@ -5,8 +5,8 @@ go 1.15
 require (
 	github.com/gobwas/httphead v0.1.0 // indirect
 	github.com/gobwas/pool v0.2.1 // indirect
-	github.com/gobwas/ws v1.0.3
-	github.com/urfave/cli/v2 v2.1.1
+	github.com/gobwas/ws v1.0.4
+	github.com/urfave/cli/v2 v2.3.0
 	github.com/xtaci/smux v2.0.1+incompatible
-	go.uber.org/zap v1.15.0
+	go.uber.org/zap v1.16.0
 )

go.sum

@@ -9,8 +9,8 @@ github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU
 github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
 github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
 github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
-github.com/gobwas/ws v1.0.3 h1:ZOigqf7iBxkA4jdQ3am7ATzdlOFp9YzA6NmuvEEZc9g=
-github.com/gobwas/ws v1.0.3/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
+github.com/gobwas/ws v1.0.4 h1:5eXU1CZhpQdq5kXbKb+sECH5Ia5KiO6CYzIzdlVx6Bs=
+github.com/gobwas/ws v1.0.4/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
@@ -31,8 +31,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
-github.com/urfave/cli/v2 v2.1.1 h1:Qt8FeAtxE/vfdrLmR3rxR6JRE0RoVmbXu8+6kZtYU4k=
-github.com/urfave/cli/v2 v2.1.1/go.mod h1:SE9GqnLQmjVa0iPEY0f1w3ygNIYcIJ0OKPMoW2caLfQ=
+github.com/urfave/cli/v2 v2.3.0 h1:qph92Y649prgesehzOrQjdWyxFOp/QVM+6imKHad91M=
+github.com/urfave/cli/v2 v2.3.0/go.mod h1:LJmUH05zAU44vOAcrfzZQKsZbVcdbOG8rtL3/XcUArI=
 github.com/xtaci/smux v2.0.1+incompatible h1:4NrCD5VzuFktMCxK08IShR0C5vKyNICJRShUzvk0U34=
 github.com/xtaci/smux v2.0.1+incompatible/go.mod h1:f+nYm6SpuHMy/SH0zpbvAFHT1QoMcgLOsWcFip5KfPw=
 go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk=
@@ -41,8 +41,8 @@ go.uber.org/multierr v1.5.0 h1:KCa4XfM8CWFCpxXRGok+Q0SS/0XBhMDbHHGABQLvD2A=
 go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
 go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee h1:0mgffUl7nfd+FpvXMVz4IDEaUSmT1ysygQC7qYo7sG4=
 go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
-go.uber.org/zap v1.15.0 h1:ZZCA22JRF2gQE5FoNmhmrf7jeJJ2uhqDUNRYKm8dvmM=
-go.uber.org/zap v1.15.0/go.mod h1:Mb2vm2krFEG5DV0W9qcHBYFtp/Wku1cvYaqPsS/WYfc=
+go.uber.org/zap v1.16.0 h1:uFRZXykJGK9lLY4HtgSw44DnIcAM+kRBP7x5m+NpAOM=
+go.uber.org/zap v1.16.0/go.mod h1:MA8QOfq0BHJwdXa996Y4dYkAqRKB8/1K1QMMZVaNZjQ=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
@@ -65,7 +65,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
-gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.3 h1:fvjTMHxHEw/mxHbtzPi3JCcKXQRAnQTBRo6YCJSVHKI=
+gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 honnef.co/go/tools v0.0.1-2019.2.3 h1:3JgtbtFHMiCmsznwGVTUWbgGov+pVqnlf1dEJTNAXeM=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=

internal/constant/constant.go

@@ -0,0 +1,21 @@
+package constant
+
+import "time"
+
+const (
+	MaxMWSSStreamCnt = 10
+	DialTimeOut      = 3 * time.Second
+	MaxConKeepAlive  = 10 * time.Minute
+
+	Listen_RAW  = "raw"
+	Listen_WS   = "ws"
+	Listen_WSS  = "wss"
+	Listen_MWSS = "mwss"
+
+	Transport_RAW  = "raw"
+	Transport_WS   = "ws"
+	Transport_WSS  = "wss"
+	Transport_MWSS = "mwss"
+
+	BUFFER_SIZE = 4 * 1024 // 4kb
+)

internal/lb/lb.go

@@ -10,10 +10,10 @@ type LBNode struct {
 	index         int
 }
 
-type LBNodeHeap []*LBNode
+type LBNodes []*LBNode
 
-func New(remotes []string) LBNodeHeap {
-	lh := make(LBNodeHeap, len(remotes))
+func New(remotes []string) *LBNodes {
+	lh := make(LBNodes, len(remotes))
 	for i, remote := range remotes {
 		lh[i] = &LBNode{
 			Remote:        remote,
@@ -22,29 +22,31 @@ func New(remotes []string) LBNodeHeap {
 		}
 	}
 	lh.HeapInit()
-	return lh
+	return &lh
 }
 
-func (lp LBNodeHeap) Len() int { return len(lp) }
+func (lp *LBNodes) Len() int { return len(*lp) }
 
-func (lp LBNodeHeap) Less(i, j int) bool {
-	return lp[i].OnLineUserCnt < lp[j].OnLineUserCnt
+func (lp *LBNodes) Less(i, j int) bool {
+	l := *lp
+	return l[i].OnLineUserCnt < l[j].OnLineUserCnt
 }
 
-func (lp LBNodeHeap) Swap(i, j int) {
-	lp[i], lp[j] = lp[j], lp[i]
-	lp[i].index = i
-	lp[j].index = j
+func (lp *LBNodes) Swap(i, j int) {
+	l := *lp
+	l[i], l[j] = l[j], l[i]
+	l[i].index = i
+	l[j].index = j
 }
 
-func (lp *LBNodeHeap) Push(x interface{}) {
+func (lp *LBNodes) Push(x interface{}) {
 	n := len(*lp)
 	node := x.(*LBNode)
 	node.index = n
 	*lp = append(*lp, node)
 }
 
-func (lp *LBNodeHeap) Pop() interface{} {
+func (lp *LBNodes) Pop() interface{} {
 	old := *lp
 	n := len(old)
 	node := old[n-1]
@@ -55,42 +57,47 @@ func (lp *LBNodeHeap) Pop() interface{} {
 }
 
 // update modifies the priority and value of an Item in the queue.
-func (lp *LBNodeHeap) update(node *LBNode, remote string, cnt int) {
+func (lp *LBNodes) update(node *LBNode, remote string, cnt int) {
 	node.Remote = remote
 	node.OnLineUserCnt = cnt
 	heap.Fix(lp, node.index)
 }
 
-func (lp *LBNodeHeap) HeapInit() {
+func (lp *LBNodes) HeapInit() {
 	heap.Init(lp)
 }
 
-func (lp *LBNodeHeap) HeapPush(node *LBNode) {
+func (lp *LBNodes) HeapPush(node *LBNode) {
 	heap.Push(lp, node)
 }
 
-func (lp *LBNodeHeap) HeapPop() *LBNode {
+func (lp *LBNodes) HeapPop() *LBNode {
 	return heap.Pop(lp).(*LBNode)
 }
 
-func (lp *LBNodeHeap) MinLBNode() *LBNode {
+func (lp *LBNodes) MinLBNode() *LBNode {
 	if lp.Len() > 0 {
 		old := *lp
 		return old[0]
 	}
 	return nil
 }
 
-func (lp *LBNodeHeap) IncrUserCnt(node *LBNode, num int) {
+func (lp *LBNodes) IncrUserCnt(node *LBNode, num int) {
 	lp.update(node, node.Remote, node.OnLineUserCnt+num)
 }
 
-func (lp *LBNodeHeap) PickMin() *LBNode {
+func (lp *LBNodes) PickMin() *LBNode {
 	node := lp.MinLBNode()
 	lp.IncrUserCnt(node, 1)
 	return node
 }
 
-func (lp *LBNodeHeap) DeferPick(node *LBNode) {
+func (lp *LBNodes) DeferPick(node *LBNode) {
 	lp.IncrUserCnt(node, -1)
 }
+
+func (lp *LBNodes) OnError(node *LBNode) {
+	// NOTE 遇到错误的时候降低这个节点的权重
+	lp.IncrUserCnt(node, 100)
+}

internal/lb/lb_test.go

@@ -10,7 +10,7 @@ func TestLBNodeHeap(t *testing.T) {
 		"1.1.1.1": 1, "4.4.4.4": 4, "3.3.3.3": 3,
 	}
 
-	lp := make(LBNodeHeap, len(nodes))
+	lp := make(LBNodes, len(nodes))
 	i := 0
 	for value, priority := range nodes {
 		lp[i] = &LBNode{