From 2c5f4e08b52120e2b6e60c9bdc5a1bc53f319fac Mon Sep 17 00:00:00 2001
From: Martin Polden <martin.polden@gmail.com>
Date: Wed, 27 Apr 2016 17:07:53 +0200
Subject: [PATCH] Refactor

---
 api/api.go      | 41 ++++++++++++++++++++++++-----------------
 api/api_test.go |  3 +++
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/api/api.go b/api/api.go
index 98f12b6..b2fe8dd 100644
--- a/api/api.go
+++ b/api/api.go
@@ -35,7 +35,7 @@ type Response struct {
 	Hostname string `json:"hostname,omitempty"`
 }
 
-type TestPortResponse struct {
+type PortResponse struct {
 	IP        net.IP `json:"ip"`
 	Port      uint64 `json:"port"`
 	Reachable bool   `json:"reachable"`
@@ -86,6 +86,27 @@ func (a *API) newResponse(r *http.Request) (Response, error) {
 	}, nil
 }
 
+func (a *API) newPortResponse(r *http.Request) (PortResponse, error) {
+	vars := mux.Vars(r)
+	port, err := strconv.ParseUint(vars["port"], 10, 16)
+	if err != nil {
+		return PortResponse{Port: port}, err
+	}
+	if port < 1 || port > 65355 {
+		return PortResponse{Port: port}, fmt.Errorf("invalid port: %d", port)
+	}
+	ip, err := ipFromRequest(a.IPHeader, r)
+	if err != nil {
+		return PortResponse{Port: port}, err
+	}
+	err = a.oracle.LookupPort(ip, port)
+	return PortResponse{
+		IP:        ip,
+		Port:      port,
+		Reachable: err == nil,
+	}, nil
+}
+
 func (a *API) CLIHandler(w http.ResponseWriter, r *http.Request) *appError {
 	ip, err := ipFromRequest(a.IPHeader, r)
 	if err != nil {
@@ -128,23 +149,9 @@ func (a *API) JSONHandler(w http.ResponseWriter, r *http.Request) *appError {
 }
 
 func (a *API) PortHandler(w http.ResponseWriter, r *http.Request) *appError {
-	vars := mux.Vars(r)
-	port, err := strconv.ParseUint(vars["port"], 10, 16)
+	response, err := a.newPortResponse(r)
 	if err != nil {
-		return badRequest(err).WithMessage("Invalid port: " + vars["port"]).AsJSON()
-	}
-	if port < 1 || port > 65355 {
-		return badRequest(nil).WithMessage("Invalid port: " + vars["port"]).AsJSON()
-	}
-	ip, err := ipFromRequest(a.IPHeader, r)
-	if err != nil {
-		return internalServerError(err).AsJSON()
-	}
-	err = a.oracle.LookupPort(ip, port)
-	response := TestPortResponse{
-		IP:        ip,
-		Port:      port,
-		Reachable: err == nil,
+		return badRequest(err).WithMessage(fmt.Sprintf("Invalid port: %d", response.Port)).AsJSON()
 	}
 	b, err := json.Marshal(response)
 	if err != nil {
diff --git a/api/api_test.go b/api/api_test.go
index c34c5ae..efe8c66 100644
--- a/api/api_test.go
+++ b/api/api_test.go
@@ -85,6 +85,9 @@ func TestJSONHandlers(t *testing.T) {
 		status int
 	}{
 		{s.URL, `{"ip":"127.0.0.1","country":"Elbonia","city":"Bornyasherk","hostname":"localhost"}`, 200},
+		{s.URL + "/port/foo", `{"error":"404 page not found"}`, 404},
+		{s.URL + "/port/0", `{"error":"Invalid port: 0"}`, 400},
+		{s.URL + "/port/65356", `{"error":"Invalid port: 65356"}`, 400},
 		{s.URL + "/port/31337", `{"ip":"127.0.0.1","port":31337,"reachable":true}`, 200},
 		{s.URL + "/foo", `{"error":"404 page not found"}`, 404},
 	}