Evan Chen 2021-11-04 01:17:58 +08:00
commit 08739831d3
8 changed files with 339 additions and 0 deletions

32
README.md Normal file
View File

@ -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)
}
}
```

13
go.mod Normal file
View File

@ -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
)

8
go.sum Normal file
View File

@ -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=

109
kmux.go Normal file
View File

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

60
net.go Normal file
View File

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

24
server.go Normal file
View File

@ -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()
}

24
test/main.go Normal file
View File

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

69
util.go Normal file
View File

@ -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
}