ksrv/kserver.go

package ksrv

import (
	"fmt"
	"io"
	"net/http"

	"kumoly.io/lib/klog"
)

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")
	l.SetPrinter(func(w io.Writer, d *klog.Ldata, l2 *klog.Logger) {
		status, ok := d.Fields["Status"].(int)
		msg := d.Message
		if msg != "" {
			msg = l.M(msg, klog.FgRed)
		}
		stack := d.Stack
		if stack != "" {
			stack = l.M(stack, klog.FgRed)
		}
		if !ok {
			status = 500
		}
		fmt.Fprintf(w, "[%s] %s |%s| %15s %s %s %s\n%s",
			l.M("KSRV ", klog.FgHiCyan),
			d.Time.Format("2006/01/02 15:04:05"),
			StatusColor(status, l),
			d.Fields["IP"],
			MethodColor(d.Fields["Method"], l),
			d.Fields["URL"], msg, stack,
		)
	})

	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) *kserver {
	k.nolog = nolog
	return k
}

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")})
	}
}

func MethodColor(method interface{}, l *klog.Logger) string {
	show := fmt.Sprintf(" %-7s", method)
	switch method {
	case http.MethodGet:
		return l.M(show, 97, 44)
	case http.MethodPost:
		return l.M(show, 97, 46)
	case http.MethodPut:
		return l.M(show, 90, 43)
	case http.MethodDelete:
		return l.M(show, 97, 41)
	case http.MethodPatch:
		return l.M(show, 97, 42)
	case http.MethodHead:
		return l.M(show, 97, 45)
	case http.MethodOptions:
		return l.M(show, 90, 47)
	default:
		return l.M(show, 0)
	}
}

func StatusColor(code int, l *klog.Logger) string {
	s := fmt.Sprintf(" %d ", code)
	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)
	}
}