quartzgun/router/router.go

package router

import (
	"context"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
)

type Router struct {
	/* This is the template for error pages */
	Fallback template.Template
	/* Routes are only filled by using the appropriate methods. */
	routes []Route
	/* StaticPaths can be filled from outside when constructing the Router.
	 * key = uri
	 * value = file path
	 */
	StaticPaths map[string]string
}

type Route struct {
	path       *regexp.Regexp
	handlerMap map[string]http.Handler
}

func (self *Router) Get(path string, h http.Handler) {
	self.AddRoute("GET", path, h)
}

func (self *Router) Post(path string, h http.Handler) {
	self.AddRoute("POST", path, h)
}

func (self *Router) Put(path string, h http.Handler) {
	self.AddRoute("PUT", path, h)
}

func (self *Router) Delete(path string, h http.Handler) {
	self.AddRoute("DELETE", path, h)
}

func (self *Router) AddRoute(method string, path string, h http.Handler) {

	exactPath := regexp.MustCompile("^" + path + "$")

	/* If the route already exists, try to add this method to the ServerTask map. */
	for _, r := range self.routes {
		if r.path == exactPath {
			r.handlerMap[method] = h
			return
		}
	}

	/* Otherwise add a new route */
	self.routes = append(self.routes, Route{
		path:       exactPath,
		handlerMap: map[string]http.Handler{method: h},
	})

}

func (self *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	/* Show the 500 error page if we panic */
	defer func() {
		if r := recover(); r != nil {
			log.Println("ERROR:", r)
			self.ErrorPage(w, req, 500, "There was an error on the server.")
		}
	}()

	/* If the request matches any our StaticPaths, try to serve a file. */
	for uri, dir := range self.StaticPaths {
		if req.Method == "GET" && strings.HasPrefix(req.URL.Path, uri) {
			restOfUri := strings.TrimPrefix(req.URL.Path, uri)
			p := path.Join(dir, restOfUri)
			p = path.Clean(p)

			/* If the file exists, try to serve it. */
			info, err := os.Stat(p)
			if err == nil {
				if !info.IsDir() {
					http.ServeFile(w, req, p)
				} else {
					indexFile := path.Join(p, "index.html")
					info2, err2 := os.Stat(indexFile)
					if err2 == nil && !info2.IsDir() {
						http.ServeFile(w, req, indexFile)
					} else {
						self.ErrorPage(w, req, 403, "Access forbidden")
					}
				}
				/* Handle the common errors */
			} else if errors.Is(err, os.ErrNotExist) || errors.Is(err, os.ErrExist) {
				self.ErrorPage(w, req, 404, "The page you requested does not exist!")
			} else if errors.Is(err, os.ErrPermission) || info.IsDir() {
				self.ErrorPage(w, req, 403, "Access forbidden")
				/* If it's some weird error, serve a 500. */
			} else {
				self.ErrorPage(w, req, 500, "Internal server error")
			}

			return
		}
	}

	fmt.Printf("%s: %s\n", req.Method, req.URL.Path)
	/* Otherwise, this is a normal route */
	for _, r := range self.routes {

		/* Pull the params out of the regex;
		 * If the path doesn't match the regex, params will be nil.
		 */
		params := r.Match(req)
		if params == nil {
			continue
		}
		for method, handler := range r.handlerMap {
			if method == req.Method {
				/* Parse the form and add the params to the context */
				req.ParseForm()
				ProcessParams(req, params)
				/* handle the request! */
				handler.ServeHTTP(w, req)
				return
			}
		}
	}
	self.ErrorPage(w, req, 404, "The page you requested does not exist!")
}

/*******************
 * Utility Methods *
 *******************/

func ProcessParams(req *http.Request, params map[string]string) {
	*req = *req.WithContext(context.WithValue(req.Context(), "params", params))
}

func (self *Route) Match(r *http.Request) map[string]string {
	match := self.path.FindStringSubmatch(r.URL.Path)
	if match == nil {
		return nil
	}

	params := map[string]string{}
	groupNames := self.path.SubexpNames()

	for i, group := range match {
		params[groupNames[i]] = group
	}

	return params
}

func (self *Router) ErrorPage(w http.ResponseWriter, req *http.Request, code int, errMsg string) {
	w.WriteHeader(code)
	params := map[string]string{
		"ErrorCode":    strconv.Itoa(code),
		"ErrorMessage": errMsg,
	}
	ProcessParams(req, params)
	self.Fallback.Execute(w, req)
}