Allow a health/ping request to be identified by User-Agent (#567)

* Add an option to allow health checks based on User-Agent.

* Formatting fix

* Rename field and avoid unnecessary interface.

* Skip the redirect fix so it can be put into a different PR.

* Add CHANGELOG entry

* Adding a couple tests for the PingUserAgent option.

CHANGELOG.md

@@ -58,6 +58,7 @@
 - [#560](https://github.com/oauth2-proxy/oauth2-proxy/pull/560) Fallback to UserInfo is User ID claim not present (@JoelSpeed)
 - [#598](https://github.com/oauth2-proxy/oauth2-proxy/pull/598) acr_values no longer sent to IdP when empty (@ScottGuymer)
 - [#548](https://github.com/oauth2-proxy/oauth2-proxy/pull/548) Separate logging options out of main options structure (@JoelSpeed)
+- [#567](https://github.com/oauth2-proxy/oauth2-proxy/pull/567) Allow health/ping request to be identified via User-Agent (@chkohner)
 - [#536](https://github.com/oauth2-proxy/oauth2-proxy/pull/536) Improvements to Session State code (@JoelSpeed)
 - [#573](https://github.com/oauth2-proxy/oauth2-proxy/pull/573) Properly parse redis urls for cluster and sentinel connections (@amnay-mo)
 - [#574](https://github.com/oauth2-proxy/oauth2-proxy/pull/574) render error page on 502 proxy status (@amnay-mo)

docs/configuration/configuration.md

@@ -88,6 +88,7 @@ An example [oauth2-proxy.cfg]({{ site.gitweb }}/contrib/oauth2-proxy.cfg.example
 | `--provider` | string | OAuth provider | google |
 | `--provider-display-name` | string | Override the provider's name with the given string; used for the sign-in page | (depends on provider) |
 | `--ping-path` | string | the ping endpoint that can be used for basic health checks | `"/ping"` |
+| `--ping-user-agent` | string | a User-Agent that can be used for basic health checks | `""` (don't check user agent) |
 | `--proxy-prefix` | string | the url root path that this proxy should be nested under (e.g. /`<oauth2>/sign_in`) | `"/oauth2"` |
 | `--proxy-websockets` | bool | enables WebSocket proxying | true |
 | `--pubjwk-url` | string | JWK pubkey access endpoint: required by login.gov | |
@@ -163,7 +164,7 @@ There are three different types of logging: standard, authentication, and HTTP r
 
 Each type of logging has their own configurable format and variables. By default these formats are similar to the Apache Combined Log.
 
-Logging of requests to the `/ping` endpoint can be disabled with `--silence-ping-logging` reducing log volume. This flag appends the `--ping-path` to `--exclude-logging-paths`.
+Logging of requests to the `/ping` endpoint (or using `--ping-user-agent`) can be disabled with `--silence-ping-logging` reducing log volume. This flag appends the `--ping-path` to `--exclude-logging-paths`.
 
 ### Auth Log Format
 Authentication logs are logs which are guaranteed to contain a username or email address of a user attempting to authenticate. These logs are output by default in the below format:

logging_handler.go

@@ -21,6 +21,7 @@ type responseLogger struct {
 	size     int
 	upstream string
 	authInfo string
+	silent   bool
 }
 
 // Header returns the ResponseWriter's Header
@@ -104,5 +105,7 @@ func (h loggingHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	url := *req.URL
 	responseLogger := &responseLogger{w: w}
 	h.handler.ServeHTTP(responseLogger, req)
-	logger.PrintReq(responseLogger.authInfo, responseLogger.upstream, req, url, t, responseLogger.Status(), responseLogger.Size())
+	if !responseLogger.silent {
+		logger.PrintReq(responseLogger.authInfo, responseLogger.upstream, req, url, t, responseLogger.Status(), responseLogger.Size())
+	}
 }

logging_handler_test.go

@@ -5,11 +5,14 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/oauth2-proxy/oauth2-proxy/pkg/apis/options"
 	"github.com/oauth2-proxy/oauth2-proxy/pkg/logger"
+	"github.com/oauth2-proxy/oauth2-proxy/pkg/validation"
 )
 
 func TestLoggingHandler_ServeHTTP(t *testing.T) {
@@ -67,3 +70,59 @@ func TestLoggingHandler_ServeHTTP(t *testing.T) {
 		}
 	}
 }
+
+func TestLoggingHandler_PingUserAgent(t *testing.T) {
+	tests := []struct {
+		ExpectedLogMessage string
+		Path               string
+		SilencePingLogging bool
+		WithUserAgent      string
+	}{
+		{"444\n", "/foo", true, "Blah"},
+		{"444\n", "/foo", false, "Blah"},
+		{"", "/ping", true, "Blah"},
+		{"200\n", "/ping", false, "Blah"},
+		{"", "/ping", true, "PingMe!"},
+		{"", "/ping", false, "PingMe!"},
+		{"", "/foo", true, "PingMe!"},
+		{"", "/foo", false, "PingMe!"},
+	}
+
+	for idx, test := range tests {
+		t.Run(fmt.Sprintf("%d", idx), func(t *testing.T) {
+			opts := options.NewOptions()
+			opts.PingUserAgent = "PingMe!"
+			opts.SkipAuthRegex = []string{"/foo"}
+			opts.Upstreams = []string{"static://444/foo"}
+			opts.Logging.SilencePing = test.SilencePingLogging
+			if test.SilencePingLogging {
+				opts.Logging.ExcludePaths = []string{"/ping"}
+			}
+			opts.RawRedirectURL = "localhost"
+			validation.Validate(opts)
+
+			p := NewOAuthProxy(opts, func(email string) bool {
+				return true
+			})
+			p.provider = NewTestProvider(&url.URL{Host: "localhost"}, "")
+
+			buf := bytes.NewBuffer(nil)
+			logger.SetOutput(buf)
+			logger.SetReqEnabled(true)
+			logger.SetReqTemplate("{{.StatusCode}}")
+
+			r, _ := http.NewRequest("GET", test.Path, nil)
+			if test.WithUserAgent != "" {
+				r.Header.Set("User-Agent", test.WithUserAgent)
+			}
+
+			h := LoggingHandler(p)
+			h.ServeHTTP(httptest.NewRecorder(), r)
+
+			actual := buf.String()
+			if !strings.Contains(actual, test.ExpectedLogMessage) {
+				t.Errorf("Log message was\n%s\ninstead of matching \n%s", actual, test.ExpectedLogMessage)
+			}
+		})
+	}
+}

oauthproxy.go

@@ -82,6 +82,8 @@ type OAuthProxy struct {
 
 	RobotsPath        string
 	PingPath          string
+	PingUserAgent     string
+	SilencePings      bool
 	SignInPath        string
 	SignOutPath       string
 	OAuthStartPath    string
@@ -312,6 +314,8 @@ func NewOAuthProxy(opts *options.Options, validator func(string) bool) *OAuthPro
 
 		RobotsPath:        "/robots.txt",
 		PingPath:          opts.PingPath,
+		PingUserAgent:     opts.PingUserAgent,
+		SilencePings:      opts.Logging.SilencePing,
 		SignInPath:        fmt.Sprintf("%s/sign_in", opts.ProxyPrefix),
 		SignOutPath:       fmt.Sprintf("%s/sign_out", opts.ProxyPrefix),
 		OAuthStartPath:    fmt.Sprintf("%s/start", opts.ProxyPrefix),
@@ -466,6 +470,11 @@ func (p *OAuthProxy) RobotsTxt(rw http.ResponseWriter) {
 
 // PingPage responds 200 OK to requests
 func (p *OAuthProxy) PingPage(rw http.ResponseWriter) {
+	if p.SilencePings {
+		if rl, ok := rw.(*responseLogger); ok {
+			rl.silent = true
+		}
+	}
 	rw.WriteHeader(http.StatusOK)
 	fmt.Fprintf(rw, "OK")
 }
@@ -675,6 +684,17 @@ func prepareNoCache(w http.ResponseWriter) {
 	}
 }
 
+// IsPingRequest will check if the request appears to be performing a health check
+// either via the path it's requesting or by a special User-Agent configuration.
+func (p *OAuthProxy) IsPingRequest(req *http.Request) bool {
+
+	if req.URL.EscapedPath() == p.PingPath {
+		return true
+	}
+
+	return p.PingUserAgent != "" && req.Header.Get("User-Agent") == p.PingUserAgent
+}
+
 func (p *OAuthProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	if strings.HasPrefix(req.URL.Path, p.ProxyPrefix) {
 		prepareNoCache(rw)
@@ -683,7 +703,7 @@ func (p *OAuthProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	switch path := req.URL.Path; {
 	case path == p.RobotsPath:
 		p.RobotsTxt(rw)
-	case path == p.PingPath:
+	case p.IsPingRequest(req):
 		p.PingPage(rw)
 	case p.IsWhitelistedRequest(req):
 		p.serveMux.ServeHTTP(rw, req)

pkg/apis/options/options.go

@@ -24,6 +24,7 @@ type SignatureData struct {
 type Options struct {
 	ProxyPrefix        string `flag:"proxy-prefix" cfg:"proxy_prefix"`
 	PingPath           string `flag:"ping-path" cfg:"ping_path"`
+	PingUserAgent      string `flag:"ping-user-agent" cfg:"ping_user_agent"`
 	ProxyWebSockets    bool   `flag:"proxy-websockets" cfg:"proxy_websockets"`
 	HTTPAddress        string `flag:"http-address" cfg:"http_address"`
 	HTTPSAddress       string `flag:"https-address" cfg:"https_address"`
@@ -245,6 +246,7 @@ func NewFlagSet() *pflag.FlagSet {
 	flagSet.String("footer", "", "custom footer string. Use \"-\" to disable default footer.")
 	flagSet.String("proxy-prefix", "/oauth2", "the url root path that this proxy should be nested under (e.g. /<oauth2>/sign_in)")
 	flagSet.String("ping-path", "/ping", "the ping endpoint that can be used for basic health checks")
+	flagSet.String("ping-user-agent", "", "special User-Agent that will be used for basic health checks")
 	flagSet.Bool("proxy-websockets", true, "enables WebSocket proxying")
 
 	flagSet.String("cookie-name", "_oauth2_proxy", "the name of the cookie that the oauth_proxy creates")