Remove port when using headers in c.RealIP

Co-authored-by: Elwin Tamminga <elwintamminga@gmail.com>
Co-committed-by: Elwin Tamminga <elwintamminga@gmail.com>

context.go

@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"encoding/json"
 	"io"
+	"net"
 	"net/http"
+	"strings"
 
 	"github.com/julienschmidt/httprouter"
 )
@@ -18,7 +20,8 @@ type Context struct {
 	store    map[string]interface{}
 }
 
-func newContext(router *Router, res http.ResponseWriter, req *http.Request, param httprouter.Params) *Context {
+// NewContext creates a new context, this function is only exported for use in tests
+func NewContext(router *Router, res http.ResponseWriter, req *http.Request, param httprouter.Params) *Context {
 	return &Context{router, req, res, param.ByName, make(map[string]interface{})}
 }
 
@@ -99,3 +102,20 @@ func (c *Context) Set(key string, value interface{}) {
 func (c *Context) Get(key string) interface{} {
 	return c.store[key]
 }
+
+// RealIP uses proxy headers for the real ip, if none exist the IP of the current connection is returned
+func (c *Context) RealIP() string {
+	reqIP := c.Request.RemoteAddr
+
+	if ip := c.Request.Header.Get(`X-Forwarded-For`); ip != `` {
+		reqIP = strings.Split(ip, `, `)[0]
+	} else if ip := c.Request.Header.Get(`X-Real-IP`); ip != `` {
+		reqIP = ip
+	}
+
+	ra, _, _ := net.SplitHostPort(reqIP)
+	if ra != `` {
+		reqIP = ra
+	}
+	return reqIP
+}

router.go

@@ -38,6 +38,7 @@ type Router struct {
 	NotFoundHandler         Handle
 	MethodNotAllowedHandler Handle
 	ErrorHandler            ErrorHandle
+	TrimTrailingSlashes     bool
 	server                  *http.Server
 }
 
@@ -102,6 +103,7 @@ func (r *Router) Start(addr string) error {
 	return r.server.ListenAndServe()
 }
 
+// StartTLS starts a TLS web server using the given key, cert and config and binds to the given address
 func (r *Router) StartTLS(addr, certFile, keyFile string, conf *tls.Config) error {
 	httpr := r.getHttpr()
 
@@ -115,7 +117,7 @@ func (r *Router) Stop() error {
 		return nil
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(time.Second*5))
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 	err := r.server.Shutdown(ctx)
 	if err == context.DeadlineExceeded {
@@ -126,7 +128,7 @@ func (r *Router) Stop() error {
 	return err
 }
 
-func (r *Router) getHttpr() *httprouter.Router {
+func (r *Router) getHttpr() http.Handler {
 	httpr := httprouter.New()
 
 	for _, v := range r.routes {
@@ -135,7 +137,11 @@ func (r *Router) getHttpr() *httprouter.Router {
 			handle = handlePOST(r, v.Handle)
 		}
 
-		httpr.Handle(v.Method, v.Path, handleReq(r, handle, append(r.middleware, v.Middleware...)))
+		middleware := make([]Middleware, len(r.middleware)+len(v.Middleware))
+		copy(middleware, r.middleware)
+		copy(middleware[len(r.middleware):], v.Middleware)
+
+		httpr.Handle(v.Method, v.Path, handleReq(r, handle, middleware))
 	}
 
 	httpr.NotFound = http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
@@ -147,10 +153,22 @@ func (r *Router) getHttpr() *httprouter.Router {
 	})
 
 	httpr.PanicHandler = func(res http.ResponseWriter, req *http.Request, err interface{}) {
-		c := newContext(r, res, req, nil)
+		c := NewContext(r, res, req, nil)
 		r.ErrorHandler(c, err)
 	}
 
+	if r.TrimTrailingSlashes {
+		httpr.RedirectTrailingSlash = false
+
+		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			l := len(req.URL.Path)
+			if l > 1 && req.URL.Path[l-1] == '/' {
+				req.URL.Path = req.URL.Path[:l-1]
+			}
+			httpr.ServeHTTP(w, req)
+		})
+	}
+
 	return httpr
 }
 
@@ -176,8 +194,6 @@ func checkInterfaceHandle(f interface{}) {
 	if rt.In(0) != reflect.TypeOf(&Context{}) {
 		panic(`handle should accept Context as first argument`)
 	}
-
-	return
 }
 
 func handlePOST(r *Router, f interface{}) Handle {
@@ -208,10 +224,10 @@ func handlePOST(r *Router, f interface{}) Handle {
 
 func handleReq(r *Router, handle Handle, m []Middleware) httprouter.Handle {
 	return func(res http.ResponseWriter, req *http.Request, param httprouter.Params) {
-		c := newContext(r, res, req, param)
+		c := NewContext(r, res, req, param)
 
 		f := handle
-		for i := len(m) - 1; i >= 0; i-- { // TODO: 1,2,3 of 3,2,1
+		for i := len(m) - 1; i >= 0; i-- {
 			f = m[i](f)
 		}