Extend ServeMux to register method handler

By Audren Bouëssel du Bourg

When building a web api with the stdlib, a typical handler may look like this:

mux := http.NewServeMux()
mux.HandleFunc("/ping", handlePing)

func handlePing(w http.ResponseWriter, r *http.Method) {
    switch r.Method {
        case http.MethodGet: 
            fmt.Fprintf(w, "pong")
        default:
            http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
            return
    }
}

It is quite powerful as is and you might not need anything else. It is very clear how methods are handled.

One thing I really appreciate when building a rest api is the ability to register routes based on method provided, such as:

r := newRouter() // a router not implemented yet
r.GET("/ping", handlePing)

func handlePing(w http.ResponseWriter, r *http.Method) {
-    switch r.Method {
-        case http.MethodGet: 
            fmt.Fprintf(w, "pong")
-        default:
-            http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
-            return
-    }
}

Let’s try to reproduce that behavior with http.ServeMux to extend it a little bit!

You can find the github for this article here: https://github.com/audrenbdb/router

First of all we want a router struct, keeping track of our endpoints and their handlers for each methods:

type router struct {
	endpoints []endpoint
}

type endpoint struct {
	pattern string
	handler methodHandler
}

// methodHandler is a map of request method : http.Handler
type methodHandler map[string]http.Handler

The router shall have the ability to register a an endpoint method and its handler.

func (r *router) POST(pattern string, handlerFunc http.HandlerFunc) {
	r.registerEndpoint(http.MethodPost, pattern, handlerFunc)
}

func (r *router) PUT(pattern string, handlerFunc http.HandlerFunc) {
	r.registerEndpoint(http.MethodPut, pattern, handlerFunc)
}

func (r *router) PATCH(pattern string, handlerFunc http.HandlerFunc) {
	r.registerEndpoint(http.MethodPatch, pattern, handlerFunc)
}

func (r *router) DELETE(pattern string, handlerFunc http.HandlerFunc) {
	r.registerEndpoint(http.MethodDelete, pattern, handlerFunc)
}


// registerEndpoint adds an endpoint handler 
// from given method to router.
// 
// an endpoint with same method cannot be registered twice
func (r *router) registerEndpoint(method, pattern string, handler http.Handler) {
	ep := r.findEndpoint(pattern)
	if ep == nil {
		newEndpoint := endpoint{
			pattern: pattern,
			handler: map[string]http.Handler{},
		}
		r.endpoints = append(r.endpoints, newEndpoint)
		ep = &newEndpoint
	}
	_, methodFound := ep.handler[method]
	if methodFound {
		msg := fmt.Sprintf("method %s already registered for pattern %s", method, pattern)
		log.Fatal(msg)
	}
	ep.handler[method] = handler
}

func (r *router) findEndpoint(pattern string) *endpoint {
	for _, ep := range r.endpoints {
		if ep.pattern == pattern {
			return &ep
		}
	}
	return nil
}

Our method handler will dispatch request based on its method, or error with http.StatusMethodNotAllowed if request method is not registered on given endpoint. See below:

func (m methodHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h, ok := m[r.Method]
	if !ok {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	h.ServeHTTP(w, r)
}

Finally, we need to ensure that are router implements aswell the http.Handler interface, so that we can easily use it with http.ListenAndServe(":8080", router).

Remember, in the end, a regular http.ServeMux will be used.

For that we need a little adjustment on our router structure:

type router struct {
+   muxMutex  sync.Mutex
+	mux       *http.ServeMux    
	endpoints []*endpoint
}

Then, we can implement http.Handler. The router will set up mux on the first request.

To prevent race condition, we use a mutex so that mux is only going to be initialized once.

Once initialized, router will use ServeMux of our router to handle request.

func (r *router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if r.mux == nil {
		r.initializeMux()
	}
	r.mux.ServeHTTP(w, req)
}

func (r *router) initializeMux() {
	r.muxMutex.Lock()
	defer r.muxMutex.Unlock()
	if r.mux != nil {
		return
	}
	r.mux = http.NewServeMux()
	for _, ep := range r.endpoints {
		r.mux.Handle(ep.pattern, ep.handler)
	}
}

Usage

func main() {
    r := router.New()
    r.GET("/foo", handleFoo)
    
    http.ListenAndServe(":8080", r)
}


func handleFoo(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "bar")
}

Code: https://github.com/audrenbdb/router

Closing thoughts

Other feature I like from popular routers is the named parameters such as /posts/:id/authors.

When not a requirement, stdlib is perfect.