From 082447f5170445e0bc69865fa37d756646c17511 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Mon, 16 Jul 2018 01:21:29 +0800
Subject: [PATCH] frpc: support health check

---
 client/health.go       | 123 +++++++++++++++++++++++++++++++++++++++--
 models/config/proxy.go |   2 +
 2 files changed, 121 insertions(+), 4 deletions(-)

diff --git a/client/health.go b/client/health.go
index ad58554d..8e84a6f8 100644
--- a/client/health.go
+++ b/client/health.go
@@ -15,18 +15,133 @@
 package client
 
 import (
-	"github.com/fatedier/frp/models/config"
+	"context"
+	"net"
+	"net/http"
+	"time"
 )
 
 type HealthCheckMonitor struct {
-	cfg config.HealthCheckConf
+	checkType      string
+	interval       time.Duration
+	timeout        time.Duration
+	maxFailedTimes int
+
+	// For tcp
+	addr string
+
+	// For http
+	url string
+
+	failedTimes    uint64
+	statusOK       bool
+	statusNormalFn func()
+	statusFailedFn func()
+
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
-func NewHealthCheckMonitor(cfg *config.HealthCheckConf) *HealthCheckMonitor {
+func NewHealthCheckMonitor(checkType string, intervalS int, timeoutS int, maxFailedTimes int, addr string, url string,
+	statusNormalFn func(), statusFailedFn func()) *HealthCheckMonitor {
+
+	if intervalS <= 0 {
+		intervalS = 10
+	}
+	if timeoutS <= 0 {
+		timeoutS = 3
+	}
+	if maxFailedTimes <= 0 {
+		maxFailedTimes = 1
+	}
+	ctx, cancel := context.WithCancel(context.Background())
 	return &HealthCheckMonitor{
-		cfg: *cfg,
+		checkType:      checkType,
+		interval:       time.Duration(intervalS) * time.Second,
+		timeout:        time.Duration(timeoutS) * time.Second,
+		maxFailedTimes: maxFailedTimes,
+		addr:           addr,
+		url:            url,
+		statusOK:       false,
+		statusNormalFn: statusNormalFn,
+		statusFailedFn: statusFailedFn,
+		ctx:            ctx,
+		cancel:         cancel,
 	}
 }
 
 func (monitor *HealthCheckMonitor) Start() {
+	go monitor.checkWorker()
+}
+
+func (monitor *HealthCheckMonitor) Stop() {
+	monitor.cancel()
+}
+
+func (monitor *HealthCheckMonitor) checkWorker() {
+	for {
+		ctx, cancel := context.WithDeadline(monitor.ctx, time.Now().Add(monitor.timeout))
+		ok := monitor.doCheck(ctx)
+
+		// check if this monitor has been closed
+		select {
+		case <-ctx.Done():
+			cancel()
+			return
+		default:
+			cancel()
+		}
+
+		if ok {
+			if !monitor.statusOK && monitor.statusNormalFn != nil {
+				monitor.statusOK = true
+				monitor.statusNormalFn()
+			}
+		} else {
+			monitor.failedTimes++
+			if monitor.statusOK && int(monitor.failedTimes) >= monitor.maxFailedTimes && monitor.statusFailedFn != nil {
+				monitor.statusOK = false
+				monitor.statusFailedFn()
+			}
+		}
+
+		time.Sleep(monitor.interval)
+	}
+}
+
+func (monitor *HealthCheckMonitor) doCheck(ctx context.Context) bool {
+	switch monitor.checkType {
+	case "tcp":
+		return monitor.doTcpCheck(ctx)
+	case "http":
+		return monitor.doHttpCheck(ctx)
+	default:
+		return false
+	}
+}
+
+func (monitor *HealthCheckMonitor) doTcpCheck(ctx context.Context) bool {
+	var d net.Dialer
+	conn, err := d.DialContext(ctx, "tcp", monitor.addr)
+	if err != nil {
+		return false
+	}
+	conn.Close()
+	return true
+}
+
+func (monitor *HealthCheckMonitor) doHttpCheck(ctx context.Context) bool {
+	req, err := http.NewRequest("GET", monitor.url, nil)
+	if err != nil {
+		return false
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return false
+	}
+
+	if resp.StatusCode/100 != 2 {
+		return false
+	}
+	return true
 }
diff --git a/models/config/proxy.go b/models/config/proxy.go
index b600be5c..9ea680a4 100644
--- a/models/config/proxy.go
+++ b/models/config/proxy.go
@@ -381,6 +381,8 @@ func (cfg *LocalSvrConf) checkForCli() (err error) {
 // Health check info
 type HealthCheckConf struct {
 	HealthCheckType      string `json:"health_check_type"` // tcp | http
+	HealthCheckTimeout   int    `json:"health_check_timeout"`
+	HealthCheckMaxFailed int    `json:"health_check_max_failed"`
 	HealthCheckIntervalS int    `json:"health_check_interval_s"`
 	HealthCheckUrl       string `json:"health_check_url"`