Support getting information about another IP (#94)

Signed-off-by: Vincent Batts <vbatts@hashbangbash.com>

4 files changed, 45 insertions, 22 deletions

diff --git a/README.md b/README.md
index 9fb2b54..e38fd7e 100644
--- a/README.md
+++ b/README.md
@@ -81,6 +81,7 @@ between IPv4 and IPv6 lookup.
 * JSON output
 * ASN, country and city lookup using the MaxMind GeoIP database
 * Port testing
+* All endpoints (except `/port`) can return information about a custom IP address specified via `?ip=` query parameter
 * Open source under the [BSD 3-Clause license](https://opensource.org/licenses/BSD-3-Clause)
 
 ## Why?
diff --git a/http/http.go b/http/http.go
index 7c9263a..b5d275d 100644
--- a/http/http.go
+++ b/http/http.go
@@ -69,15 +69,27 @@ func ipFromForwardedForHeader(v string) string {
 	return v[:sep]
 }
 
-func ipFromRequest(headers []string, r *http.Request) (net.IP, error) {
+// ipFromRequest detects the IP address for this transaction.
+//
+// * `headers` - the specific HTTP headers to trust
+// * `r` - the incoming HTTP request
+// * `customIP` - whether to allow the IP to be pulled from query parameters
+func ipFromRequest(headers []string, r *http.Request, customIP bool) (net.IP, error) {
 	remoteIP := ""
-	for _, header := range headers {
-		remoteIP = r.Header.Get(header)
-		if http.CanonicalHeaderKey(header) == "X-Forwarded-For" {
-			remoteIP = ipFromForwardedForHeader(remoteIP)
+	if customIP && r.URL != nil {
+		if v, ok := r.URL.Query()["ip"]; ok {
+			remoteIP = v[0]
 		}
-		if remoteIP != "" {
-			break
+	}
+	if remoteIP == "" {
+		for _, header := range headers {
+			remoteIP = r.Header.Get(header)
+			if http.CanonicalHeaderKey(header) == "X-Forwarded-For" {
+				remoteIP = ipFromForwardedForHeader(remoteIP)
+			}
+			if remoteIP != "" {
+				break
+			}
 		}
 	}
 	if remoteIP == "" {
@@ -105,7 +117,7 @@ func userAgentFromRequest(r *http.Request) *useragent.UserAgent {
 }
 
 func (s *Server) newResponse(r *http.Request) (Response, error) {
-	ip, err := ipFromRequest(s.IPHeaders, r)
+	ip, err := ipFromRequest(s.IPHeaders, r, true)
 	if err != nil {
 		return Response{}, err
 	}
@@ -127,7 +139,6 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
 	if asn.AutonomousSystemNumber > 0 {
 		autonomousSystemNumber = fmt.Sprintf("AS%d", asn.AutonomousSystemNumber)
 	}
-	userAgent := userAgentFromRequest(r)
 	response = &Response{
 		IP:         ip,
 		IPDecimal:  ipDecimal,
@@ -145,9 +156,9 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
 		ASN:        autonomousSystemNumber,
 		ASNOrg:     asn.AutonomousSystemOrganization,
 		Hostname:   hostname,
-		UserAgent:  userAgent,
 	}
 	s.cache.Set(ip, response)
+	response.UserAgent = userAgentFromRequest(r)
 	return *response, nil
 }
 
@@ -157,7 +168,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
 	if err != nil || port < 1 || port > 65535 {
 		return PortResponse{Port: port}, fmt.Errorf("invalid port: %s", lastElement)
 	}
-	ip, err := ipFromRequest(s.IPHeaders, r)
+	ip, err := ipFromRequest(s.IPHeaders, r, false)
 	if err != nil {
 		return PortResponse{Port: port}, err
 	}
@@ -170,7 +181,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
 }
 
 func (s *Server) CLIHandler(w http.ResponseWriter, r *http.Request) *appError {
-	ip, err := ipFromRequest(s.IPHeaders, r)
+	ip, err := ipFromRequest(s.IPHeaders, r, true)
 	if err != nil {
 		return internalServerError(err)
 	}
diff --git a/http/http_test.go b/http/http_test.go
index ab3a789..00e58a4 100644
--- a/http/http_test.go
+++ b/http/http_test.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"testing"
 
 	"github.com/mpolden/echoip/iputil/geo"
@@ -139,6 +140,8 @@ func TestJSONHandlers(t *testing.T) {
 		{s.URL + "/port/0", `{"error":"invalid port: 0"}`, 400},
 		{s.URL + "/port/65537", `{"error":"invalid port: 65537"}`, 400},
 		{s.URL + "/port/31337", `{"ip":"127.0.0.1","port":31337,"reachable":true}`, 200},
+		{s.URL + "/port/80", `{"ip":"127.0.0.1","port":80,"reachable":true}`, 200},            // checking that our test server is reachable on port 80
+		{s.URL + "/port/80?ip=1.3.3.7", `{"ip":"127.0.0.1","port":80,"reachable":true}`, 200}, // ensuring that the "ip" parameter is not usable to check remote host ports
 		{s.URL + "/foo", `{"error":"404 page not found"}`, 404},
 		{s.URL + "/health", `{"status":"OK"}`, 200},
 	}
@@ -165,22 +168,29 @@ func TestIPFromRequest(t *testing.T) {
 		trustedHeaders []string
 		out            string
 	}{
-		{"127.0.0.1:9999", "", "", nil, "127.0.0.1"},                                                          // No header given
-		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", nil, "127.0.0.1"},                                          // Trusted header is empty
-		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Foo-Bar"}, "127.0.0.1"},                        // Trusted header does not match
-		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"},       // Trusted header matches
-		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"}, // Second trusted header matches
-		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7,4.2.4.2", []string{"X-Forwarded-For"}, "1.3.3.7"},      // X-Forwarded-For with multiple entries (commas separator)
-		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7, 4.2.4.2", []string{"X-Forwarded-For"}, "1.3.3.7"},     // X-Forwarded-For with multiple entries (space+comma separator)
-		{"127.0.0.1:9999", "X-Forwarded-For", "", []string{"X-Forwarded-For"}, "127.0.0.1"},                   // Empty header
+		{"127.0.0.1:9999", "", "", nil, "127.0.0.1"},                                                                // No header given
+		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", nil, "127.0.0.1"},                                                // Trusted header is empty
+		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Foo-Bar"}, "127.0.0.1"},                              // Trusted header does not match
+		{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"},             // Trusted header matches
+		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"},       // Second trusted header matches
+		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7,4.2.4.2", []string{"X-Forwarded-For"}, "1.3.3.7"},            // X-Forwarded-For with multiple entries (commas separator)
+		{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7, 4.2.4.2", []string{"X-Forwarded-For"}, "1.3.3.7"},           // X-Forwarded-For with multiple entries (space+comma separator)
+		{"127.0.0.1:9999", "X-Forwarded-For", "", []string{"X-Forwarded-For"}, "127.0.0.1"},                         // Empty header
+		{"127.0.0.1:9999?ip=1.2.3.4", "", "", nil, "1.2.3.4"},                                                       // passed in "ip" parameter
+		{"127.0.0.1:9999?ip=1.2.3.4", "X-Forwarded-For", "1.3.3.7,4.2.4.2", []string{"X-Forwarded-For"}, "1.2.3.4"}, // ip parameter wins over X-Forwarded-For with multiple entries
 	}
 	for _, tt := range tests {
+		u, err := url.Parse("http://" + tt.remoteAddr)
+		if err != nil {
+			t.Fatal(err)
+		}
 		r := &http.Request{
-			RemoteAddr: tt.remoteAddr,
+			RemoteAddr: u.Host,
 			Header:     http.Header{},
+			URL:        u,
 		}
 		r.Header.Add(tt.headerKey, tt.headerValue)
-		ip, err := ipFromRequest(tt.trustedHeaders, r)
+		ip, err := ipFromRequest(tt.trustedHeaders, r, true)
 		if err != nil {
 			t.Fatal(err)
 		}
diff --git a/index.html b/index.html
index b26d8d3..95ac56b 100644
--- a/index.html
+++ b/index.html
@@ -135,6 +135,7 @@ $ http {{ .Host }}/asn
 $ http {{ .Host }}/json
 {{ .JSON }}</pre>
         <p>Setting the <code>Accept: application/json</code> header also works as expected.</p>
+	<p>All endpoints (except <code>/port</code>) can return information about a custom IP address specified via <code>?ip=</code> query parameter.</p>
         <h2>Plain output</h2>
         <p>Always returns the IP address including a trailing newline, regardless of user agent.</p>
         <pre>