Feat: done tcp

tcp/router.go

@@ -0,0 +1,114 @@
+package tcp
+
+import (
+	"encoding/json"
+	"io"
+	"net"
+
+	"go.uber.org/zap"
+)
+
+type Handler func(w ResponseWriter, req *Request) error
+type Middleware func(next Handler) Handler
+
+type Router struct {
+	middlewares []Middleware
+	routes      map[Method]map[string]Handler
+}
+
+func NewRouter() *Router {
+	return &Router{
+		routes: make(map[Method]map[string]Handler),
+	}
+}
+
+func (self *Router) Use(middleware Middleware) {
+	self.middlewares = append(self.middlewares, middleware)
+}
+
+func (self *Router) Register(method Method, route string, handler Handler) {
+	_, ok := self.routes[method]
+	if !ok {
+		self.routes[method] = make(map[string]Handler)
+	}
+
+	for _, middleware := range self.middlewares {
+		handler = middleware(handler)
+	}
+	self.routes[method][route] = handler
+}
+
+func (self *Router) run(conn net.Conn, req *Request) {
+	handler, ok := self.routes[req.Method][req.Route]
+	if !ok {
+		zap.L().Warn("route not exist",
+			zap.String("route", req.Route))
+		return
+	}
+
+	w := NewResponseBuilder(conn)
+	err := handler(w, req)
+	if err != nil {
+		zap.L().Error("failed to run handler",
+			zap.Error(err))
+		return
+	}
+
+	res := w.Build(req)
+	header, err := res.Header()
+	if err != nil {
+		zap.L().Error("failed to marshal header",
+			zap.Error(err))
+		return
+	}
+	if err := sendFrame(conn, header); err != nil {
+		zap.L().Error("failed to write header",
+			zap.Error(err))
+		return
+	}
+	if err := sendFrame(conn, res.Body); err != nil {
+		zap.L().Error("failed to write body",
+			zap.Error(err))
+		return
+	}
+}
+
+func (self *Router) Listen(addr string) error {
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+	defer listener.Close()
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			return err
+		}
+
+		rawHeader, err := readFrame(conn)
+		if err != nil {
+			if err == io.EOF {
+				return nil
+			}
+			return err
+		}
+
+		body, err := readFrame(conn)
+		if err != nil {
+			if err == io.EOF {
+				return nil
+			}
+			return err
+		}
+
+		var header RequestHeader
+		if err := json.Unmarshal(rawHeader, &header); err != nil {
+			return err
+		}
+		req := NewRequest(conn, header, body)
+		go self.run(conn, req)
+	}
+
+	return nil
+}