Merge branch 'main' of github.com:nmasse-itix/wimip

main.go

@@ -5,6 +5,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"strings"
 
 	"github.com/chrismarget/arp"
 	"github.com/spf13/viper"
@@ -15,6 +16,38 @@ type IPInfo struct {
 	MacAddress string `json:"mac,omitempty"`
 }
 
+type IPInfoHandler struct {
+
+}
+
+func (h *IPInfoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	var info IPInfo
+
+	ip, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		w.WriteHeader(500)
+	}
+	
+	fwdAddress := r.Header.Get("X-Forwarded-For")
+	if fwdAddress != "" {
+		ip = fwdAddress // If it's a single IP, then awesome!
+
+		// If we got an array... grab the first IP
+		ips := strings.Split(fwdAddress, ", ")
+		if len(ips) > 1 {
+			ip = ips[0]
+		}
+	}
+
+	info.IP = ip
+
+	arp.CacheUpdate()
+	info.MacAddress = arp.Search(ip)
+
+	out, _ := json.MarshalIndent(info, "", " ")
+	w.Write(out)
+}
+
 func main() {
 	log.SetFlags(log.Flags() &^ (log.Ldate | log.Ltime))
 
@@ -22,21 +55,7 @@ func main() {
 	viper.SetEnvPrefix("WIMIP")
 	viper.AutomaticEnv()
 
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	var handler *IPInfoHandler
-		var info IPInfo
-		ip, _, err := net.SplitHostPort(r.RemoteAddr)
-		if err != nil {
-			w.WriteHeader(500)
-		}
-		info.IP = ip
-
-		arp.CacheUpdate()
-		info.MacAddress = arp.Search(ip)
-
-		out, _ := json.MarshalIndent(info, "", " ")
-		w.Write(out)
-	})
-
 	log.Printf("Listening on %s...", viper.GetString("ListenAddr"))
-	log.Fatal(http.ListenAndServe(viper.GetString("ListenAddr"), nil))
+	log.Fatal(http.ListenAndServe(viper.GetString("ListenAddr"), handler))
 }