From 7a4a202a53ffcd5285c94733cf568f7e271e652c Mon Sep 17 00:00:00 2001
From: Renuka Fernando <renukapiyumal@gmail.com>
Date: Thu, 27 Jun 2024 22:34:01 +0530
Subject: [PATCH] refactor: Add command-line option to set response body via
 flag

---
 README.md |  4 +++-
 main.go   | 23 ++++++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 6455484..b2ca749 100644
--- a/README.md
+++ b/README.md
@@ -24,8 +24,10 @@ docker run --rm -p 8080:8080 -e "NAME=Service A" renukafernando/request-info:lat
 - set `-cert` argument to set the path to the cert file
 - set `-mtls` argument to enable mTLS
 - set `-ca` argument to set the path to the CA file
+- set `-setResponseBody` argument to set the response body via the flag `-responseBody`
+- set `-responseBody` argument to set the response body
 - set `-disable-access-logs` argument to disable access logs
-- ser `-wait-before-graceful-shutdown-ms` argument to set the time to wait before graceful shutdown in milliseconds
+- set `-wait-before-graceful-shutdown-ms` argument to set the time to wait before graceful shutdown in milliseconds
 
 Resource Usage
 ```yaml
diff --git a/main.go b/main.go
index 0618358..b64c576 100644
--- a/main.go
+++ b/main.go
@@ -207,6 +207,10 @@ func setResponseHandler(handler func(http.ResponseWriter, *http.Request)) func(h
 	}
 }
 
+func setResponseBodyFromArgsHandler(w http.ResponseWriter, req *http.Request) {
+	_, _ = fmt.Fprintln(w, responseBody)
+}
+
 func healthzHandler(w http.ResponseWriter, req *http.Request) {
 	if atomic.LoadInt32(&healthy) == 1 {
 		w.WriteHeader(http.StatusOK)
@@ -257,6 +261,8 @@ func main() {
 	flag.StringVar(&clientCA, "ca", "ca.crt", "CA certificate file for client verification")
 	flag.IntVar(&delayMs, "delayMs", 0, "Time to wait (ms) before responding to request")
 	flag.IntVar(&statusCode, "status", 200, "HTTP status code to respond")
+	flag.BoolVar(&setResponseBody, "setResponseBody", false, "Response the body set via -responseBody flag")
+	flag.StringVar(&responseBody, "responseBody", "", "The response body to return")
 	flag.BoolVar(&disableAccessLogs, "disable-access-logs", false, "Disable access logs")
 	flag.IntVar(&waitBeforeGracefulShutdownMs, "wait-before-graceful-shutdown-ms", 0, "Time to wait (ms) before graceful shutdown")
 	flag.Parse()
@@ -276,12 +282,17 @@ func main() {
 	// Create a new HTTP server
 	server := http.NewServeMux()
 
-	// Set up a handler for the desired route
-	server.HandleFunc("/req-info/response", reqInfoSetPayloadHandler)
-	server.HandleFunc("/empty", setResponseHandler(empty))
-	server.HandleFunc("/echo", setResponseHandler(echo))
 	server.HandleFunc("/healthz", healthzHandler)
-	server.HandleFunc("/", setResponseHandler(reqInfo))
+
+	if setResponseBody {
+		log.Println("[INFO] Responding the body set via -responseBody flag")
+		server.HandleFunc("/", setResponseBodyFromArgsHandler)
+	} else {
+		server.HandleFunc("/req-info/response", reqInfoSetPayloadHandler)
+		server.HandleFunc("/empty", setResponseHandler(empty))
+		server.HandleFunc("/echo", setResponseHandler(echo))
+		server.HandleFunc("/", setResponseHandler(reqInfo))
+	}
 
 	// Enable CORS middleware
 	corsHandler := enableCORS(server)
@@ -373,5 +384,7 @@ var key string
 var clientCA string
 var delayMs int
 var statusCode int
+var setResponseBody bool
+var responseBody string
 var disableAccessLogs bool
 var waitBeforeGracefulShutdownMs int