Flexible handler and middleware

Signed-off-by: Vishal Rana <vr@labstack.com>

README.md

@@ -3,8 +3,9 @@ Bolt is a fast HTTP router (zero memory allocation) + micro web framework in Go.
 
 ### Features
 - Zippy router.
+- Extensible middleware.
+- Bring you own handler / middleware.
 - Serve static files, including index.
-- Extensible middleware which also allows you to use third party handler / middleware.
 
 ### Example
 https://github.com/labstack/bolt/tree/master/example

bolt.go

@@ -9,14 +9,17 @@ import (
 type (
 	Bolt struct {
 		Router                     *router
-		handlers                   []HandlerFunc
+		middleware                 []MiddlewareFunc
 		maxParam                   byte
 		notFoundHandler            HandlerFunc
 		methodNotAllowedHandler    HandlerFunc
 		internalServerErrorHandler HandlerFunc
 		pool                       sync.Pool
 	}
-	HandlerFunc func(*Context)
+	Handler        interface{}
+	HandlerFunc    func(*Context)
+	Middleware     interface{}
+	MiddlewareFunc func(HandlerFunc) HandlerFunc
 )
 
 const (
@@ -28,34 +31,21 @@ const (
 	HeaderContentType        = "Content-Type"
 )
 
-// Methods is a map for looking up HTTP method index.
-var Methods = map[string]uint8{
-	"CONNECT": 0,
-	"DELETE":  1,
-	"GET":     2,
-	"HEAD":    3,
-	"OPTIONS": 4,
-	"PATCH":   5,
-	"POST":    6,
-	"PUT":     7,
-	"TRACE":   8,
-}
-
 // New creates a bolt instance.
 func New() (b *Bolt) {
 	b = &Bolt{
 		maxParam: 5,
 		notFoundHandler: func(c *Context) {
 			http.Error(c.Response, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-			c.Halt()
+			// c.Halt()
 		},
 		methodNotAllowedHandler: func(c *Context) {
 			http.Error(c.Response, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
-			c.Halt()
+			// c.Halt()
 		},
 		internalServerErrorHandler: func(c *Context) {
 			http.Error(c.Response, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
-			c.Halt()
+			// c.Halt()
 		},
 	}
 	b.Router = NewRouter(b)
@@ -64,13 +54,24 @@ func New() (b *Bolt) {
 			Response: &response{},
 			params:   make(Params, b.maxParam),
 			store:    make(store),
-			i:        -1,
-			bolt:     b,
+			// i:        -1,
+			bolt: b,
 		}
 	}
 	return
 }
 
+// NOP
+func (h HandlerFunc) ServeHTTP(r http.ResponseWriter, w *http.Request) {
+}
+
+func (b *Bolt) Sub(prefix string, m ...MiddlewareFunc) *Bolt {
+	return &Bolt{
+		// prefix:   b.prefix + prefix,
+		// middleware: append(b.handlers, handlers...),
+	}
+}
+
 // MaxParam sets the max path params allowed. Default is 5, good enough for
 // many users.
 func (b *Bolt) MaxParam(n uint8) {
@@ -92,74 +93,68 @@ func (b *Bolt) InternalServerErrorHandler(h HandlerFunc) {
 	b.internalServerErrorHandler = h
 }
 
-// Chain adds middleware to the chain.
-func (b *Bolt) Chain(h ...HandlerFunc) {
-	b.handlers = append(b.handlers, h...)
-}
-
-// Wrap wraps any http.Handler into bolt.HandlerFunc. It facilitates to use
-// third party handler / middleware with bolt.
-func (b *Bolt) Wrap(h http.Handler) HandlerFunc {
-	return func(c *Context) {
-		h.ServeHTTP(c.Response, c.Request)
-		c.Next()
+// Use adds handler to the middleware chain.
+func (b *Bolt) Use(m ...Middleware) {
+	for _, h := range m {
+		b.middleware = append(b.middleware, wrapM(h))
 	}
 }
 
 // Connect adds a CONNECT route.
-func (b *Bolt) Connect(path string, h ...HandlerFunc) {
+func (b *Bolt) Connect(path string, h Handler) {
 	b.Handle("CONNECT", path, h)
 }
 
 // Delete adds a DELETE route.
-func (b *Bolt) Delete(path string, h ...HandlerFunc) {
+func (b *Bolt) Delete(path string, h Handler) {
 	b.Handle("DELETE", path, h)
 }
 
 // Get adds a GET route.
-func (b *Bolt) Get(path string, h ...HandlerFunc) {
+func (b *Bolt) Get(path string, h Handler) {
 	b.Handle("GET", path, h)
 }
 
 // Head adds a HEAD route.
-func (b *Bolt) Head(path string, h ...HandlerFunc) {
+func (b *Bolt) Head(path string, h Handler) {
 	b.Handle("HEAD", path, h)
 }
 
 // Options adds an OPTIONS route.
-func (b *Bolt) Options(path string, h ...HandlerFunc) {
+func (b *Bolt) Options(path string, h Handler) {
 	b.Handle("OPTIONS", path, h)
 }
 
 // Patch adds a PATCH route.
-func (b *Bolt) Patch(path string, h ...HandlerFunc) {
+func (b *Bolt) Patch(path string, h Handler) {
 	b.Handle("PATCH", path, h)
 }
 
 // Post adds a POST route.
-func (b *Bolt) Post(path string, h ...HandlerFunc) {
+func (b *Bolt) Post(path string, h Handler) {
 	b.Handle("POST", path, h)
 }
 
 // Put adds a PUT route.
-func (b *Bolt) Put(path string, h ...HandlerFunc) {
+func (b *Bolt) Put(path string, h Handler) {
 	b.Handle("PUT", path, h)
 }
 
 // Trace adds a TRACE route.
-func (b *Bolt) Trace(path string, h ...HandlerFunc) {
+func (b *Bolt) Trace(path string, h Handler) {
 	b.Handle("TRACE", path, h)
 }
 
 // Handle adds method, path  handler to the router.
-func (b *Bolt) Handle(method, path string, h []HandlerFunc) {
-	h = append(b.handlers, h...)
-	l := len(h)
-	b.Router.Add(method, path, func(c *Context) {
-		c.handlers = h
-		c.l = l
-		c.Next()
-	})
+func (b *Bolt) Handle(method, path string, h Handler) {
+	b.Router.Add(method, path, b.wrapH(h))
+	// hs := append(b.middleware, wrap(h, false))
+	// l := len(hs)
+	// b.Router.Add(method, path, func(c *Context) {
+	// 	c.handlers = hs
+	// 	c.l = l
+	// c.Next()
+	// })
 }
 
 // Static serves static files.
@@ -183,10 +178,14 @@ func (b *Bolt) Index(file string) {
 }
 
 func (b *Bolt) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
-	// Find and execute handler
 	h, c, s := b.Router.Find(r.Method, r.URL.Path)
 	c.reset(rw, r)
 	if h != nil {
+		// Middleware
+		for i := len(b.middleware) - 1; i >= 0; i-- {
+			h = b.middleware[i](h)
+		}
+		// Handler
 		h(c)
 	} else {
 		if s == NotFound {
@@ -202,6 +201,40 @@ func (b *Bolt) Run(addr string) {
 	log.Fatal(http.ListenAndServe(addr, b))
 }
 
-func (b *Bolt) Stop(addr string) {
-	panic("implement it")
+// wraps Handler
+func (b *Bolt) wrapH(h Handler) HandlerFunc {
+	switch h := h.(type) {
+	case func(*Context):
+		return HandlerFunc(h)
+	case http.HandlerFunc:
+		return func(c *Context) {
+			h.ServeHTTP(c.Response, c.Request)
+		}
+	default:
+		panic("bolt: unknown handler")
+	}
+}
+
+// wraps Middleware
+func wrapM(m Middleware) MiddlewareFunc {
+	switch m := m.(type) {
+	case func(HandlerFunc) HandlerFunc:
+		return MiddlewareFunc(m)
+	case func(http.ResponseWriter, *http.Request):
+		return func(h HandlerFunc) HandlerFunc {
+			return func(c *Context) {
+				m(c.Response, c.Request)
+				h(c)
+			}
+		}
+	case func(http.Handler) http.Handler:
+		return func(h HandlerFunc) HandlerFunc {
+			return func(c *Context) {
+				m(h).ServeHTTP(c.Response, c.Request)
+				h(c)
+			}
+		}
+	default:
+		panic("bolt: unknown middleware")
+	}
 }

context.go

@@ -16,9 +16,9 @@ type (
 		params   Params
 		handlers []HandlerFunc
 		store    map[string]interface{}
-		l        int // Handlers' length
-		i        int // Current handler index
-		bolt     *Bolt
+		// l        int // Handlers' length
+		// i        int // Current handler index
+		bolt *Bolt
 	}
 	store map[string]interface{}
 )
@@ -64,12 +64,12 @@ func (c *Context) JSON(n int, i interface{}) {
 // }
 
 // Next executes the next handler in the chain.
-func (c *Context) Next() {
-	c.i++
-	if c.i < c.l {
-		c.handlers[c.i](c)
-	}
-}
+// func (c *Context) Next() {
+// 	c.i++
+// 	if c.i < c.l {
+// 		c.handlers[c.i](c)
+// 	}
+// }
 
 // Get retrieves data from the context.
 func (c *Context) Get(key string) interface{} {
@@ -89,10 +89,10 @@ func (c *Context) Redirect(n int, url string) {
 func (c *Context) reset(rw http.ResponseWriter, r *http.Request) {
 	c.Response.reset(rw)
 	c.Request = r
-	c.i = -1
+	// c.i = -1
 }
 
 // Halt halts the current request.
-func (c *Context) Halt() {
-	c.i = c.l
-}
+// func (c *Context) Halt() {
+// 	c.i = c.l
+// }

example/main.go

@@ -5,6 +5,8 @@ import (
 
 	"github.com/labstack/bolt"
 	mw "github.com/labstack/bolt/middleware"
+	"github.com/rs/cors"
+	"github.com/thoas/stats"
 )
 
 type user struct {
@@ -41,7 +43,7 @@ func getUser(c *bolt.Context) {
 
 func main() {
 	b := bolt.New()
-	b.Chain(mw.Logger())
+	b.Use(mw.Logger)
 	b.Index("public/index.html")
 	b.Static("/js", "public/js")
 	b.Post("/users", createUser)

middleware/auth.go

@@ -10,10 +10,10 @@ import (
 )
 
 type (
-	BasicAuthFunc       func(usr, pwd string) bool
+	BasicAuthFunc       func(string, string) bool
 	AuthorizedHandler   bolt.HandlerFunc
-	UnauthorizedHandler func(c *bolt.Context, err error)
-	JwtKeyFunc          func(kid string) ([]byte, error)
+	UnauthorizedHandler func(*bolt.Context, error)
+	JwtKeyFunc          func(string) ([]byte, error)
 	Claims              map[string]interface{}
 )
 
@@ -57,7 +57,7 @@ func JwtAuth(ah AuthorizedHandler, uah UnauthorizedHandler, fn JwtKeyFunc) bolt.
 			if t.Valid {
 				c.Set("claims", Claims(t.Claims))
 				ah(c)
-				c.Next()
+				// c.Next()
 			} else {
 				// TODO: capture errors
 				uah(c, err)

middleware/logger.go

@@ -8,10 +8,10 @@ import (
 	"github.com/labstack/gommon/color"
 )
 
-func Logger() bolt.HandlerFunc {
-	return func(c *bolt.Context) {
+func Logger(h bolt.HandlerFunc) bolt.HandlerFunc {
+	return bolt.HandlerFunc(func(c *bolt.Context) {
 		start := time.Now()
-		c.Next()
+		h(c)
 		end := time.Now()
 		col := color.Green
 		m := c.Request.Method
@@ -28,5 +28,5 @@ func Logger() bolt.HandlerFunc {
 		}
 
 		log.Printf("%s %s %s %s", m, p, col(s), end.Sub(start))
-	}
+	})
 }

response.go

@@ -17,14 +17,16 @@ type (
 	}
 )
 
-func (r *response) WriteHeader(c int) {
+
+func (r *response) WriteHeader(n int) {
+	// TODO: fix when halted.
 	if r.committed {
 		// TODO: Warning
 		log.Println("bolt: response already committed")
 		return
 	}
-	r.status = c
-	r.ResponseWriter.WriteHeader(c)
+	r.status = n
+	r.ResponseWriter.WriteHeader(n)
 	r.committed = true
 }
 

router.go

@@ -39,11 +39,24 @@ const (
 	NotAllowed
 )
 
+//  methods is a map for looking up HTTP method index.
+var methods = map[string]uint8{
+	"CONNECT": 0,
+	"DELETE":  1,
+	"GET":     2,
+	"HEAD":    3,
+	"OPTIONS": 4,
+	"PATCH":   5,
+	"POST":    6,
+	"PUT":     7,
+	"TRACE":   8,
+}
+
 func NewRouter(b *Bolt) (r *router) {
 	r = &router{
 		root: &node{
 			prefix:   "",
-			handlers: make([]HandlerFunc, len(Methods)),
+			handlers: make([]HandlerFunc, len(methods)),
 			edges:    edges{},
 		},
 		bolt: b,
@@ -95,7 +108,7 @@ func (r *router) insert(method, path string, h HandlerFunc, has ntype) {
 			cn.prefix = search
 			cn.has = has
 			if h != nil {
-				cn.handlers[Methods[method]] = h
+				cn.handlers[methods[method]] = h
 			}
 			return
 		} else if l < pl {
@@ -107,15 +120,15 @@ func (r *router) insert(method, path string, h HandlerFunc, has ntype) {
 			cn.label = cn.prefix[0]
 			cn.prefix = cn.prefix[:l]
 			cn.has = snode
-			cn.handlers = make([]HandlerFunc, len(Methods))
+			cn.handlers = make([]HandlerFunc, len(methods))
 
 			if l == sl {
 				// At parent node
-				cn.handlers[Methods[method]] = h
+				cn.handlers[methods[method]] = h
 			} else {
 				// Need to fork a node
 				n = newNode(search[l:], has, nil, nil)
-				n.handlers[Methods[method]] = h
+				n.handlers[methods[method]] = h
 				cn.edges = append(cn.edges, n)
 			}
 			break
@@ -125,7 +138,7 @@ func (r *router) insert(method, path string, h HandlerFunc, has ntype) {
 			if e == nil {
 				n := newNode(search, has, nil, nil)
 				if h != nil {
-					n.handlers[Methods[method]] = h
+					n.handlers[methods[method]] = h
 				}
 				cn.edges = append(cn.edges, n)
 				break
@@ -135,7 +148,7 @@ func (r *router) insert(method, path string, h HandlerFunc, has ntype) {
 		} else {
 			// Node already exists
 			if h != nil {
-				cn.handlers[Methods[method]] = h
+				cn.handlers[methods[method]] = h
 			}
 			break
 		}
@@ -151,7 +164,7 @@ func newNode(pfx string, has ntype, h []HandlerFunc, e edges) (n *node) {
 		edges:    e,
 	}
 	if h == nil {
-		n.handlers = make([]HandlerFunc, len(Methods))
+		n.handlers = make([]HandlerFunc, len(methods))
 	}
 	if e == nil {
 		n.edges = edges{}
@@ -168,7 +181,7 @@ func (r *router) Find(method, path string) (handler HandlerFunc, c *Context, s S
 	for {
 		if search == "" || search == cn.prefix {
 			// Node found
-			h := cn.handlers[Methods[method]]
+			h := cn.handlers[methods[method]]
 			if h != nil {
 				// Handler found
 				handler = h
@@ -252,6 +265,7 @@ func (r *router) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	r.bolt.pool.Put(c)
 }
 
+// Get returns path parameter by name.
 func (ps Params) Get(n string) (v string) {
 	for _, p := range ps {
 		if p.Name == n {