refactor: code review.

app.go

@@ -6,22 +6,18 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"strconv"
-	"strings"
-	"sync"
-
-	"github.com/gorilla/websocket"
 	"log"
 	"net"
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
+	"strings"
+	"sync"
 	"time"
-)
 
-var cert string
+	"github.com/gorilla/websocket"
-var key string
+)
-var port string
 
 const (
 	_        = iota
@@ -31,6 +27,10 @@ const (
 	TB
 )
 
+var cert string
+var key string
+var port string
+
 func init() {
 	flag.StringVar(&cert, "cert", "", "give me a certificate")
 	flag.StringVar(&key, "key", "", "give me a key")
@@ -44,19 +44,49 @@ var upgrader = websocket.Upgrader{
 
 func main() {
 	flag.Parse()
+
 	http.HandleFunc("/data", dataHandler)
 	http.HandleFunc("/echo", echoHandler)
 	http.HandleFunc("/bench", benchHandler)
-	http.HandleFunc("/", whoami)
+	http.HandleFunc("/", whoamiHandler)
-	http.HandleFunc("/api", api)
+	http.HandleFunc("/api", apiHandler)
 	http.HandleFunc("/health", healthHandler)
+
 	fmt.Println("Starting up on port " + port)
+
 	if len(cert) > 0 && len(key) > 0 {
 		log.Fatal(http.ListenAndServeTLS(":"+port, cert, key, nil))
 	}
 	log.Fatal(http.ListenAndServe(":"+port, nil))
 }
 
+func benchHandler(w http.ResponseWriter, _ *http.Request) {
+	w.Header().Set("Connection", "keep-alive")
+	w.Header().Set("Content-Type", "text/plain")
+	fmt.Fprint(w, "1")
+}
+
+func echoHandler(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+
+	for {
+		messageType, p, err := conn.ReadMessage()
+		if err != nil {
+			return
+		}
+
+		printBinary(p)
+		err = conn.WriteMessage(messageType, p)
+		if err != nil {
+			return
+		}
+	}
+}
+
 func printBinary(s []byte) {
 	fmt.Printf("Received b:")
 	for n := 0; n < len(s); n++ {
@@ -64,29 +94,7 @@ func printBinary(s []byte) {
 	}
 	fmt.Printf("\n")
 }
-func benchHandler(w http.ResponseWriter, r *http.Request) {
+
-	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("Content-Type", "text/plain")
-	fmt.Fprint(w, "1")
-}
-func echoHandler(w http.ResponseWriter, r *http.Request) {
-	conn, err := upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	for {
-		messageType, p, err := conn.ReadMessage()
-		if err != nil {
-			return
-		}
-		printBinary(p)
-		err = conn.WriteMessage(messageType, p)
-		if err != nil {
-			return
-		}
-	}
-}
 func dataHandler(w http.ResponseWriter, r *http.Request) {
 	u, _ := url.Parse(r.URL.String())
 	queryParams := u.Query()
@@ -100,11 +108,15 @@ func dataHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	unit := queryParams.Get("unit")
-	switch (strings.ToLower(unit)) {
+	switch strings.ToLower(unit) {
-		case "kb": size = size * KB
+	case "kb":
-		case "mb": size = size * MB
+		size = size * KB
-		case "gb": size = size * GB
+	case "mb":
-		case "tb": size = size * TB
+		size = size * MB
+	case "gb":
+		size = size * GB
+	case "tb":
+		size = size * TB
 	}
 
 	attachment, err := strconv.ParseBool(queryParams.Get("attachment"))
@@ -115,25 +127,30 @@ func dataHandler(w http.ResponseWriter, r *http.Request) {
 	content := fillContent(size)
 
 	if attachment {
-		const name = "data.txt"
 		w.Header().Add("Content-Disposition", "Attachment")
-		http.ServeContent(w, r, name, time.Now(), content)
+		http.ServeContent(w, r, "data.txt", time.Now(), content)
-	} else {
+		return
-		io.Copy(w, content)
+	}
+
+	if _, err := io.Copy(w, content); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
 	}
 }
-func whoami(w http.ResponseWriter, req *http.Request) {
+
+func whoamiHandler(w http.ResponseWriter, req *http.Request) {
 	u, _ := url.Parse(req.URL.String())
-	queryParams := u.Query()
+	wait := u.Query().Get("wait")
-	wait := queryParams.Get("wait")
 	if len(wait) > 0 {
 		duration, err := time.ParseDuration(wait)
 		if err == nil {
 			time.Sleep(duration)
 		}
 	}
+
 	hostname, _ := os.Hostname()
 	fmt.Fprintln(w, "Hostname:", hostname)
+
 	ifaces, _ := net.Interfaces()
 	for _, i := range ifaces {
 		addrs, _ := i.Addrs()
@@ -149,11 +166,16 @@ func whoami(w http.ResponseWriter, req *http.Request) {
 			fmt.Fprintln(w, "IP:", ip)
 		}
 	}
-	req.Write(w)
+
+	if err := req.Write(w); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
 }
 
-func api(w http.ResponseWriter, req *http.Request) {
+func apiHandler(w http.ResponseWriter, req *http.Request) {
 	hostname, _ := os.Hostname()
+
 	data := struct {
 		Hostname string      `json:"hostname,omitempty"`
 		IP       []string    `json:"ip,omitempty"`
@@ -162,12 +184,12 @@ func api(w http.ResponseWriter, req *http.Request) {
 		Host     string      `json:"host,omitempty"`
 		Method   string      `json:"method,omitempty"`
 	}{
-		hostname,
+		Hostname: hostname,
-		[]string{},
+		IP:       []string{},
-		req.Header,
+		Headers:  req.Header,
-		req.URL.RequestURI(),
+		URL:      req.URL.RequestURI(),
-		req.Host,
+		Host:     req.Host,
-		req.Method,
+		Method:   req.Method,
 	}
 
 	ifaces, _ := net.Interfaces()
@@ -185,29 +207,34 @@ func api(w http.ResponseWriter, req *http.Request) {
 			data.IP = append(data.IP, ip.String())
 		}
 	}
-	json.NewEncoder(w).Encode(data)
+
+	if err := json.NewEncoder(w).Encode(data); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
 }
 
 type healthState struct {
 	StatusCode int
 }
 
-var currentHealthState = healthState{200}
+var currentHealthState = healthState{http.StatusOK}
 var mutexHealthState = &sync.RWMutex{}
 
 func healthHandler(w http.ResponseWriter, req *http.Request) {
 	if req.Method == http.MethodPost {
 		var statusCode int
-		err := json.NewDecoder(req.Body).Decode(&statusCode)
+
-		if err != nil {
+		if err := json.NewDecoder(req.Body).Decode(&statusCode); err != nil {
-			w.WriteHeader(http.StatusBadRequest)
+			http.Error(w, err.Error(), http.StatusBadRequest)
-			w.Write([]byte(err.Error()))
+			return
-		} else {
+		}
+
 		fmt.Printf("Update health check status code [%d]\n", statusCode)
+
 		mutexHealthState.Lock()
 		defer mutexHealthState.Unlock()
 		currentHealthState.StatusCode = statusCode
-		}
 	} else {
 		mutexHealthState.RLock()
 		defer mutexHealthState.RUnlock()