From 7a99835dab64f7864186185761bbf5194216f8b6 Mon Sep 17 00:00:00 2001
From: Matthew Holt <mholt@users.noreply.github.com>
Date: Thu, 4 Jun 2020 12:06:38 -0600
Subject: reverseproxy: Enable changing only the status code (close #2920)

---
 modules/caddyhttp/caddyhttp.go                 | 39 ++++++++++++++++++++++++++
 modules/caddyhttp/reverseproxy/reverseproxy.go | 24 ++++++++++++++--
 modules/caddyhttp/subroute.go                  | 30 --------------------
 3 files changed, 60 insertions(+), 33 deletions(-)

diff --git a/modules/caddyhttp/caddyhttp.go b/modules/caddyhttp/caddyhttp.go
index fda7a92..a7ac889 100644
--- a/modules/caddyhttp/caddyhttp.go
+++ b/modules/caddyhttp/caddyhttp.go
@@ -92,6 +92,45 @@ var errorEmptyHandler Handler = HandlerFunc(func(w http.ResponseWriter, r *http.
 	return nil
 })
 
+// ResponseHandler pairs a response matcher with custom handling
+// logic. Either the status code can be changed to something else
+// while using the original response body, or, if a status code
+// is not set, it can execute a custom route list; this is useful
+// for executing handler routes based on the properties of an HTTP
+// response that has not been written out to the client yet.
+//
+// To use this type, provision it at module load time, then when
+// ready to use, match the response against its matcher; if it
+// matches (or doesn't have a matcher), change the status code on
+// the response if configured; otherwise invoke the routes by
+// calling `rh.Routes.Compile(next).ServeHTTP(rw, req)` (or similar).
+type ResponseHandler struct {
+	// The response matcher for this handler. If empty/nil,
+	// it always matches.
+	Match *ResponseMatcher `json:"match,omitempty"`
+
+	// To write the original response body but with a different
+	// status code, set this field to the desired status code.
+	// If set, this takes priority over routes.
+	StatusCode WeakString `json:"status_code,omitempty"`
+
+	// The list of HTTP routes to execute if no status code is
+	// specified. If evaluated, the original response body
+	// will not be written.
+	Routes RouteList `json:"routes,omitempty"`
+}
+
+// Provision sets up the routse in rh.
+func (rh *ResponseHandler) Provision(ctx caddy.Context) error {
+	if rh.Routes != nil {
+		err := rh.Routes.Provision(ctx)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // WeakString is a type that unmarshals any JSON value
 // as a string literal, with the following exceptions:
 //
diff --git a/modules/caddyhttp/reverseproxy/reverseproxy.go b/modules/caddyhttp/reverseproxy/reverseproxy.go
index 06802a0..7971348 100644
--- a/modules/caddyhttp/reverseproxy/reverseproxy.go
+++ b/modules/caddyhttp/reverseproxy/reverseproxy.go
@@ -24,6 +24,7 @@ import (
 	"net"
 	"net/http"
 	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -531,15 +532,32 @@ func (h *Handler) reverseProxy(rw http.ResponseWriter, req *http.Request, di Dia
 		}
 	}
 
+	// see if any response handler is configured for this response from the backend
 	for i, rh := range h.HandleResponse {
-		if len(rh.Routes) == 0 {
+		if rh.Match != nil && !rh.Match.Match(res.StatusCode, res.Header) {
 			continue
 		}
-		if rh.Match != nil && !rh.Match.Match(res.StatusCode, res.Header) {
+
+		repl := req.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
+
+		// if configured to only change the status code, do that then continue regular proxy response
+		if statusCodeStr := rh.StatusCode.String(); statusCodeStr != "" {
+			statusCode, err := strconv.Atoi(repl.ReplaceAll(statusCodeStr, ""))
+			if err != nil {
+				return caddyhttp.Error(http.StatusInternalServerError, err)
+			}
+			if statusCode != 0 {
+				res.StatusCode = statusCode
+			}
+			break
+		}
+
+		// otherwise, if there are any routes configured, execute those as the
+		// actual response instead of what we got from the proxy backend
+		if len(rh.Routes) == 0 {
 			continue
 		}
 		res.Body.Close()
-		repl := req.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
 		repl.Set("http.reverse_proxy.status_code", res.StatusCode)
 		repl.Set("http.reverse_proxy.status_text", res.Status)
 		h.logger.Debug("handling response", zap.Int("handler", i))
diff --git a/modules/caddyhttp/subroute.go b/modules/caddyhttp/subroute.go
index b1700f5..2e80d88 100644
--- a/modules/caddyhttp/subroute.go
+++ b/modules/caddyhttp/subroute.go
@@ -80,36 +80,6 @@ func (sr *Subroute) ServeHTTP(w http.ResponseWriter, r *http.Request, next Handl
 	return err
 }
 
-// ResponseHandler pairs a response matcher with a route list.
-// It is useful for executing handler routes based on the
-// properties of an HTTP response that has not been written
-// out to the client yet.
-//
-// To use this type, provision it at module load time, then
-// when ready to use, match the response against its matcher;
-// if it matches (or doesn't have a matcher), invoke the routes
-// by calling `rh.Routes.Compile(next).ServeHTTP(rw, req)` (or
-// similar).
-type ResponseHandler struct {
-	// The response matcher for this handler. If empty/nil,
-	// it always matches.
-	Match *ResponseMatcher `json:"match,omitempty"`
-
-	// The list of HTTP routes to execute.
-	Routes RouteList `json:"routes,omitempty"`
-}
-
-// Provision sets up the routse in rh.
-func (rh *ResponseHandler) Provision(ctx caddy.Context) error {
-	if rh.Routes != nil {
-		err := rh.Routes.Provision(ctx)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // Interface guards
 var (
 	_ caddy.Provisioner = (*Subroute)(nil)
-- 
cgit v1.2.3