From 2aacb311bc79a89d7bfca918a9e10bcc705816ec Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Tue, 5 Oct 2021 22:19:33 +0200
Subject: [PATCH] internal: add internal healthchecking to prevent websocket
 errors

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 cmd/server/main.go              |  2 +-
 internal/gounicorn/gounicorn.go | 38 +++++++++++++++++++++++++++++++++
 internal/web/proxy.go           | 10 ++++++++-
 internal/web/web.go             |  5 ++++-
 lifecycle/ak                    |  2 +-
 5 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 91277a0f1..06e452dd4 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -54,7 +54,7 @@ func main() {
 	u, _ := url.Parse("http://localhost:8000")
 
 	g := gounicorn.NewGoUnicorn()
-	ws := web.NewWebServer()
+	ws := web.NewWebServer(g)
 	defer g.Kill()
 	defer ws.Shutdown()
 	go web.RunMetricsServer()
diff --git a/internal/gounicorn/gounicorn.go b/internal/gounicorn/gounicorn.go
index 335aa48b9..422d15421 100644
--- a/internal/gounicorn/gounicorn.go
+++ b/internal/gounicorn/gounicorn.go
@@ -1,10 +1,13 @@
 package gounicorn
 
 import (
+	"net/http"
 	"os"
 	"os/exec"
+	"time"
 
 	log "github.com/sirupsen/logrus"
+	"goauthentik.io/internal/outpost/ak"
 )
 
 type GoUnicorn struct {
@@ -12,6 +15,7 @@ type GoUnicorn struct {
 	p       *exec.Cmd
 	started bool
 	killed  bool
+	alive   bool
 }
 
 func NewGoUnicorn() *GoUnicorn {
@@ -20,6 +24,7 @@ func NewGoUnicorn() *GoUnicorn {
 		log:     logger,
 		started: false,
 		killed:  false,
+		alive:   false,
 	}
 	g.initCmd()
 	return g
@@ -35,6 +40,10 @@ func (g *GoUnicorn) initCmd() {
 	g.p.Stderr = os.Stderr
 }
 
+func (g *GoUnicorn) IsRunning() bool {
+	return g.alive
+}
+
 func (g *GoUnicorn) Start() error {
 	if g.killed {
 		g.log.Debug("Not restarting gunicorn since we're killed")
@@ -44,9 +53,38 @@ func (g *GoUnicorn) Start() error {
 		g.initCmd()
 	}
 	g.started = true
+	go g.healthcheck()
 	return g.p.Run()
 }
 
+func (g *GoUnicorn) healthcheck() {
+	g.log.Debug("starting healthcheck")
+	h := &http.Client{
+		Transport: ak.NewUserAgentTransport("goauthentik.io go proxy healthcheck", http.DefaultTransport),
+	}
+	check := func() bool {
+		res, err := h.Get("http://localhost:8000/-/health/live/")
+		if err == nil && res.StatusCode == 204 {
+			g.alive = true
+			return true
+		}
+		return false
+	}
+
+	// Default healthcheck is every 1 second on startup
+	// once we've been healthy once, increase to 30 seconds
+	for range time.Tick(time.Second) {
+		if check() {
+			g.log.Info("backend is alive, backing off with healthchecks")
+			break
+		}
+		g.log.Debug("backend not alive yet")
+	}
+	for range time.Tick(30 * time.Second) {
+		check()
+	}
+}
+
 func (g *GoUnicorn) Kill() {
 	g.killed = true
 	err := g.p.Process.Kill()
diff --git a/internal/web/proxy.go b/internal/web/proxy.go
index c33302894..e85196edd 100644
--- a/internal/web/proxy.go
+++ b/internal/web/proxy.go
@@ -40,6 +40,10 @@ func (ws *WebServer) configureProxy() {
 		ws.proxyErrorHandler(rw, r, fmt.Errorf("proxy not running"))
 	})
 	ws.m.PathPrefix("/").HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		if !ws.p.IsRunning() {
+			ws.proxyErrorHandler(rw, r, fmt.Errorf("authentik core not running yet"))
+			return
+		}
 		host := web.GetHost(r)
 		before := time.Now()
 		if ws.ProxyServer != nil {
@@ -59,8 +63,12 @@ func (ws *WebServer) configureProxy() {
 }
 
 func (ws *WebServer) proxyErrorHandler(rw http.ResponseWriter, req *http.Request, err error) {
-	ws.log.WithError(err).Warning("proxy error")
+	ws.log.Warning(err.Error())
 	rw.WriteHeader(http.StatusBadGateway)
+	_, err = rw.Write([]byte("authentik starting..."))
+	if err != nil {
+		ws.log.WithError(err).Warning("failed to write error message")
+	}
 }
 
 func (ws *WebServer) proxyModifyResponse(r *http.Response) error {
diff --git a/internal/web/web.go b/internal/web/web.go
index 2c199711b..60ca4afed 100644
--- a/internal/web/web.go
+++ b/internal/web/web.go
@@ -11,6 +11,7 @@ import (
 	"github.com/pires/go-proxyproto"
 	log "github.com/sirupsen/logrus"
 	"goauthentik.io/internal/config"
+	"goauthentik.io/internal/gounicorn"
 	"goauthentik.io/internal/outpost/proxyv2"
 )
 
@@ -27,9 +28,10 @@ type WebServer struct {
 	m   *mux.Router
 	lh  *mux.Router
 	log *log.Entry
+	p   *gounicorn.GoUnicorn
 }
 
-func NewWebServer() *WebServer {
+func NewWebServer(g *gounicorn.GoUnicorn) *WebServer {
 	l := log.WithField("logger", "authentik.g.web")
 	mainHandler := mux.NewRouter()
 	if config.G.ErrorReporting.Enabled {
@@ -46,6 +48,7 @@ func NewWebServer() *WebServer {
 		m:   mainHandler,
 		lh:  logginRouter,
 		log: l,
+		p:   g,
 	}
 	ws.configureStatic()
 	ws.configureProxy()
diff --git a/lifecycle/ak b/lifecycle/ak
index d0beb5e15..f64320d98 100755
--- a/lifecycle/ak
+++ b/lifecycle/ak
@@ -49,7 +49,7 @@ elif [[ "$1" == "test" ]]; then
 elif [[ "$1" == "healthcheck" ]]; then
     mode=$(cat $MODE_FILE)
     if [[ $mode == "server" ]]; then
-        curl --user-agent "authentik Healthcheck" -I http://localhost:9000/-/health/ready/
+        curl --user-agent "goauthentik.io lifecycle Healthcheck" -I http://localhost:9000/-/health/ready/
     elif [[ $mode == "worker" ]]; then
         celery -A authentik.root.celery inspect ping -d celery@$HOSTNAME
     fi