commit 08739831d327d65ff4d97666ef3b6fac74d62755 Author: Evan Chen Date: Thu Nov 4 01:17:58 2021 +0800 update diff --git a/README.md b/README.md new file mode 100644 index 0000000..ec09814 --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# kmux + +A extended gorilla mux with logging and panic recovery + +## Examples + +```go +package main + +import ( + "errors" + "fmt" + "net/http" + + "kumoly.io/core/log" + "kumoly.io/lib/kmux" +) + +func main() { + log.PROD = false + mux := kmux.NewRouter() + mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) { rw.Write([]byte("ok")) }) + mux.HandleFunc("/err", func(rw http.ResponseWriter, r *http.Request) { kmux.Abort(rw, errors.New("small err")) }) + mux.HandleFunc("/panic", func(rw http.ResponseWriter, r *http.Request) { panic(500) }) + mux.HandleFunc("/out", func(rw http.ResponseWriter, r *http.Request) { arr := []int{0, 1}; fmt.Print(arr[9]) }) + log.Info("start") + err := mux.Server().Listen("0.0.0.0:8081").Serve() + if err != nil { + panic(err) + } +} +``` \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..ff8adc7 --- /dev/null +++ b/go.mod @@ -0,0 +1,13 @@ +module kumoly.io/lib/kmux + +go 1.17 + +require ( + github.com/gorilla/mux v1.8.0 + kumoly.io/core/log v0.1.7 +) + +require ( + github.com/mattn/go-isatty v0.0.14 // indirect + golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..68ccfa6 --- /dev/null +++ b/go.sum @@ -0,0 +1,8 @@ +github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI= +github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= +github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c h1:F1jZWGFhYfh0Ci55sIpILtKKK8p3i2/krTr0H1rg74I= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +kumoly.io/core/log v0.1.7 h1:eAOitZUYyeivFlqcdNsITZT79vVXCzRckuZozqNwx44= +kumoly.io/core/log v0.1.7/go.mod h1:DrKjL3xGEEJ7OC+NHuuv0JSMV7ZqDSw2GiSFFWxilT4= diff --git a/kmux.go b/kmux.go new file mode 100644 index 0000000..aebef27 --- /dev/null +++ b/kmux.go @@ -0,0 +1,109 @@ +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}) + } +} diff --git a/net.go b/net.go new file mode 100644 index 0000000..91b0135 --- /dev/null +++ b/net.go @@ -0,0 +1,60 @@ +package kmux + +import ( + "net/http" + "strconv" +) + +type responseWriter struct { + http.ResponseWriter + StatueCode int + err string +} + +func (w *responseWriter) WriteHeader(statusCode int) { + if w.StatueCode != 0 { + return + } + w.StatueCode = statusCode + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *responseWriter) Write(body []byte) (int, error) { + if w.StatueCode >= 500 { + w.err = string(body) + } + if w.StatueCode == 0 { + w.WriteHeader(200) + } + return w.ResponseWriter.Write(body) +} + +// Response shorthand for set status code and write body +func Response(w http.ResponseWriter, status int, body []byte) (int, error) { + w.WriteHeader(status) + return w.Write(body) +} + +// Abort shorthand for aborting with error, strings and status code could also be passed +func Abort(w http.ResponseWriter, errs ...interface{}) (int, error) { + code := 500 + msg := []byte{} + for _, err := range errs { + switch v := err.(type) { + case int: + if v >= 100 || v < 600 { + code = v + } + case string: + msg = []byte(v) + case error: + msg = []byte(v.Error()) + } + } + if len(msg) == 0 { + msg = []byte(strconv.Itoa(code)) + } + + w.WriteHeader(code) + return w.Write(msg) +} diff --git a/server.go b/server.go new file mode 100644 index 0000000..8fcc8ea --- /dev/null +++ b/server.go @@ -0,0 +1,24 @@ +package kmux + +import ( + "net/http" +) + +type kserver struct { + http.Server +} + +// Listen to addr +func (s *kserver) Listen(addr string) *kserver { + s.Addr = addr + return s +} + +// Serve alias to ListenAndServe +func (s *kserver) Serve() error { + return s.ListenAndServe() +} + +func (s *kserver) ListenAndServe() error { + return s.Server.ListenAndServe() +} diff --git a/test/main.go b/test/main.go new file mode 100644 index 0000000..f36da1b --- /dev/null +++ b/test/main.go @@ -0,0 +1,24 @@ +package main + +import ( + "errors" + "fmt" + "net/http" + + "kumoly.io/core/log" + "kumoly.io/lib/kmux" +) + +func main() { + log.PROD = false + mux := kmux.NewRouter() + mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) { rw.Write([]byte("ok")) }) + mux.HandleFunc("/err", func(rw http.ResponseWriter, r *http.Request) { kmux.Abort(rw, errors.New("small err")) }) + mux.HandleFunc("/panic", func(rw http.ResponseWriter, r *http.Request) { panic(500) }) + mux.HandleFunc("/out", func(rw http.ResponseWriter, r *http.Request) { arr := []int{0, 1}; fmt.Print(arr[9]) }) + log.Info("start") + err := mux.Server().Listen("0.0.0.0:8081").Serve() + if err != nil { + panic(err) + } +} diff --git a/util.go b/util.go new file mode 100644 index 0000000..d35e979 --- /dev/null +++ b/util.go @@ -0,0 +1,69 @@ +package kmux + +import ( + "net" + "net/http" + "strings" +) + +// MatchIPGlob match ip to glob pattern, ex. * 192.168.* 192.* 192.168.51.*2* +func MatchIPGlob(ip, pattern string) bool { + parts := strings.Split(pattern, ".") + seg := strings.Split(ip, ".") + for i, part := range parts { + + // normalize pattern to 3 digits + switch len(part) { + case 1: + if part == "*" { + part = "***" + } else { + part = "00" + part + } + case 2: + if strings.HasPrefix(part, "*") { + part = "*" + part + } else if strings.HasSuffix(part, "*") { + part = part + "*" + } else { + part = "0" + part + } + } + + // normalize ip to 3 digits + switch len(seg[i]) { + case 1: + seg[i] = "00" + seg[i] + case 2: + seg[i] = "0" + seg[i] + } + + for j := range part { + if string(part[j]) == "*" { + continue + } + if part[j] != seg[i][j] { + return false + } + } + } + return true +} + +// GetIP gets the real ip (could still be tricked by proxy) +func GetIP(r *http.Request) string { + ip := r.Header.Get("X-Real-Ip") + if ip == "" { + ips := r.Header.Get("X-Forwarded-For") + ipArr := strings.Split(ips, ",") + ip = strings.Trim(ipArr[len(ipArr)-1], " ") + } + if ip == "" { + var err error + ip, _, err = net.SplitHostPort(r.RemoteAddr) + if err != nil { + ip = r.RemoteAddr + } + } + return ip +}