一个简单的golang http.ServeMux实现restful

用来了解http.ServeMux的大概工作过程和实现方法,代码如下:

// go 1.7

package main

import (
	"io"
	"log"
	"net/http"
)

type restful struct {
	m map[string]func(w http.ResponseWriter, r *http.Request)
}

func newRestful() *restful {
	return &restful{make(map[string]func(w http.ResponseWriter, r *http.Request))}
}

func (mux *restful) get(pattern string, callback func(w http.ResponseWriter, r *http.Request)) {
	mux.m["GET"+pattern] = callback
}

func (mux *restful) post(pattern string, callback func(w http.ResponseWriter, r *http.Request)) {
	mux.m["POST"+pattern] = callback
}

func (mux *restful) notFound(callback func(w http.ResponseWriter, r *http.Request)) {
	mux.m["notFound"] = callback
}

func (mux *restful) before(callback func(w http.ResponseWriter, r *http.Request)) {
	mux.m["before"] = callback
}

func (mux *restful) after(callback func(w http.ResponseWriter, r *http.Request)) {
	mux.m["after"] = callback
}

func (mux *restful) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	method := r.Method + r.URL.Path
	if _, k := mux.m[method]; k {
		if mux.m["before"] != nil {
			mux.m["before"](w, r)
		}
		mux.m[method](w, r)
	} else {
		if mux.m["notFound"] != nil {
			mux.m["notFound"](w, r)
		} else {
			w.WriteHeader(404)
		}
		return
	}
	if mux.m["after"] != nil {
		mux.m["after"](w, r)
	}
}

func notFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	return
}

func before(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "before runing...\n")
}

func after(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "after runing...\n")
}

func test(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "test running....\n")
}

func index(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "index running....\n")
}

func main() {
	mux := newRestful()
	mux.get("/", index)
	mux.get("/test", test)
	mux.before(before)
	mux.after(after)
	mux.notFound(notFound)

	err := http.ListenAndServe(":8091", mux)
	if err != nil {
		log.Fatal(err)
	}
}