diff --git a/main.go b/main.go index 9de6e42..f399ae1 100644 --- a/main.go +++ b/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 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) - }) - + 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)) }