115 lines
2.1 KiB
Go
115 lines
2.1 KiB
Go
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
|
|
}
|