context: implement handler chain

framework/context.go

@@ -19,7 +19,10 @@ type Context struct {
 	ctx            context.Context
 	request        *http.Request
 	responseWriter http.ResponseWriter
-	handler        ControllerHandler
+
+	handlers []ControllerHandler
+	// current handler index
+	index int
 
 	hasTimeout bool
 	writerMux  *sync.Mutex
@@ -32,6 +35,7 @@ func NewContext(w http.ResponseWriter, r *http.Request) *Context {
 		request:        r,
 		responseWriter: w,
 		writerMux:      &sync.Mutex{},
+		index:          -1, // will be set to 0 when at the beginning
 	}
 }
 
@@ -97,6 +101,25 @@ func (ctx *Context) Value(key any) any {
 	return ctx.BaseContext().Value(key)
 }
 
+// Next runs the next function in the function chain
+func (ctx *Context) Next() error {
+	ctx.index++
+	if ctx.index >= len(ctx.handlers) {
+		// This is the end of the chain
+		return nil
+	}
+	// Run this handler
+	if err := ctx.handlers[ctx.index](ctx); err != nil {
+		return err
+	}
+	return nil
+}
+
+// SetHandlers sets handlers for context
+func (ctx *Context) SetHandlers(handlers []ControllerHandler) {
+	ctx.handlers = handlers
+}
+
 // }}}
 // {{{ Implements request functions
 

framework/core.go

