Draft: poop

tcp/router.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"io"
 	"net"
+	"strings"
 
 	"go.uber.org/zap"
 )
@@ -14,11 +15,41 @@ type Middleware func(next Handler) Handler
 type Router struct {
 	middlewares []Middleware
 	routes      map[Method]map[string]Handler
+
+	listenAddr string
+	doneCh     chan error
 }
 
 func NewRouter() *Router {
 	return &Router{
 		routes: make(map[Method]map[string]Handler),
+		doneCh: make(chan error),
+	}
+}
+
+type Group struct {
+	pRouter *Router
+	pGroup  *Group
+
+	path        string
+	middlewares []Middleware
+}
+
+func (self *Router) Group(path string) *Group {
+	return &Group{
+		pRouter: self,
+		pGroup:  nil,
+
+		path: path,
+	}
+}
+
+func (self *Group) Group(path string) *Group {
+	return &Group{
+		pRouter: nil,
+		pGroup:  self,
+
+		path: path,
 	}
 }
 
@@ -27,6 +58,11 @@ func (self *Router) Use(middleware Middleware) *Router {
 	return self
 }
 
+func (self *Group) Use(middleware Middleware) *Group {
+	self.middlewares = append(self.middlewares, middleware)
+	return self
+}
+
 func (self *Router) Register(method Method, route string, handler Handler) {
 	_, ok := self.routes[method]
 	if !ok {
@@ -39,26 +75,59 @@ func (self *Router) Register(method Method, route string, handler Handler) {
 	self.routes[method][route] = handler
 }
 
+func (self *Group) Register(method Method, route string, handler Handler) {
+	for _, middleware := range self.middlewares {
+		handler = middleware(handler)
+	}
+
+	if self.pRouter != nil {
+		self.pRouter.Register(method, self.path+route, handler)
+	}
+	if self.pGroup != nil {
+		self.pGroup.Register(method, self.path+route, handler)
+	}
+}
+
 func (self *Router) GET(route string, handler Handler) {
 	self.Register(MethodGET, route, handler)
 }
 
+func (self *Group) GET(route string, handler Handler) {
+	self.Register(MethodGET, route, handler)
+}
+
 func (self *Router) POST(route string, handler Handler) {
 	self.Register(MethodPOST, route, handler)
 }
 
+func (self *Group) POST(route string, handler Handler) {
+	self.Register(MethodPOST, route, handler)
+}
+
 func (self *Router) PUT(route string, handler Handler) {
 	self.Register(MethodPUT, route, handler)
 }
 
+func (self *Group) PUT(route string, handler Handler) {
+	self.Register(MethodPUT, route, handler)
+}
+
 func (self *Router) DELETE(route string, handler Handler) {
 	self.Register(MethodDELETE, route, handler)
 }
 
+func (self *Group) DELETE(route string, handler Handler) {
+	self.Register(MethodDELETE, route, handler)
+}
+
 func (self *Router) SOCKET(route string, handler Handler) {
 	self.Register(MethodSOCKET, route, handler)
 }
 
+func (self *Group) SOCKET(route string, handler Handler) {
+	self.Register(MethodSOCKET, route, handler)
+}
+
 func (self *Router) run(conn net.Conn, req *Request) {
 	handler, ok := self.routes[req.Method][req.Route]
 	if !ok {
@@ -94,42 +163,64 @@ func (self *Router) run(conn net.Conn, req *Request) {
 	}
 }
 
-func (self *Router) Listen(addr string) error {
-	listener, err := net.Listen("tcp", addr)
-	if err != nil {
-		return err
-	}
+func (self *Router) listen(listener net.Listener) {
 	defer listener.Close()
-
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
-			return err
+			self.doneCh <- err
+			return
 		}
 
 		rawHeader, err := readFrame(conn)
 		if err != nil {
 			if err == io.EOF {
-				return nil
+				self.doneCh <- nil
+				return
 			}
-			return err
+			self.doneCh <- err
+			return
 		}
 
 		body, err := readFrame(conn)
 		if err != nil {
 			if err == io.EOF {
-				return nil
+				self.doneCh <- nil
+				return
 			}
-			return err
+			self.doneCh <- err
+			return
 		}
 
 		var header RequestHeader
 		if err := json.Unmarshal(rawHeader, &header); err != nil {
-			return err
+			self.doneCh <- err
+			return
 		}
 		req := NewRequest(conn, header, body)
 		go self.run(conn, req)
 	}
+}
 
+func (self *Router) WaitFor() error {
+	return <-self.doneCh
+}
+
+func (self *Router) ListenAddr() string {
+	conn, _ := net.Dial("udp", "8.8.8.8:80")
+	addr := conn.LocalAddr().(*net.UDPAddr).IP.String()
+	conn.Close()
+	ss2 := strings.Split(self.listenAddr, ":")
+	return addr + ":" + ss2[len(ss2)-1]
+}
+
+func (self *Router) Listen(addr string) error {
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+
+	self.listenAddr = listener.Addr().String()
+	go self.listen(listener)
 	return nil
 }