if port is not defined for live check, mount the liveness check on th…

…e main router protect liveness check allow only GET requests
TykTechnologies · Mar 22, 2019 · daee449 · daee449
1 parent 832ff32
commit daee449
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 67 deletions.
diff --git a/api_healthcheck.go b/api_healthcheck.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"net/http"
 	"strconv"
 	"strings"
 	"sync"
@@ -138,3 +139,42 @@ func (h *DefaultHealthChecker) ApiHealthValues() (HealthCheckValues, error) {
 	values.AvgUpstreamLatency = roundValue(float64(runningTotal / len(vals)))
 	return values, nil
 }
+
+func liveCheck(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		doJSONWrite(w, http.StatusMethodNotAllowed, apiError(http.StatusText(http.StatusMethodNotAllowed)))
+		return
+	}
+
+	redisStore := storage.RedisCluster{KeyPrefix: "livenesscheck-"}
+
+	key := "tyk-liveness-probe"
+
+	err := redisStore.SetRawKey(key, key, 10)
+	if err != nil {
+		mainLog.WithField("liveness-check", true).Error(err)
+		doJSONWrite(w, 500, apiError("Gateway is not connected to Redis. An error occurred while writing key to Redis"))
+		return
+	}
+
+	redisStore.DeleteRawKey(key)
+
+	if config.Global().UseDBAppConfigs {
+		if err = DashService.Ping(); err != nil {
+			doJSONWrite(w, 500, apiError("Dashboard is down. Gateway cannot connect to the dashboard"))
+			return
+		}
+	}
+
+	if config.Global().Policies.PolicySource == "rpc" {
+		rpcStore := RPCStorageHandler{KeyPrefix: "livenesscheck-"}
+
+		if !rpcStore.Connect() {
+			doJSONWrite(w, 500, apiError("RPC connection is down!!!"))
+			return
+		}
+	}
+
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte(`Gateway is alive!!!!`))
+}
diff --git a/api_loader.go b/api_loader.go
@@ -4,6 +4,7 @@ import (
 	"crypto/md5"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"path/filepath"
@@ -668,6 +669,32 @@ func loadApps(specs []*APISpec, muxer *mux.Router) {
 		fmt.Fprint(w, "Hello Tiki")
 	})
 
+	if config.Global().LivenessCheck.Enabled {
+		handler := checkIsAPIOwner(http.HandlerFunc(liveCheck))
+
+		if config.Global().LivenessCheck.Port <= 0 {
+			// If not port was defined, mount the liveness check on
+			// the main router
+			muxer.Handle("/status", handler)
+		} else {
+			l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.Global().ListenAddress, config.Global().LivenessCheck.Port))
+			if err != nil {
+				mainLog.Errorf("an error occurred while creating the liveness check router.... %v", err)
+			} else {
+				livenessRouter := mux.NewRouter()
+				livenessRouter.Handle("/status", handler)
+
+				srv := &http.Server{
+					ReadTimeout:  time.Second * 5,
+					WriteTimeout: time.Second * 5,
+					Handler:      livenessRouter,
+				}
+
+				go srv.Serve(l)
+			}
+		}
+	}
+
 	// Swap in the new register
 	apisMu.Lock()
 	apisByID = tmpSpecRegister

diff --git a/main.go b/main.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"crypto/tls"
-	"fmt"
 	"html/template"
 	"io/ioutil"
 	stdlog "log"
@@ -1391,73 +1390,7 @@ func listen(listener, controlListener net.Listener, err error) {
 	mainLog.Info("--> Listening on port: ", config.Global().ListenPort)
 	mainLog.Info("--> PID: ", hostDetails.PID)
 
-	mainRouter.HandleFunc("/"+config.Global().HealthCheckEndpointName, func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprintf(w, "Hello Tiki")
-	})
-
-	var livenessRouter *mux.Router
-	if config.Global().LivenessCheck.Enabled {
-		var l net.Listener
-
-		if config.Global().LivenessCheck.Port <= 0 {
-			l = listener
-		} else {
-			l, err = net.Listen("tcp", fmt.Sprintf("%s:%d", config.Global().ListenAddress, config.Global().LivenessCheck.Port))
-			if err != nil {
-				mainLog.Errorf("an error occurred while creating the liveness check router.... %v", err)
-			} else {
-				livenessRouter = mux.NewRouter()
-				livenessRouter.HandleFunc("/status", liveCheck)
-			}
-		}
-
-		if l != nil {
-			srv := &http.Server{
-				ReadTimeout:  time.Second * 5,
-				WriteTimeout: time.Second * 5,
-				Handler:      livenessRouter,
-			}
-
-			go srv.Serve(l)
-		}
-	}
-
 	if !rpc.IsEmergencyMode() {
 		doReload()
 	}
 }
-
-func liveCheck(w http.ResponseWriter, r *http.Request) {
-
-	redisStore := storage.RedisCluster{KeyPrefix: "livenesscheck-"}
-
-	key := "tyk-liveness-probe"
-
-	err := redisStore.SetRawKey(key, key, 10)
-	if err != nil {
-		mainLog.WithField("liveness-check", true).Error(err)
-		doJSONWrite(w, 500, apiError("Gateway is not connected to Redis. An error occurred while writing key to Redis"))
-		return
-	}
-
-	redisStore.DeleteRawKey(key)
-
-	if config.Global().UseDBAppConfigs {
-		if err = DashService.Ping(); err != nil {
-			doJSONWrite(w, 500, apiError("Dashboard is down. Gateway cannot connect to the dashboard"))
-			return
-		}
-	}
-
-	if config.Global().Policies.PolicySource == "rpc" {
-		rpcStore := RPCStorageHandler{KeyPrefix: "livenesscheck-"}
-
-		if !rpcStore.Connect() {
-			doJSONWrite(w, 500, apiError("RPC connection is down!!!"))
-			return
-		}
-	}
-
-	w.WriteHeader(http.StatusOK)
-	w.Write([]byte(`Gateway is alive!!!!`))
-}