ksrv/kserver.go

package ksrv

import (
	"net/http"

	"kumoly.io/lib/klog"
)

const (
	DEFAULT_INFO = `[{{"KSRV"|cyan}}] {{Time}} {{with .Fields}}|{{printf " %3d " .Status|statcol .Status}}| ` +
		`{{printf "%15s" .IP}} {{printf " %-7s " .Method|methcol .Method}} {{.URL}}{{"\n"}}{{end}}`
	DEFAULT_ERR = `[{{"KSRV"|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}}`
)

func (k *kserver) SetLogger(l *klog.Logger) {
	k.l = l
}

func (k *kserver) GetLogger() *klog.Logger {
	return k.l
}

// New returns a extended http.Server
func New() *kserver {
	l := klog.Sub("KSRV")
	tmpl := klog.NewLogFormater()
	tmpl.InfoTmplStr = DEFAULT_INFO
	tmpl.ErrTmplStr = DEFAULT_ERR

	fMap := klog.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 := &kserver{
		l:     l,
		nolog: func(r *http.Request) bool { return false },
	}
	return k
}

func (k *kserver) 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 *kserver) SetNoLogCondition(nolog func(r *http.Request) bool) {
	k.nolog = nolog
}

func (k *kserver) catch(rw *responseWriter, r *http.Request) {
	ex := recover()
	if ex != nil {
		Abort(rw, ex)
		k.l.ErrorF(klog.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path, "Agent": r.Header.Get("User-Agent")}, ex)
	} else if rw.StatueCode >= 500 {
		k.l.ErrorF(klog.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path, "Agent": r.Header.Get("User-Agent")}, rw.err)
	} else if k.nolog == nil || !k.nolog(r) {
		k.l.InfoF(klog.H{"Status": rw.StatueCode, "IP": GetIP(r), "Method": r.Method, "URL": r.URL.Path, "Agent": r.Header.Get("User-Agent")})
	}
}