From 783221bb99d6a2a2ef26e91dd1df696d0920415b Mon Sep 17 00:00:00 2001
From: Yi Yang <yiy@google.com>
Date: Thu, 20 Oct 2022 17:44:33 +0000
Subject: [PATCH] Add statusCodeHandler to 404-server-with-metrics which serves
 status code given in URL path

---
 .../server-with-metrics.go                    | 36 +++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/cmd/404-server-with-metrics/server-with-metrics.go b/cmd/404-server-with-metrics/server-with-metrics.go
index 617e2b0a11..f2f94507de 100644
--- a/cmd/404-server-with-metrics/server-with-metrics.go
+++ b/cmd/404-server-with-metrics/server-with-metrics.go
@@ -26,6 +26,8 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
+	"strconv"
+	"strings"
 	"time"
 
 	"k8s.io/klog/v2"
@@ -47,6 +49,8 @@ var (
 	isProd            = flag.Bool("is_prod", true, "Indicates if the server is running in production.")
 )
 
+const statusCodePrefix string = "/statuscode/"
+
 func main() {
 	flag.Parse()
 	klog.InitFlags(nil)
@@ -61,7 +65,7 @@ func main() {
 	server.registerHandlers()
 	klog.Infof("Default 404 server is running with GOMAXPROCS(%d) on %s:%d\n", runtime.GOMAXPROCS(-1), hostName, *port)
 
-	// The main http server for handling NotFound and healthzrequests
+	// The main http server for handling NotFound, StatusCode and healthz requests
 	go func() {
 		err := server.httpServer.ListenAndServe()
 		if err != nil {
@@ -186,12 +190,13 @@ func (s *server) registerHandlers() {
 	httpMux := http.NewServeMux()
 	metricsMux := http.NewServeMux()
 
-	// Register the default notFoundHandler, healthz with the main http server
+	// Register the default notFoundHandler, statusCodeHandler, healthz with the main http server
 	httpMux.HandleFunc("/", s.notFoundHandler())
 	// enable shutdown handler only for non-prod environments
 	if *isProd == false {
 		httpMux.HandleFunc("/shutdown", s.shutdownHandler())
 	}
+	httpMux.HandleFunc(statusCodePrefix, s.statusCodeHandler())
 	httpMux.HandleFunc("/healthz", s.healthzHandler())
 
 	// Register the healthz and metrics handlers with the metrics server
@@ -245,6 +250,33 @@ func (s *server) notFoundHandler() http.HandlerFunc {
 	}
 }
 
+// statusCodeHandler returns a status code given by URL path in the request.
+// If URL path does not contain a valid status code, returns 404 (NotFound).
+func (s *server) statusCodeHandler() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		p := strings.TrimPrefix(r.URL.Path, statusCodePrefix)
+		code, err := strconv.Atoi(p)
+		if err != nil {
+			w.WriteHeader(http.StatusNotFound)
+			fmt.Fprintf(w, "response 404, service rules for [ %s ] non-existent, [ %s ] cannot be converted into status code \n", r.URL.Path, p)
+			return
+		}
+		if code < 400 || code > 599 {
+			w.WriteHeader(http.StatusNotFound)
+			fmt.Fprintf(w, "response 404, service rules for [ %s ] non-existent, [ %s ] is unsupported, status code must be in range [400, 599] \n", r.URL.Path, p)
+			return
+		}
+		statusText := http.StatusText(code)
+		if statusText == "" {
+			w.WriteHeader(http.StatusNotFound)
+			fmt.Fprintf(w, "response 404, service rules for [ %s ] non-existent, [ %s ] is an unknown status code \n", r.URL.Path, p)
+			return
+		}
+		w.WriteHeader(code)
+		fmt.Fprintf(w, "response %d \n", code)
+	}
+}
+
 // graceful shutdown handler
 func gracefulShutdown(s *server) {
 	// have a small buffered channel so as not to lose signal sent when we were not ready.