Add support for Kubernetes liveness/readiness probes.

fingerproxy.go

@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
@@ -139,6 +140,7 @@ func Run() {
 		envWithDefault("CERT_FILENAME", "tls.crt"),
 		"TLS certificate filename, equivalent to $CERT_FILENAME",
 	)
+
 	flagKeyFilename := flag.String(
 		"certkey-filename",
 		envWithDefault("CERTKEY_FILENAME", "tls.key"),
@@ -151,7 +153,18 @@ func Run() {
 		"Listening address of Prometheus metrics, equivalent to $METRICS_LISTEN_ADDR",
 	)
 
-	flagVerboseLogs := flag.Bool("verbose", false, "Enable verbose logs")
+	flagEnableKubernetesProbe := flag.Bool(
+		"enable-kubernetes-probe",
+		envWithDefaultBool("ENABLE_KUBERNETES_PROBE", true),
+		"Enable kubernetes liveness/readiness probe support, equivalent to $ENABLE_KUBERNETES_PROBE",
+	)
+
+	flagVerboseLogs := flag.Bool(
+		"verbose",
+		envWithDefaultBool("VERBOSE", false),
+		"Enable verbose logs, equivalent to $VERBOSE",
+	)
+
 	flagVersion := flag.Bool("version", false, "Print version and exit")
 	flag.Parse()
 
@@ -173,8 +186,13 @@ func Run() {
 
 	InitFingerprint(*flagVerboseLogs)
 
+	handler := GetReverseProxyHTTPHandler(forwardTo)
+	if *flagEnableKubernetesProbe {
+		handler.IsProbeRequest = reverseproxy.IsKubernetesProbeRequest
+	}
+
 	server := DefaultProxyServer(
-		GetReverseProxyHTTPHandler(forwardTo),
+		handler,
 		tlsConfig,
 		*flagVerboseLogs,
 	)
@@ -193,3 +211,14 @@ func envWithDefault(key string, defaultVal string) string {
 	}
 	return defaultVal
 }
+
+func envWithDefaultBool(key string, defaultVal bool) bool {
+	if envVal, ok := os.LookupEnv(key); ok {
+		if strings.ToLower(envVal) == "true" {
+			return true
+		} else if strings.ToLower(envVal) == "false" {
+			return false
+		}
+	}
+	return defaultVal
+}

pkg/reverseproxy/handler.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"strings"
 )
 
 type HTTPHandler struct {
@@ -19,8 +20,18 @@ type HTTPHandler struct {
 
 	// optional, but in fact required, injecting fingerprint headers to outbound requests
 	HeaderInjectors []HeaderInjector
+
+	// optional, if IsProbeRequest returns true, handler will respond with
+	// a HTTP 200 OK instead of forwarding requests, useful for kubernetes
+	// liveness/readiness probes. defaults to nil, which disables this behavior
+	IsProbeRequest func(*http.Request) bool
 }
 
+const (
+	ProbeStatusCode = http.StatusOK
+	ProbeResponse   = "OK"
+)
+
 func NewHTTPHandler(to *url.URL, reverseProxy *httputil.ReverseProxy, headerInjectors []HeaderInjector) *HTTPHandler {
 	f := &HTTPHandler{
 		To:              to,
@@ -55,10 +66,19 @@ func (f *HTTPHandler) rewriteFunc(r *httputil.ProxyRequest) {
 }
 
 func (f *HTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	// TODO: add liveness probe
+	if f.IsProbeRequest != nil && f.IsProbeRequest(req) {
+		w.WriteHeader(ProbeStatusCode)
+		w.Write([]byte(ProbeResponse))
+		return
+	}
 	f.ReverseProxy.ServeHTTP(w, req)
 }
 
+func IsKubernetesProbeRequest(r *http.Request) bool {
+	// https://github.com/kubernetes/kubernetes/blob/656cb1028ea5af837e69b5c9c614b008d747ab63/pkg/probe/http/request.go#L91
+	return strings.HasPrefix(r.UserAgent(), "kube-probe/")
+}
+
 func (f *HTTPHandler) logf(format string, args ...any) {
 	if f.ReverseProxy.ErrorLog != nil {
 		f.ReverseProxy.ErrorLog.Printf(format, args...)