健康检查

lib/common/util.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"github.com/cnlh/nps/lib/crypt"
 	"github.com/cnlh/nps/lib/pool"
+	"html/template"
 	"io"
 	"io/ioutil"
 	"net"
@@ -276,3 +277,29 @@ func GetLocalUdpAddr() (net.Conn, error) {
 	}
 	return tmpConn, tmpConn.Close()
 }
+
+func ParseStr(str string) (string, error) {
+	tmp := template.New("npc")
+	var err error
+	w := new(bytes.Buffer)
+	if tmp, err = tmp.Parse(str); err != nil {
+		return "", err
+	}
+	if err = tmp.Execute(w, GetEnvMap()); err != nil {
+		return "", err
+	}
+	return w.String(), nil
+}
+
+//get env
+func GetEnvMap() map[string]string {
+	m := make(map[string]string)
+	environ := os.Environ()
+	for i := range environ {
+		tmp := strings.Split(environ[i], "=")
+		if len(tmp) == 2 {
+			m[tmp[0]] = tmp[1]
+		}
+	}
+	return m
+}

lib/config/config.go

@@ -38,7 +38,9 @@ func NewConfig(path string) (c *Config, err error) {
 	if b, err = common.ReadAllFromFile(path); err != nil {
 		return
 	} else {
-		c.content = string(b)
+		if c.content, err = common.ParseStr(string(b)); err != nil {
+			return nil, err
+		}
 		if c.title, err = getAllTitle(c.content); err != nil {
 			return
 		}
@@ -153,6 +155,14 @@ func dealHost(s string) *file.Host {
 			h.HostChange = item[1]
 		case "location":
 			h.Location = item[1]
+		case "health_check_timeout":
+			h.HealthCheckTimeout = common.GetIntNoErrByStr(item[1])
+		case "health_check_max_failed":
+			h.HealthMaxFail = common.GetIntNoErrByStr(item[1])
+		case "health_check_interval":
+			h.HealthCheckInterval = common.GetIntNoErrByStr(item[1])
+		case "health_http_url":
+			h.HttpHealthUrl = item[1]
 		default:
 			if strings.Contains(item[0], "header") {
 				headerChange += strings.Replace(item[0], "header_", "", -1) + ":" + item[1] + "\n"
@@ -178,7 +188,7 @@ func dealTunnel(s string) *file.Tunnel {
 		case "mode":
 			t.Mode = item[1]
 		case "target":
-			t.Target = item[1]
+			t.Target = strings.Replace(item[1], ",", "\n", -1)
 		case "targetAddr":
 			t.TargetAddr = item[1]
 		case "password":
@@ -187,6 +197,12 @@ func dealTunnel(s string) *file.Tunnel {
 			t.LocalPath = item[1]
 		case "strip_pre":
 			t.StripPre = item[1]
+		case "health_check_timeout":
+			t.HealthCheckTimeout = common.GetIntNoErrByStr(item[1])
+		case "health_check_max_failed":
+			t.HealthMaxFail = common.GetIntNoErrByStr(item[1])
+		case "health_check_interval":
+			t.HealthCheckInterval = common.GetIntNoErrByStr(item[1])
 		}
 	}
 	return t

lib/file/obj.go

@@ -4,6 +4,7 @@ import (
 	"github.com/cnlh/nps/lib/rate"
 	"strings"
 	"sync"
+	"time"
 )
 
 type Flow struct {
@@ -96,14 +97,15 @@ func (s *Client) HasHost(h *Host) bool {
 }
 
 type Tunnel struct {
-	Id         int     //Id
-	Port       int     //服务端监听端口
-	Mode       string  //启动方式
-	Target     string  //目标
-	Status     bool    //设置是否开启
-	RunStatus  bool    //当前运行状态
-	Client     *Client //所属客户端id
-	Ports      string  //客户端与服务端传递
+	Id         int      //Id
+	Port       int      //服务端监听端口
+	Mode       string   //启动方式
+	Target     string   //目标
+	TargetArr  []string //目标
+	Status     bool     //设置是否开启
+	RunStatus  bool     //当前运行状态
+	Client     *Client  //所属客户端id
+	Ports      string   //客户端与服务端传递
 	Flow       *Flow
 	Password   string //私密模式密码，唯一
 	Remark     string //备注
@@ -111,6 +113,35 @@ type Tunnel struct {
 	NoStore    bool
 	LocalPath  string
 	StripPre   string
+	NowIndex   int
+	Health
+	sync.RWMutex
+}
+
+type Health struct {
+	HealthCheckTimeout  int
+	HealthMaxFail       int
+	HealthCheckInterval int
+	HealthNextTime      time.Time
+	HealthMap           map[string]int
+	HttpHealthUrl       string
+	HealthRemoveArr     []string
+}
+
+func (s *Tunnel) GetRandomTarget() string {
+	if s.TargetArr == nil {
+		s.TargetArr = strings.Split(s.Target, "\n")
+	}
+	if len(s.TargetArr) == 1 {
+		return s.TargetArr[0]
+	}
+	s.Lock()
+	defer s.Unlock()
+	if s.NowIndex >= len(s.TargetArr)-1 {
+		s.NowIndex = -1
+	}
+	s.NowIndex++
+	return s.TargetArr[s.NowIndex]
 }
 
 type Config struct {
@@ -133,7 +164,8 @@ type Host struct {
 	NowIndex     int
 	TargetArr    []string
 	NoStore      bool
-	Scheme        string //http https all
+	Scheme       string //http https all
+	Health
 	sync.RWMutex
 }
 
@@ -141,10 +173,13 @@ func (s *Host) GetRandomTarget() string {
 	if s.TargetArr == nil {
 		s.TargetArr = strings.Split(s.Target, "\n")
 	}
+	if len(s.TargetArr) == 1 {
+		return s.TargetArr[0]
+	}
 	s.Lock()
 	defer s.Unlock()
 	if s.NowIndex >= len(s.TargetArr)-1 {
-		s.NowIndex = 0
+		s.NowIndex = -1
 	} else {
 		s.NowIndex++
 	}

lib/mux/pmux.go

@@ -11,18 +11,20 @@ import (
 	"net"
 	"strconv"
 	"strings"
+	"time"
 )
 
 const (
-	HTTP_GET     = 716984
-	HTTP_POST    = 807983
-	HTTP_HEAD    = 726965
-	HTTP_PUT     = 808585
-	HTTP_DELETE  = 686976
-	HTTP_CONNECT = 677978
-	HTTP_OPTIONS = 798084
-	HTTP_TRACE   = 848265
-	CLIENT       = 848384
+	HTTP_GET        = 716984
+	HTTP_POST       = 807983
+	HTTP_HEAD       = 726965
+	HTTP_PUT        = 808585
+	HTTP_DELETE     = 686976
+	HTTP_CONNECT    = 677978
+	HTTP_OPTIONS    = 798084
+	HTTP_TRACE      = 848265
+	CLIENT          = 848384
+	ACCEPT_TIME_OUT = 10
 )
 
 type PortMux struct {
@@ -122,7 +124,11 @@ func (pMux *PortMux) process(conn net.Conn) {
 	if len(rs) == 0 {
 		rs = buf
 	}
-	ch <- newPortConn(conn, rs)
+	timer := time.NewTimer(ACCEPT_TIME_OUT)
+	select {
+	case <-timer.C:
+	case ch <- newPortConn(conn, rs):
+	}
 }
 
 func (pMux *PortMux) Close() error {

lib/sheap/heap.go

@@ -0,0 +1,21 @@
+package sheap
+
+type IntHeap []int64
+
+func (h IntHeap) Len() int           { return len(h) }
+func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
+func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *IntHeap) Push(x interface{}) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	*h = append(*h, x.(int64))
+}
+
+func (h *IntHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}