From 77d21b681bdd477b95d039c93a4bb4c149822fb9 Mon Sep 17 00:00:00 2001
From: Ross Smith <rsmith@pomerium.com>
Date: Tue, 4 Feb 2025 19:19:15 -0500
Subject: [PATCH] add a simple check for databroker record as a readiness check

---
 proxy/handlers.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/proxy/handlers.go b/proxy/handlers.go
index 3b0b19ec4..3d0b4ec2f 100644
--- a/proxy/handlers.go
+++ b/proxy/handlers.go
@@ -17,6 +17,7 @@ import (
 	"github.com/pomerium/pomerium/internal/middleware"
 	"github.com/pomerium/pomerium/internal/telemetry/trace"
 	"github.com/pomerium/pomerium/internal/urlutil"
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 )
 
 // registerDashboardHandlers returns the proxy service's ServeMux
@@ -40,6 +41,15 @@ func (p *Proxy) registerDashboardHandlers(r *mux.Router, opts *config.Options) *
 	c := r.PathPrefix(dashboardPath + "/callback").Subrouter()
 	c.Path("/").Handler(httputil.HandlerFunc(p.Callback)).Methods(http.MethodGet)
 
+	// Handlers for checking whether the proxy is ready to serve traffic
+	r.Path(dashboardPath + "/readyz").
+		Handler(httputil.HandlerFunc(func(w http.ResponseWriter, r *http.Request) error {
+			if r.Method != http.MethodGet {
+				http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
+			}
+			return p.ReadyZ(w, r)
+		}))
+
 	// Programmatic API handlers and middleware
 	// gorilla mux has a bug that prevents HTTP 405 errors from being returned properly so we do all this manually
 	// https://github.com/gorilla/mux/issues/739
@@ -122,6 +132,29 @@ func (p *Proxy) Callback(w http.ResponseWriter, r *http.Request) error {
 	return p.state.Load().authenticateFlow.Callback(w, r)
 }
 
+// Readyz provides some information about whether the proxy is configured and ready
+// to serve requests
+func (p *Proxy) ReadyZ(w http.ResponseWriter, r *http.Request) error {
+	client := p.state.Load().dataBrokerClient
+
+	resp, err := client.Get(r.Context(), &databroker.GetRequest{
+		Type: "type.googleapis.com/pomerium.config.Config",
+		Id:   "dashboard-settings",
+	})
+	if err != nil {
+		return httputil.NewError(http.StatusInternalServerError, err)
+	}
+
+	if resp.GetRecord().GetData() != nil {
+		w.WriteHeader(http.StatusOK)
+		_, _ = io.WriteString(w, "data record from console")
+		return nil
+	}
+
+	http.Error(w, "No data record from console", http.StatusTeapot)
+	return nil
+}
+
 // ProgrammaticLogin returns a signed url that can be used to login
 // using the authenticate service.
 func (p *Proxy) ProgrammaticLogin(w http.ResponseWriter, r *http.Request) error {