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,19 +16,29 @@ type IPInfo struct {
 	MacAddress string `json:"mac,omitempty"`
 }
 
-func main() {
+type IPInfoHandler struct {
-	log.SetFlags(log.Flags() &^ (log.Ldate | log.Ltime))
 
-	viper.SetDefault("ListenAddr", ":8080")
+}
-	viper.SetEnvPrefix("WIMIP")
-	viper.AutomaticEnv()
 
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+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()
@@ -35,8 +46,16 @@ func main() {
 
 	out, _ := json.MarshalIndent(info, "", " ")
 	w.Write(out)
-	})
+}
+
+func main() {
+	log.SetFlags(log.Flags() &^ (log.Ldate | log.Ltime))
+
+	viper.SetDefault("ListenAddr", ":8080")
+	viper.SetEnvPrefix("WIMIP")
+	viper.AutomaticEnv()
 
+	var handler *IPInfoHandler
 	log.Printf("Listening on %s...", viper.GetString("ListenAddr"))
-	log.Fatal(http.ListenAndServe(viper.GetString("ListenAddr"), nil))
+	log.Fatal(http.ListenAndServe(viper.GetString("ListenAddr"), handler))
 }