From 65da260ddb5ffa38cfb07a4abacde1ec649db2a3 Mon Sep 17 00:00:00 2001
From: Bel LaPointe <breel@qualtrics.com>
Date: Wed, 14 Oct 2020 09:16:23 -0600
Subject: [PATCH] block on http shouldnt change echo ness

---
 main.go | 64 +++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 49 insertions(+), 15 deletions(-)

diff --git a/main.go b/main.go
index f77848d..d1a29ce 100755
--- a/main.go
+++ b/main.go
@@ -29,17 +29,21 @@ func main() {
 	crt := flag.String("crt", envOrDefault("CRT", ""), "path to crt")
 	key := flag.String("key", envOrDefault("KEY", ""), "path to key")
 	tcp := flag.Bool("tcp", false, "accept TCP")
+	udp := flag.Bool("udp", false, "accept udp")
 	fail := flag.Bool("fail", false, "fail connections forever")
 	block := flag.Int("block", 0, "seconds to block connections")
 	status := flag.Int("status", http.StatusOK, "status to always return")
+	body := flag.String("body", "", "body to always return")
 	jwt := flag.String("jwt", envOrDefault("SECRET", ""), "secret for JWT")
 	flag.Parse()
 
 	// start server
-	if *tcp {
+	if *udp {
+		startUDP(*port, *block)
+	} else if *tcp {
 		startTCP(*port, *fail, *block)
 	} else {
-		startHTTP(*port, *fail, *block, *jwt, *crt, *key, *status)
+		startHTTP(*port, *fail, *block, *jwt, *crt, *key, *status, *body)
 	}
 
 	// catch stop
@@ -48,6 +52,29 @@ func main() {
 	<-stop
 }
 
+func startUDP(port string, block int) {
+	pc, err := net.ListenPacket("udp", ":"+port)
+	if err != nil {
+		panic(err)
+	}
+	log.Println("echoing", port)
+	go func() {
+		defer pc.Close()
+		buff := make([]byte, 1024)
+		for {
+			n, addr, err := pc.ReadFrom(buff)
+			if err != nil {
+				panic(err)
+			}
+			log.Printf("%s: %s", addr, buff[:n])
+			_, err = pc.WriteTo(buff[:n], addr)
+			if err != nil {
+				panic(err)
+			}
+		}
+	}()
+}
+
 func startTCP(port string, fail bool, block int) {
 	l, err := net.Listen("tcp", ":"+port)
 	if err != nil {
@@ -72,10 +99,10 @@ func startTCP(port string, fail bool, block int) {
 	}()
 }
 
-func startHTTP(port string, fail bool, block int, secret, crt, key string, status int) {
+func startHTTP(port string, fail bool, block int, secret, crt, key string, status int, body string) {
 	srv := &http.Server{
 		Addr:    ":" + port,
-		Handler: makeHTTPHandler(fail, block, secret, status),
+		Handler: makeHTTPHandler(fail, block, secret, status, body),
 	}
 	go func() {
 		if crt == "" || key == "" {
@@ -116,20 +143,20 @@ func startHTTP(port string, fail bool, block int, secret, crt, key string, statu
 	log.Println("Listening on", srv.Addr)
 }
 
-func makeHTTPHandler(fail bool, block int, secret string, status int) http.HandlerFunc {
+func makeHTTPHandler(fail bool, block int, secret string, status int, body string) http.HandlerFunc {
 	if block > 0 {
-		return blockH(block, fail)
+		return blockH(block, fail, secret, status, body)
 	} else if fail {
 		return http.HandlerFunc(failH)
 	}
-	return echoH(secret, status)
+	return echoH(secret, status, body)
 }
 
 func failH(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusInternalServerError)
 }
 
-func blockH(block int, fail bool) http.HandlerFunc {
+func blockH(block int, fail bool, secret string, status int, body string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		select {
 		case <-time.After(time.Second * time.Duration(block)):
@@ -137,14 +164,17 @@ func blockH(block int, fail bool) http.HandlerFunc {
 			log.Println(r.Context().Err())
 			fmt.Fprintln(w, r.Context().Err())
 			w.WriteHeader(http.StatusInternalServerError)
+			return
 		}
 		if fail {
 			w.WriteHeader(http.StatusInternalServerError)
+		} else {
+			echoH(secret, status, body)(w, r)
 		}
 	}
 }
 
-func echoH(secret string, status int) http.HandlerFunc {
+func echoH(secret string, status int, body string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		log.Println("METHOD:", r.Method)
 		log.Println("URL:", r.URL.String())
@@ -167,11 +197,15 @@ func echoH(secret string, status int) http.HandlerFunc {
 		log.Println("BODY:", s)
 
 		w.WriteHeader(status)
-		fmt.Fprintf(w, "===\nMETHOD: %q\nURL: %q\nHEADER: %q\nBODY: %q\n===\n",
-			r.Method,
-			r.URL.String(),
-			fmt.Sprintf("%v", r.Header),
-			s,
-		)
+		if len(body) == 0 {
+			fmt.Fprintf(w, "===\nMETHOD: %q\nURL: %q\nHEADER: %q\nBODY: %q\n===\n",
+				r.Method,
+				r.URL.String(),
+				fmt.Sprintf("%v", r.Header),
+				s,
+			)
+		} else {
+			fmt.Fprintln(w, body)
+		}
 	}
 }