From 3476daf32226f890fdc9a1dae2ace8d8bed401e3 Mon Sep 17 00:00:00 2001
From: timothy-spencer <timothy.spencer@gsa.gov>
Date: Wed, 20 Mar 2019 14:29:44 -0700
Subject: [PATCH] added an option to enable GCP healthcheck endpoints

---
 http.go    | 17 +++++++++++++++++
 main.go    | 10 +++++++++-
 options.go |  9 +++++----
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/http.go b/http.go
index 4456e39..3f9ba57 100644
--- a/http.go
+++ b/http.go
@@ -24,6 +24,23 @@ func (s *Server) ListenAndServe() {
 	}
 }
 
+// gcpHealthcheck handles healthcheck queries from GCP
+func gcpHealthcheck(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.EscapedPath() == "/liveness_check" {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("OK"))
+			return
+		}
+		if r.URL.EscapedPath() == "/readiness_check" {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("OK"))
+			return
+		}
+		h.ServeHTTP(w, r)
+	})
+}
+
 // ServeHTTP constructs a net.Listener and starts handling HTTP requests
 func (s *Server) ServeHTTP() {
 	HTTPAddress := s.Opts.HTTPAddress
diff --git a/main.go b/main.go
index 5625259..efe336c 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"math/rand"
+	"net/http"
 	"os"
 	"runtime"
 	"strings"
@@ -92,6 +93,7 @@ func main() {
 	flagSet.String("acr-values", "http://idmanagement.gov/ns/assurance/loa/1", "acr values string:  optional, used by login.gov")
 	flagSet.String("jwt-key", "", "private key used to sign JWT: required by login.gov")
 	flagSet.String("pubjwk-url", "", "JWK pubkey access endpoint: required by login.gov")
+	flagSet.Bool("gcp-healthchecks", false, "Enable GCP healthcheck endpoints")
 
 	flagSet.Parse(os.Args[1:])
 
@@ -139,8 +141,14 @@ func main() {
 
 	rand.Seed(time.Now().UnixNano())
 
+	var myhandler http.Handler
+	if opts.GCPHealthChecks {
+		myhandler = gcpHealthcheck(LoggingHandler(os.Stdout, oauthproxy, opts.RequestLogging, opts.RequestLoggingFormat))
+	} else {
+		myhandler = LoggingHandler(os.Stdout, oauthproxy, opts.RequestLogging, opts.RequestLoggingFormat)
+	}
 	s := &Server{
-		Handler: LoggingHandler(os.Stdout, oauthproxy, opts.RequestLogging, opts.RequestLoggingFormat),
+		Handler: myhandler,
 		Opts:    opts,
 	}
 	s.ListenAndServe()
diff --git a/options.go b/options.go
index 90af3d3..b736521 100644
--- a/options.go
+++ b/options.go
@@ -86,10 +86,11 @@ type Options struct {
 	RequestLogging       bool   `flag:"request-logging" cfg:"request_logging" env:"OAUTH2_PROXY_REQUEST_LOGGING"`
 	RequestLoggingFormat string `flag:"request-logging-format" cfg:"request_logging_format" env:"OAUTH2_PROXY_REQUEST_LOGGING_FORMAT"`
 
-	SignatureKey string `flag:"signature-key" cfg:"signature_key" env:"OAUTH2_PROXY_SIGNATURE_KEY"`
-	AcrValues    string `flag:"acr-values" cfg:"acr_values" env:"OAUTH2_PROXY_ACR_VALUES"`
-	JWTKey       string `flag:"jwt-key" cfg:"jwt_key" env:"OAUTH2_PROXY_JWT_KEY"`
-	PubJWKURL    string `flag:"pubjwk-url" cfg:"pubjwk_url" env:"OAUTH2_PROXY_PUBJWK_URL"`
+	SignatureKey    string `flag:"signature-key" cfg:"signature_key" env:"OAUTH2_PROXY_SIGNATURE_KEY"`
+	AcrValues       string `flag:"acr-values" cfg:"acr_values" env:"OAUTH2_PROXY_ACR_VALUES"`
+	JWTKey          string `flag:"jwt-key" cfg:"jwt_key" env:"OAUTH2_PROXY_JWT_KEY"`
+	PubJWKURL       string `flag:"pubjwk-url" cfg:"pubjwk_url" env:"OAUTH2_PROXY_PUBJWK_URL"`
+	GCPHealthChecks bool   `flag:"gcp-healthchecks" cfg:"gcp_healthchecks" env:"OAUTH2_PROXY_GCP_HEALTHCHECKS"`
 
 	// internal values that are set after config validation
 	redirectURL   *url.URL