add a health-check server

Author: metachris

cmd/proxy-server/main.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
-	"log"
+	"log/slog"
 	"net/http"
 	"os"
 	"os/signal"
@@ -24,6 +24,12 @@ var flags []cli.Flag = []cli.Flag{
 		Value:   "127.0.0.1:8080",
 		Usage:   "address to listen on",
 	},
+	&cli.StringFlag{
+		Name:    "listen-addr-healthcheck",
+		EnvVars: []string{"LISTEN_ADDR_HEALTHCHECK"},
+		Value:   "",
+		Usage:   "address to listen on for health checks",
+	},
 	&cli.StringFlag{
 		Name:    "target-addr",
 		EnvVars: []string{"TARGET_ADDR"},
@@ -39,12 +45,12 @@ var flags []cli.Flag = []cli.Flag{
 	&cli.StringFlag{
 		Name:    "tls-certificate",
 		EnvVars: []string{"TLS_CERTIFICATE"},
-		Usage:   "Certificate to present (PEM). Only valid for --server-attestation-type=none and with --tls-private-key.",
+		Usage:   "Path to TLS certificate (PEM). Only valid for --server-attestation-type=none and with --tls-private-key.",
 	},
 	&cli.StringFlag{
 		Name:    "tls-private-key",
 		EnvVars: []string{"TLS_PRIVATE_KEY"},
-		Usage:   "Private key for the certificate (PEM). Only valid with --tls-certificate.",
+		Usage:   "Path to private key for the certificate. Only valid with --tls-certificate.",
 	},
 	&cli.StringFlag{
 		Name:    "client-attestation-type",
@@ -71,6 +77,8 @@ var flags []cli.Flag = []cli.Flag{
 	},
 }
 
+var log *slog.Logger
+
 func main() {
 	app := &cli.App{
 		Name:   "proxy-server",
@@ -80,7 +88,11 @@ func main() {
 	}
 
 	if err := app.Run(os.Args); err != nil {
-		log.Fatal(err)
+		if log != nil {
+			log.Error("error running app", "err", err)
+		} else {
+			panic(err)
+		}
 	}
 }
 
@@ -95,7 +107,7 @@ func runServer(cCtx *cli.Context) error {
 	certFile := cCtx.String("tls-certificate")
 	keyFile := cCtx.String("tls-private-key")
 
-	log := common.SetupLogger(&common.LoggingOpts{
+	log = common.SetupLogger(&common.LoggingOpts{
 		Debug:   logDebug,
 		JSON:    logJSON,
 		Service: "proxy-server",
@@ -196,6 +208,12 @@ func runServer(cCtx *cli.Context) error {
 		}
 	}()
 
+	// Start the health check server
+	listenAddrHealthCheck := cCtx.String("listen-addr-healthcheck")
+	if listenAddrHealthCheck != "" {
+		go startHealthCheckServer(listenAddrHealthCheck)
+	}
+
 	log.With("listenAddr", listenAddr).Info("Starting proxy server")
 	err = server.Serve(tlsListener)
 	if err != nil {
@@ -205,3 +223,17 @@ func runServer(cCtx *cli.Context) error {
 
 	return nil
 }
+
+func startHealthCheckServer(listenAddrHealthCheck string) {
+	log.With("healthCheckListenAddr", listenAddrHealthCheck).Info("Starting health check server")
+	healthCheckServer := &http.Server{
+		Addr: listenAddrHealthCheck,
+		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}),
+	}
+	err := healthCheckServer.ListenAndServe()
+	if err != nil {
+		log.Error("could not start health check server", "err", err)
+	}
+}