diff --git a/context.go b/context.go index e6cc4b3..e9c2fd6 100644 --- a/context.go +++ b/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 +} diff --git a/router.go b/router.go index a878555..d869fc6 100644 --- a/router.go +++ b/router.go @@ -1,8 +1,11 @@ package router import ( + "context" + "crypto/tls" "net/http" "reflect" + "time" "github.com/julienschmidt/httprouter" ) @@ -35,6 +38,8 @@ type Router struct { NotFoundHandler Handle MethodNotAllowedHandler Handle ErrorHandler ErrorHandle + TrimTrailingSlashes bool + server *http.Server } // New returns a new Router @@ -94,10 +99,36 @@ func (r *Router) OPTIONS(path string, handle Handle, middleware ...Middleware) { func (r *Router) Start(addr string) error { httpr := r.getHttpr() - return http.ListenAndServe(addr, httpr) + r.server = &http.Server{Addr: addr, Handler: httpr} + return r.server.ListenAndServe() } -func (r *Router) getHttpr() *httprouter.Router { +// 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() + + r.server = &http.Server{Addr: addr, Handler: httpr, TLSConfig: conf} + return r.server.ListenAndServeTLS(certFile, keyFile) +} + +// Stop stops the web server +func (r *Router) Stop() error { + if r.server == nil { + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) + defer cancel() + err := r.server.Shutdown(ctx) + if err == context.DeadlineExceeded { + err = r.server.Close() + } + r.server = nil + + return err +} + +func (r *Router) getHttpr() http.Handler { httpr := httprouter.New() for _, v := range r.routes { @@ -106,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) { @@ -118,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 } @@ -147,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 { @@ -179,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) }