Use headers to check ip addresses for use behind proxy

• Break file up to multiple files
 • Can pass in headers through command line arguments, defaults for
   checking X-Real-IP and X-Forwarded-For
 • Add logging for requests with timing
 • Add 404 support for missing URLs

ip.go

@@ -0,0 +1,33 @@
+package main
+
+import (
+	"errors"
+	"net"
+)
+
+// IP is the net.IP with a version string to allow storing the version
+type IP struct {
+	net.IP
+	version string
+}
+
+func (ip *IP) isV6() bool {
+	return ip.To4() == nil
+}
+
+// NewIP returns a new instance of an IP address
+func NewIP(addr string) (*IP, error) {
+	parsedIP := net.ParseIP(addr)
+	if parsedIP == nil {
+		return nil, errors.New("Could not parse address " + addr)
+	}
+
+	result := &IP{parsedIP, ""}
+	if result.isV6() {
+		result.version = "v6"
+	} else {
+		result.version = "v4"
+	}
+
+	return result, nil
+}

logger.go

@@ -0,0 +1,35 @@
+package main
+
+import (
+	"log"
+	"net/http"
+	"os"
+	"time"
+)
+
+// Logger interface
+type Logger interface {
+	Println(v ...interface{})
+	Printf(format string, v ...interface{})
+}
+
+// LoggerHandler interface
+type LoggerHandler interface {
+	Logger
+	ServeHTTP(http.ResponseWriter, *http.Request, http.HandlerFunc)
+}
+
+// NewLogger returns a logger interface
+func NewLogger() LoggerHandler {
+	return &logger{log.New(os.Stdout, appLabel, log.Ldate|log.Lmicroseconds)}
+}
+
+type logger struct {
+	Logger
+}
+
+func (l *logger) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
+	start := time.Now()
+	next(rw, r)
+	l.Printf("Completed %s to %s in %v", r.Method, r.URL.Path, time.Since(start))
+}

server.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"io"
+	"net"
+	"net/http"
+)
+
+// Server that handles http responses
+type Server struct {
+	headerNames []string
+	logger      LoggerHandler
+}
+
+func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	for _, headerName := range s.headerNames {
+		possibleIP := req.Header.Get(headerName)
+		if possibleIP == "" {
+			continue
+		}
+
+		ip, err := NewIP(possibleIP)
+		if err != nil {
+			continue
+		}
+
+		s.sendResponse(w, ip)
+		return
+	}
+
+	addr, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err != nil {
+		s.handleError(w, http.StatusBadRequest, err, "Could not get IP address from request")
+		return
+	}
+
+	ip, err := NewIP(addr)
+	if err != nil {
+		s.handleError(w, http.StatusBadRequest, err, "Could not parse IP address"+addr)
+		return
+	}
+
+	s.sendResponse(w, ip)
+}
+
+func (s *Server) sendResponse(resp http.ResponseWriter, ip *IP) {
+	io.WriteString(resp, ip.String()+"\n")
+	s.logger.Printf("Request from %s %s\n", ip.version, ip)
+}
+
+func (s *Server) handleError(resp http.ResponseWriter, status int, err error, message string) {
+	resp.WriteHeader(status)
+	io.WriteString(resp, message)
+	s.logger.Printf("Error handling request: %s (%s)", message, err)
+}

what-is-my-ip.go

@@ -1,20 +1,24 @@
 package main
 
 import (
-	"fmt"
-	"io"
-	"log"
-	"net"
+	"errors"
 	"net/http"
 	"os"
 
+	"github.com/gorilla/pat"
 	"github.com/urfave/cli"
+	"github.com/urfave/negroni"
+)
+
+const (
+	appName  = "what-is-my-ip"
+	appLabel = "[" + appName + "] "
 )
 
 func main() {
 	app := cli.NewApp()
-	app.Name = "what-is-my-ip"
-	app.Version = "0.2.0"
+	app.Name = appName
+	app.Version = "0.3.0"
 	app.Usage = "Webapp that gives the IP address for incoming requests"
 	app.Action = action
 
@@ -26,6 +30,12 @@ func main() {
 			Usage:  "Address and port to bind to",
 		},
 
+		cli.StringSliceFlag{
+			Name:  "header, H",
+			Value: &cli.StringSlice{"X-Real-IP", "X-Forwarded-For"},
+			Usage: "Headers to check in before checking the request itself",
+		},
+
 		cli.StringFlag{
 			Name:   "certificate, c",
 			Value:  "cert.pem",
@@ -45,26 +55,38 @@ func main() {
 }
 
 func action(c *cli.Context) error {
-	http.HandleFunc("/", getIP)
 	address := c.String("address")
 	certificate := c.String("certificate")
 	key := c.String("key")
-	fmt.Printf("listening on %s\n", address)
-	return http.ListenAndServeTLS(address, certificate, key, nil)
+	handler, service := getHandler(c.StringSlice("header"))
+	service.logger.Printf("listening on %s\n", address)
+	return http.ListenAndServeTLS(address, certificate, key, handler)
 }
 
-func getIP(w http.ResponseWriter, req *http.Request) {
-	ip, _, err := net.SplitHostPort(req.RemoteAddr)
-	if err != nil {
-		handleError(w, err, "Could not get IP address from request")
-		return
+func getHandler(headers []string) (http.Handler, *Server) {
+	service := &Server{
+		headerNames: headers,
+		logger:      NewLogger(),
 	}
-	io.WriteString(w, ip+"\n")
-	log.Printf("Give response IP %s\n", ip)
-}
 
-func handleError(resp http.ResponseWriter, err error, message string) {
-	resp.WriteHeader(http.StatusBadRequest)
-	io.WriteString(resp, message)
-	log.Printf("Error handling request: %s (%s)", message, err)
+	recover := negroni.NewRecovery()
+	recover.Logger = service.logger
+
+	router := pat.New()
+
+	router.Get("/", func(resp http.ResponseWriter, req *http.Request) {
+		if req.URL.Path == "/" {
+			service.ServeHTTP(resp, req)
+		} else {
+			err := errors.New("path \"" + req.URL.Path + "\" not found")
+			service.handleError(resp, http.StatusNotFound, err, "404 page not found\n")
+		}
+	})
+
+	n := negroni.New()
+	n.Use(recover)
+	n.Use(service.logger)
+	n.UseHandler(router)
+
+	return n, service
 }