consider X-Forwarded-For header, pass port/path as parameters

1 files changed, 32 insertions, 20 deletions

diff --git a/main.go b/main.go
index 1c00d14..4b9b9cf 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"encoding/json"
+	"flag"
 	"fmt"
 	"log"
 	"net"
@@ -11,14 +12,30 @@ import (
 	"github.com/oschwald/geoip2-golang"
 )
 
-type geodb struct {
-	db *geoip2.Reader
-}
-
 func floatToString(num float64) string {
 	return strconv.FormatFloat(num, 'f', 6, 64)
 }
 
+func getRemoteIP(req *http.Request) string {
+	forward := req.Header.Get("X-Forwarded-For")
+	ipstr := ""
+	if forward != "" {
+		ipstr = forward
+	} else {
+		ip, _, err := net.SplitHostPort(req.RemoteAddr)
+		if err != nil {
+			log.Fatal(err)
+		}
+		netIP := net.ParseIP(ip)
+		ipstr = netIP.String()
+	}
+	return ipstr
+}
+
+type geodb struct {
+	db *geoip2.Reader
+}
+
 func (g *geodb) getRecordForIP(ipstr string) *geoip2.City {
 	ip := net.ParseIP(ipstr)
 	record, err := g.db.City(ip)
@@ -33,12 +50,7 @@ type jsonHandler struct {
 }
 
 func (jh *jsonHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	ip, _, err := net.SplitHostPort(req.RemoteAddr)
-	if err != nil {
-		log.Fatal(err)
-	}
-	netIP := net.ParseIP(ip)
-	ipstr := netIP.String()
+	ipstr := getRemoteIP(req)
 	record := jh.geoipdb.getRecordForIP(ipstr)
 	data := map[string]string{
 		"ip":   ipstr,
@@ -56,24 +68,23 @@ type txtHandler struct {
 }
 
 func (th *txtHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	ip, _, err := net.SplitHostPort(req.RemoteAddr)
-	if err != nil {
-		log.Fatal(err)
-	}
-	netIP := net.ParseIP(ip)
-	ipstr := netIP.String()
+	ipstr := getRemoteIP(req)
 	record := th.geoipdb.getRecordForIP(ipstr)
 
 	fmt.Fprintf(w, "Your IP: %s\n", ipstr)
 	fmt.Fprintf(w, "Your Country: %s\n", record.Country.IsoCode)
 	fmt.Fprintf(w, "Your City: %s\n", record.City.Names["en"])
 	fmt.Fprintf(w, "Your Coordinates: %s, %s\n",
-        floatToString(record.Location.Latitude),
+		floatToString(record.Location.Latitude),
 		floatToString(record.Location.Longitude))
 }
 
 func main() {
-	db, err := geoip2.Open("/var/lib/GeoIP/GeoLite2-City.mmdb")
+	var port = flag.Int("port", 9001, "port where the service listens on")
+	var dbpath = flag.String("geodb", "/var/lib/GeoIP/GeoLite2-City.mmdb", "path to the GeoLite2-City database")
+	flag.Parse()
+
+	db, err := geoip2.Open(*dbpath)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -88,6 +99,7 @@ func main() {
 	th := &txtHandler{&geoipdb}
 	mux.Handle("/", th)
 
-	log.Println("Listening...")
-	http.ListenAndServe(":9001", mux)
+	log.Println("Started Geolocation Service")
+	log.Printf("Listening on port %v...\n", *port)
+	http.ListenAndServe(":"+strconv.Itoa(*port), mux)
 }