package server import ( "context" "net/http" "strings" "github.com/samber/lo" ) type RequestCtxKey string // 可以嵌套的Rest http server mux type RestMux struct { Path string imux *http.ServeMux rmuxPaths []string middlewares []Middleware } func (mux *RestMux) Use(m Middleware) { mux.middlewares = append(mux.middlewares, m) } func (mux *RestMux) ServeHTTP(w http.ResponseWriter, r *http.Request) { canContinue := false c := r.Context() data := map[string]interface{}{} cn := context.WithValue(c, RequestCtxKey("data"), data) newRequest := r.WithContext(cn) if len(mux.middlewares) > 0 { for _, m := range mux.middlewares { canContinue = false m(w, newRequest, func() { canContinue = true }) if !canContinue { return } } } _, has := lo.Find(mux.rmuxPaths, func(s string) bool { return strings.HasPrefix(newRequest.URL.Path, s) }) if has { mux.imux.ServeHTTP(w, newRequest) return } newRequest.URL.Path = "/" + strings.ToLower(newRequest.Method) + newRequest.URL.Path newRequest.RequestURI = "/" + strings.ToLower(newRequest.Method) + newRequest.RequestURI h, _ := mux.imux.Handler(newRequest) h.ServeHTTP(w, newRequest) } func (mux *RestMux) HandleFunc(method string, path string, f func(http.ResponseWriter, *http.Request)) { m := path if !strings.HasPrefix(path, "/") { m = "/" + path } mux.imux.HandleFunc("/"+strings.ToLower(method)+m, f) } func (mux *RestMux) Get(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("GET", path, f) } func (mux *RestMux) Post(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("POST", path, f) } func (mux *RestMux) Put(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("PUT", path, f) } func (mux *RestMux) Delete(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("DELETE", path, f) } func (mux *RestMux) Patch(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("PATCH", path, f) } func (mux *RestMux) Head(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("HEAD", path, f) } func (mux *RestMux) Option(path string, f func(http.ResponseWriter, *http.Request)) { mux.HandleFunc("OPTION", path, f) } func (mux *RestMux) HandleMux(nmux *RestMux) { p := nmux.Path if !strings.HasSuffix(p, "/") { p = p + "/" } mux.imux.Handle(p, http.StripPrefix(nmux.Path, nmux)) mux.rmuxPaths = append(mux.rmuxPaths, nmux.Path) } func NewRestMux(path string) *RestMux { return &RestMux{ Path: path, imux: http.NewServeMux(), rmuxPaths: make([]string, 0), } }