@@ -30,37 +30,37 @@ func NewCore() *Core {
 // Get is a simple get router
 func (c *Core) Get(url string, handler ControllerHandler) {
 	upperUrl := strings.ToUpper(url)
-    if err := c.router["GET"].AddRouter(upperUrl, handler); err != nil{
-        log.Println(err)
-    }
+	if err := c.router["GET"].AddRouter(upperUrl, handler); err != nil {
+		log.Println(err)
+	}
 }
 
 // Post is a simple post router
 func (c *Core) Post(url string, handler ControllerHandler) {
 	upperUrl := strings.ToUpper(url)
-    if err := c.router["POST"].AddRouter(upperUrl, handler); err != nil{
-        log.Println(err)
-    }
+	if err := c.router["POST"].AddRouter(upperUrl, handler); err != nil {
+		log.Println(err)
+	}
 }
 
 // Put is a simple put router
 func (c *Core) Put(url string, handler ControllerHandler) {
 	upperUrl := strings.ToUpper(url)
-    if err := c.router["PUT"].AddRouter(upperUrl, handler); err != nil{
-        log.Println(err)
-    }
+	if err := c.router["PUT"].AddRouter(upperUrl, handler); err != nil {
+		log.Println(err)
+	}
 }
 
 // Delete is a simple delete router
 func (c *Core) Delete(url string, handler ControllerHandler) {
 	upperUrl := strings.ToUpper(url)
-    if err := c.router["DELETE"].AddRouter(upperUrl, handler); err != nil{
-        log.Println(err)
-    }
+	if err := c.router["DELETE"].AddRouter(upperUrl, handler); err != nil {
+		log.Println(err)
+	}
 }
 
 // FindRouteByRequest finds route using the request
-func (c *Core) FindRouteByRequest(r *http.Request) ControllerHandler {
+func (c *Core) FindRouteByRequest(r *http.Request) []ControllerHandler {
 	upperUri := strings.ToUpper(r.URL.Path)
 	upperMethod := strings.ToUpper(r.Method)
 
@@ -70,13 +70,13 @@ func (c *Core) FindRouteByRequest(r *http.Request) ControllerHandler {
 		return nil
 	}
 
-	controller := mapper.FindRoute(upperUri)
-	if controller == nil {
+	controllers := mapper.FindRoute(upperUri)
+	if controllers == nil {
 		log.Printf("URI %q is not recognized\n", r.URL.Path)
 		return nil
 	}
 
-	return controller
+	return controllers
 }
 
 func (c *Core) Group(prefix string) IGroup {
@@ -92,14 +92,15 @@ func (c *Core) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	ctx := NewContext(w, r)
 
-	router := c.FindRouteByRequest(r)
-	if router == nil {
+	handlers := c.FindRouteByRequest(r)
+	if handlers == nil {
 		ctx.WriteJSON(http.StatusNotFound, "Request not found")
 		return
 	}
 
-	err := router(ctx)
-	if err != nil {
+	ctx.SetHandlers(handlers)
+
+	if err := ctx.Next(); err != nil {
 		ctx.WriteJSON(http.StatusInternalServerError, "Internal error")
 		return
 	}

framework/middleware/timeout.go

@@ -0,0 +1,48 @@
+package middleware
+
+import (
+	"context"
+	"log"
+	"net/http"
+	"time"
+
+	"git.vinchent.xyz/vinchent/go-web/framework"
+)
+
+func Timeout(d time.Duration) framework.ControllerHandler {
+	return func(c *framework.Context) error {
+		finish := make(chan struct{}, 1)
+		panicChan := make(chan interface{}, 1)
+
+		durationCtx, cancel := context.WithTimeout(c.BaseContext(), d)
+		defer cancel()
+
+		go func() {
+			// Handle panic
+			defer func() {
+				if p := recover(); p != nil {
+					panicChan <- p
+				}
+			}()
+
+			// Run the next middleware or the business logic
+			c.Next()
+
+			finish <- struct{}{}
+		}()
+
+		select {
+		case p := <-panicChan:
+			// panic
+			log.Println(p)
+			c.GetResponseWriter().WriteHeader(http.StatusInternalServerError)
+		case <-finish:
+			// finish normally
+			log.Println("finish")
+		case <-durationCtx.Done():
+			c.SetHasTimeout()
+			c.GetResponseWriter().Write([]byte("time out"))
+		}
+		return nil
+	}
+}

framework/middleware/timeout_test.go

@@ -0,0 +1,25 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"git.vinchent.xyz/vinchent/go-web/framework"
+)
+
+func TestTimeout(t *testing.T) {
+	t.Run("Test timeout handler", func(t *testing.T) {
+		timeoutHandler := Timeout(1 * time.Millisecond)
+
+		request := httptest.NewRequest(http.MethodGet, "/", nil)
+		response := httptest.NewRecorder()
+		c := framework.NewContext(response, request)
+
+		err := timeoutHandler(c)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+}

framework/trie.go

@@ -14,10 +14,10 @@ func NewTrie() *Trie {
 	return &Trie{root: newNode("")}
 }
 
-func (t *Trie) FindRoute(uri string) ControllerHandler {
+func (t *Trie) FindRoute(uri string) []ControllerHandler {
 	uri = strings.TrimPrefix(uri, "/")
 	if uri == "" {
-		return t.root.handler
+		return t.root.handlers
 	}
 
 	found := t.root.findRoute(uri)
@@ -25,14 +25,14 @@ func (t *Trie) FindRoute(uri string) ControllerHandler {
 		return nil
 	}
 
-	return found.handler
+	return found.handlers
 }
 
 func (t *Trie) AddRouter(uri string, handler ControllerHandler) error {
 	uri = strings.TrimPrefix(uri, "/")
 	if uri == "" {
 		t.root.isLast = true
-		t.root.handler = handler
+		t.root.handlers = append(t.root.handlers, handler)
 		return nil
 	}
 
@@ -54,7 +54,7 @@ func (t *Trie) AddRouter(uri string, handler ControllerHandler) error {
 type node struct {
 	isLast   bool
 	segment  string
-	handler  ControllerHandler
+	handlers []ControllerHandler
 	children []*node
 }
 
@@ -125,7 +125,7 @@ func (n *node) addRoute(uri string, handler ControllerHandler) error {
 				} else {
 					// otherwise, set the child
 					child.isLast = true
-					child.handler = handler
+					child.handlers = append(child.handlers, handler)
 					return nil
 				}
 			}
@@ -138,7 +138,7 @@ func (n *node) addRoute(uri string, handler ControllerHandler) error {
 	new := newNode(splitted[0])
 	if isLast {
 		// this is the end
-		new.handler = handler
+		new.handlers = append(new.handlers, handler)
 		new.isLast = true
 		n.children = append(n.children, new)
 		return nil