支持正则路由

Author: c-wind

http/server/handler.go

@@ -1,10 +1,12 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/http"
 	"reflect"
+	"regexp"
 	"runtime/debug"
 	"strings"
 	"sync"
@@ -18,12 +20,24 @@ var (
 	server = newHTTPServer()
 )
 
+type uriRegexp struct {
+	path   string
+	keys   []string
+	exp    *regexp.Regexp
+	source reflect.Type
+}
+
 // iface 对外服务接口, path格式：Method/URI
 type iface struct {
 	path   string
 	source reflect.Type
 }
 
+type prefix struct {
+	path    string
+	uriExps []uriRegexp
+}
+
 type httpServer struct {
 	path     map[string]iface
 	prefix   *btree.BTree
@@ -43,8 +57,8 @@ func newHTTPServer() *httpServer {
 // Filter 请求过滤， 如果返回结果为nil,直接返回，不再进行后续处理.
 type Filter func(http.ResponseWriter, *http.Request) *http.Request
 
-func (i *iface) Less(bi btree.Item) bool {
-	return strings.Compare(i.path, bi.(*iface).path) == 1
+func (p *prefix) Less(bi btree.Item) bool {
+	return strings.Compare(p.path, bi.(*prefix).path) == 1
 }
 
 //NameToPath 类名转路径
@@ -74,9 +88,40 @@ func RegisterPrefix(obj interface{}, path string) error {
 	return register(obj, path, true)
 }
 
-//RegisterPath 注册url完全匹配.
-func RegisterPath(obj interface{}, path string) error {
-	return register(obj, path, true)
+var (
+	keysExp *regexp.Regexp
+)
+
+func init() {
+	exp, err := regexp.Compile("{(\\w+)?}")
+	if err != nil {
+		panic(err.Error())
+	}
+	keysExp = exp
+}
+
+func newURIRegexp(path string, source reflect.Type) uriRegexp {
+	ur := uriRegexp{
+		path:   path,
+		source: source,
+	}
+
+	//log.Debugf("path:%v", path)
+	for _, m := range keysExp.FindAllStringSubmatch(path, -1) {
+		//log.Debugf("path:%v, key:%#v", path, m)
+		ur.keys = append(ur.keys, m[1])
+	}
+
+	np := keysExp.ReplaceAllString(path, "(.+)")
+	exp, err := regexp.Compile(np)
+	if err != nil {
+		panic(err.Error())
+	}
+	ur.exp = exp
+
+	//	log.Debugf("uriRegexp:%#v", ur)
+
+	return ur
 }
 
 func register(obj interface{}, path string, isPrefix bool) error {
@@ -110,21 +155,37 @@ func register(obj interface{}, path string, isPrefix bool) error {
 			continue
 		}
 
-		ifc := iface{
-			path:   fmt.Sprintf("%v%v", method, path),
-			source: rt.Elem(),
-		}
-
 		//前缀匹配
 		if isPrefix {
-			if server.prefix.Has(&ifc) {
-				panic(fmt.Sprintf("exist url:%v %v", method, path))
+			p := prefix{
+				path: fmt.Sprintf("%v%v", method, path),
+			}
+
+			if idx := strings.Index(path, "{"); idx > 0 {
+				p.path = fmt.Sprintf("%v%v", method, path[:idx])
+			}
+
+			//	log.Debugf("path:%v", p.path)
+
+			if server.prefix.Has(&p) {
+				p = *(server.prefix.Get(&p).(*prefix))
 			}
-			server.prefix.ReplaceOrInsert(&ifc)
-			log.Infof("add prefix %v %v %v", method, path, rt)
+
+			exp := newURIRegexp(path, rt.Elem())
+
+			p.uriExps = append(p.uriExps, exp)
+
+			server.prefix.ReplaceOrInsert(&p)
+
+			log.Infof("add prefix %v %v %v %v", method, exp.path, exp.keys, rt)
+
 			continue
 		}
 
+		ifc := iface{
+			path:   fmt.Sprintf("%v%v", method, path),
+			source: rt.Elem(),
+		}
 		//全路径匹配
 		if _, ok := server.path[ifc.path]; ok {
 			panic(fmt.Sprintf("exist url:%v %v", method, path))
@@ -144,27 +205,43 @@ func AddFilter(filter Filter) {
 	server.filter = filter
 }
 
-func getInterface(w http.ResponseWriter, r *http.Request) (i iface, ok bool) {
+func parseRequestValues(path string, ur uriRegexp) context.Context {
+	ctx := context.Background()
+	for i, v := range ur.exp.FindAllStringSubmatch(path, -1) {
+		//log.Debugf("i:%v, v:%#v, keys:%#v", i, v, ur.keys)
+		ctx = context.WithValue(ctx, ur.keys[i], v[1])
+	}
+	return ctx
+}
+
+func getInterface(method, path string) (reflect.Type, context.Context) {
 	server.RLock()
 	defer server.RUnlock()
 
-	path := r.Method + r.URL.Path
+	path = method + path
 
-	if i, ok = server.path[path]; ok {
+	if i, ok := server.path[path]; ok {
 		//log.Debugf("find path:%v", path)
-		return
+		return i.source, nil
 	}
 
+	var p prefix
+	var ok bool
 	//如果完全匹配没找到，再找前缀的
-	server.prefix.AscendGreaterOrEqual(&iface{path: path}, func(item btree.Item) bool {
-		i = *(item.(*iface))
-		ok = strings.HasPrefix(path, i.path)
-		// log.Debugf("path:%v, ipath:%v, ok:%v", path, i.path, ok)
+	server.prefix.AscendGreaterOrEqual(&prefix{path: path}, func(item btree.Item) bool {
+		p = *(item.(*prefix))
+		ok = strings.HasPrefix(path, p.path)
+		//log.Debugf("path:%v, prefix:%v, ok:%v", path, p.path, ok)
 		return !ok
 	})
 
-	//	log.Debugf("find prefix:%v, ok:%v", path, ok)
-	return
+	for _, ue := range p.uriExps {
+		if ue.exp.MatchString(path) {
+			//	log.Debugf("uri exp:%v, path:%v", ue, path)
+			return ue.source, parseRequestValues(path, ue)
+		}
+	}
+	return nil, nil
 }
 
 //ServeHTTP 真正对外服务接口
@@ -183,16 +260,20 @@ func (s *httpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	i, ok := getInterface(w, r)
-	if !ok {
+	tp, ctx := getInterface(r.Method, r.URL.Path)
+	if tp == nil {
 		log.Errorf("%v %v %v not found.", r.RemoteAddr, r.Method, r.URL)
 		SendResponse(w, http.StatusNotFound, "invalid request")
 		return
 	}
 
-	log.Debugf("%v %v %v %v", r.RemoteAddr, r.Method, r.URL, i.source)
+	if ctx != nil {
+		nr = nr.WithContext(ctx)
+	}
+
+	log.Debugf("%v %v %v %v", r.RemoteAddr, r.Method, r.URL, tp)
 
-	callback := reflect.New(i.source).MethodByName(r.Method).Interface().(func(http.ResponseWriter, *http.Request))
+	callback := reflect.New(tp).MethodByName(r.Method).Interface().(func(http.ResponseWriter, *http.Request))
 	callback(w, nr)
 
 	return

main.go

@@ -11,6 +11,13 @@ import (
 	_ "github.com/dearcode/crab/server"
 )
 
+type regexpTest struct {
+}
+
+func (rt *regexpTest) GET(w http.ResponseWriter, req *http.Request) {
+	fmt.Fprintf(w, "%#v", req.Context())
+}
+
 type index struct {
 	r *http.Request
 }
@@ -34,6 +41,7 @@ func main() {
 	flag.Parse()
 
 	server.Register(&index{})
+	server.RegisterPrefix(&regexpTest{}, "/regexp/{user}/test/")
 
 	go func() {
 		for i := 0; i < 5; i++ {

server/init.go

@@ -8,7 +8,7 @@ func init() {
 	server.RegisterPrefix(&staticServer{}, "/")
 	server.RegisterPrefix(&debugServer{}, "/debug/")
 
-	server.RegisterPath(&testServer{}, "/test/")
+	server.RegisterPrefix(&testServer{}, "/test/")
 
 	server.Register(&user{})
 }