From 1af2589c73da35a11c56ef1fea0e56a3366e32e6 Mon Sep 17 00:00:00 2001
From: Wilton Rodrigues <wiltonsr94@gmail.com>
Date: Thu, 28 Apr 2022 11:20:25 -0300
Subject: [PATCH] Make /data path configurable

- Permit disable /data path to prevent DoS attack
- Fixes #17
---
 README.md | 1 +
 app.go    | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/README.md b/README.md
index 72d714d..b123c31 100644
--- a/README.md
+++ b/README.md
@@ -55,6 +55,7 @@ Heath check.
 | `port`    | `WHOAMI_PORT_NUMBER` | Give me a port number. (default: `80`)  |
 | `name`    | `WHOAMI_NAME`        | Give me a name.                         |
 | `verbose` |                      | Enable verbose logging.                 |
+| `no-data` |                      | Disable data handler.                   |
 
 ## Examples
 
diff --git a/app.go b/app.go
index 49ce421..eb2587e 100644
--- a/app.go
+++ b/app.go
@@ -43,10 +43,12 @@ var (
 	port    string
 	name    string
 	verbose bool
+	no_data bool
 )
 
 func init() {
 	flag.BoolVar(&verbose, "verbose", false, "Enable verbose logging")
+	flag.BoolVar(&no_data, "no-data", false, "Disable data handler")
 	flag.StringVar(&cert, "cert", "", "give me a certificate")
 	flag.StringVar(&key, "key", "", "give me a key")
 	flag.StringVar(&ca, "cacert", "", "give me a CA chain, enforces mutual TLS")
@@ -166,6 +168,11 @@ func printBinary(s []byte) {
 }
 
 func dataHandler(w http.ResponseWriter, r *http.Request) {
+	if no_data {
+		_, _ = fmt.Fprintln(w, "dataHandler disabled!")
+		return
+	}
+
 	queryParams := r.URL.Query()
 
 	size, err := strconv.ParseInt(queryParams.Get("size"), 10, 64)