ksrv/kmux.go

package kmux

import (
	"net/http"

	"github.com/gorilla/mux"
	"kumoly.io/core/log"
)

const (
	DEFAULT_INFO = `[{{"KMUX"|cyan}}] {{Time}} {{with .Fields}}|{{printf " %3d " .Status|statcol .Status}}| ` +
		`{{printf "%15s" .IP}} {{printf " %-7s " .Method|methcol .Method}} {{.URL}}{{"\n"}}{{end}}`
	DEFAULT_ERR = `[{{"KMUX"|cyan}}] {{Time}} {{with .Fields}}|{{printf " %3d " .Status|statcol .Status}}| ` +
		`{{printf "%15s" .IP}} {{printf " %-7s " .Method|methcol .Method}} {{.URL}}{{end}} {{.Message|red}}{{"\n"}}` +
		`{{if .Stack}}{{.Stack|redl}}{{end}}`
)

type kmux struct {
	mux.Router
	l *log.Logger
}

func (k *kmux) SetLogger(l *log.Logger) {
	k.l = l
}

func (k *kmux) GetLogger() *log.Logger {
	return k.l
}

func (k *kmux) Server() *kserver {
	return &kserver{
		http.Server{
			Handler: k.middleware(k),
			Addr:    "0.0.0.0:8080",
		},
	}
}

// NewRouter returns a extended gorilla mux
func NewRouter() *kmux {
	l := log.Sub("KMUX")
	tmpl := log.NewLogFormater()
	tmpl.InfoTmplStr = DEFAULT_INFO
	tmpl.ErrTmplStr = DEFAULT_ERR

	fMap := log.DefaultFuncMap()
	fMap["statcol"] = func(code int, s string) string {
		switch {
		case code >= http.StatusOK && code < http.StatusMultipleChoices:
			return l.M(s, 97, 42)
		case code >= http.StatusMultipleChoices && code < http.StatusBadRequest:
			return l.M(s, 90, 47)
		case code >= http.StatusBadRequest && code < http.StatusInternalServerError:
			return l.M(s, 90, 43)
		default:
			return l.M(s, 97, 41)
		}
	}
	fMap["methcol"] = func(method string, s string) string {
		switch method {
		case http.MethodGet:
			return l.M(s, 97, 44)
		case http.MethodPost:
			return l.M(s, 97, 46)
		case http.MethodPut:
			return l.M(s, 90, 43)
		case http.MethodDelete:
			return l.M(s, 97, 41)
		case http.MethodPatch:
			return l.M(s, 97, 42)
		case http.MethodHead:
			return l.M(s, 97, 45)
		case http.MethodOptions:
			return l.M(s, 90, 47)
		default:
			return l.M(s, 0)
		}
	}
	l.SetTmpl(tmpl, fMap)
	err := l.Reload()
	if err != nil {
		panic(err)
	}
	k := &kmux{*mux.NewRouter(), l}
	return k
}

func (k *kmux) middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			rw := &responseWriter{w, 0, ""}
			defer k.catch(rw, r)
			next.ServeHTTP(rw, r)
		},
	)
}

func (k *kmux) catch(rw *responseWriter, r *http.Request) {
	ex := recover()
	if ex != nil {
		Abort(rw, ex)
		k.l.ErrorF(log.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path}, ex)
	} else if rw.StatueCode >= 500 {
		k.l.ErrorF(log.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path}, rw.err)
	} else {
		k.l.InfoF(log.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path})
	}
}