From a805567810a54b187aea458d8d64c59ecf26064b Mon Sep 17 00:00:00 2001
From: Ralph Slooten <axllent@gmail.com>
Date: Sun, 31 Mar 2024 00:06:25 +1300
Subject: [PATCH] Feature: Add readyz subcommand for Docker healthcheck (#270)

---
 Dockerfile                   |  2 +
 cmd/readyz.go                | 75 ++++++++++++++++++++++++++++++++++++
 cmd/root.go                  |  2 +-
 internal/storage/database.go |  5 +++
 server/handlers/k8sready.go  |  4 +-
 5 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 cmd/readyz.go

diff --git a/Dockerfile b/Dockerfile
index 39515deaf..fa7f27fe1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,4 +25,6 @@ RUN apk add --no-cache tzdata
 
 EXPOSE 1025/tcp 1110/tcp 8025/tcp
 
+HEALTHCHECK --interval=15s CMD /mailpit readyz
+
 ENTRYPOINT ["/mailpit"]
diff --git a/cmd/readyz.go b/cmd/readyz.go
new file mode 100644
index 000000000..d65a7a2ef
--- /dev/null
+++ b/cmd/readyz.go
@@ -0,0 +1,75 @@
+package cmd
+
+import (
+	"crypto/tls"
+	"fmt"
+	"net/http"
+	"os"
+	"path"
+	"strings"
+	"time"
+
+	"github.com/axllent/mailpit/config"
+	"github.com/spf13/cobra"
+)
+
+var (
+	useHTTPS bool
+)
+
+// readyzCmd represents the healthcheck command
+var readyzCmd = &cobra.Command{
+	Use:   "readyz",
+	Short: "Run a healthcheck to test if Mailpit is running",
+	Long: `This command connects to the /readyz endpoint of a running Mailpit server
+and exits with a status of 0 if the connection is successful, else with a 
+status 1 if unhealthy.
+
+If running within Docker, it should automatically detect environment
+settings to determine the HTTP bind interface & port.
+`,
+	Run: func(cmd *cobra.Command, args []string) {
+		webroot := strings.TrimRight(path.Join("/", config.Webroot, "/"), "/") + "/"
+		proto := "http"
+		if useHTTPS {
+			proto = "https"
+		}
+
+		uri := fmt.Sprintf("%s://%s%sreadyz", proto, config.HTTPListen, webroot)
+
+		conf := &http.Transport{
+			IdleConnTimeout:       time.Second * 5,
+			ExpectContinueTimeout: time.Second * 5,
+			TLSHandshakeTimeout:   time.Second * 5,
+			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
+		}
+		client := &http.Client{Transport: conf}
+
+		res, err := client.Get(uri)
+		if err != nil || res.StatusCode != 200 {
+			os.Exit(1)
+		}
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(readyzCmd)
+
+	if len(os.Getenv("MP_UI_BIND_ADDR")) > 0 {
+		config.HTTPListen = os.Getenv("MP_UI_BIND_ADDR")
+	}
+
+	if len(os.Getenv("MP_WEBROOT")) > 0 {
+		config.Webroot = os.Getenv("MP_WEBROOT")
+	}
+
+	config.UITLSCert = os.Getenv("MP_UI_TLS_CERT")
+
+	if config.UITLSCert != "" {
+		useHTTPS = true
+	}
+
+	readyzCmd.Flags().StringVarP(&config.HTTPListen, "listen", "l", config.HTTPListen, "Set the HTTP bind interface & port")
+	readyzCmd.Flags().StringVar(&config.Webroot, "webroot", config.Webroot, "Set the webroot for web UI & API")
+	readyzCmd.Flags().BoolVar(&useHTTPS, "https", useHTTPS, "Connect via HTTPS (ignores HTTPS validation)")
+}
diff --git a/cmd/root.go b/cmd/root.go
index c2d4c27fc..792ce3b1e 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -89,7 +89,7 @@ func init() {
 	rootCmd.Flags().BoolVarP(&logger.VerboseLogging, "verbose", "v", logger.VerboseLogging, "Verbose logging")
 
 	// Web UI / API
-	rootCmd.Flags().StringVarP(&config.HTTPListen, "listen", "l", config.HTTPListen, "HTTP bind interface and port for UI")
+	rootCmd.Flags().StringVarP(&config.HTTPListen, "listen", "l", config.HTTPListen, "HTTP bind interface & port for UI")
 	rootCmd.Flags().StringVar(&config.Webroot, "webroot", config.Webroot, "Set the webroot for web UI & API")
 	rootCmd.Flags().StringVar(&config.UIAuthFile, "ui-auth-file", config.UIAuthFile, "A password file for web UI & API authentication")
 	rootCmd.Flags().StringVar(&config.UITLSCert, "ui-tls-cert", config.UITLSCert, "TLS certificate for web UI (HTTPS) - requires ui-tls-key")
diff --git a/internal/storage/database.go b/internal/storage/database.go
index 0bc45bf95..583cffb0b 100644
--- a/internal/storage/database.go
+++ b/internal/storage/database.go
@@ -115,6 +115,11 @@ func Close() {
 	}
 }
 
+// Ping the database connection and return an error if unsuccessful
+func Ping() error {
+	return db.Ping()
+}
+
 // StatsGet returns the total/unread statistics for a mailbox
 func StatsGet() MailboxStats {
 	var (
diff --git a/server/handlers/k8sready.go b/server/handlers/k8sready.go
index 875b2b4ae..9ab850046 100644
--- a/server/handlers/k8sready.go
+++ b/server/handlers/k8sready.go
@@ -3,12 +3,14 @@ package handlers
 import (
 	"net/http"
 	"sync/atomic"
+
+	"github.com/axllent/mailpit/internal/storage"
 )
 
 // ReadyzHandler is a ready probe that signals k8s to be able to retrieve traffic
 func ReadyzHandler(isReady *atomic.Value) http.HandlerFunc {
 	return func(w http.ResponseWriter, _ *http.Request) {
-		if isReady == nil || !isReady.Load().(bool) {
+		if isReady == nil || !isReady.Load().(bool) || storage.Ping() != nil {
 			http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
 			return
 		}