diff --git a/main.go b/main.go index a058a44..c2c890c 100644 --- a/main.go +++ b/main.go @@ -4,6 +4,10 @@ package main import ( "context" "flag" + "net/http" + "os" + "os/signal" + "syscall" "time" "modelRT/alert" @@ -16,6 +20,7 @@ import ( "modelRT/logger" "modelRT/middleware" "modelRT/pool" + "modelRT/router" realtimedata "modelRT/real-time-data" @@ -64,9 +69,6 @@ func main() { sqlDB.Close() }() - // init logger - logger.InitLoggerInstance(modelRTConfig.LoggerConfig) - // init alert manager _ = alert.InitAlertEventManager() @@ -120,14 +122,34 @@ func main() { // TODO 暂时屏蔽完成 swagger 启动测试 // go realtimedata.RealTimeDataComputer(ctx, nil, []string{}, "") - engine := gin.Default() - engine.Use(limiter.Middleware) + engine := gin.New() + router.RegisterRoutes(engine) + server := http.Server{ + Addr: ":8080", + Handler: engine, + } - // diagram api - engine.GET("/diagram/load", handler.CircuitDiagramLoadHandler) - engine.POST("/diagram/create", handler.CircuitDiagramCreateHandler) - engine.POST("/diagram/update", handler.CircuitDiagramUpdateHandler) - engine.POST("/diagram/delete", handler.CircuitDiagramDeleteHandler) + // creating a System Signal Receiver + done := make(chan os.Signal, 10) + signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-done + if err := server.Shutdown(context.Background()); err != nil { + logger.Error(ctx, "ShutdownServerError", "err", err) + } + }() + + logger.Info(ctx, "Starting ModelRT server...") + err = server.ListenAndServe() + if err != nil { + if err == http.ErrServerClosed { + // the service receives the shutdown signal normally and then closes + logger.Info(ctx, "Server closed under request") + } else { + // abnormal shutdown of service + logger.Error(ctx, "Server closed unexpected", "err", err) + } + } // real time data api engine.GET("/ws/rtdatas", handler.RealTimeDataReceivehandler) diff --git a/middleware/trace.go b/middleware/trace.go index 12b957e..ef2dad9 100644 --- a/middleware/trace.go +++ b/middleware/trace.go @@ -1,6 +1,12 @@ package middleware import ( + "bytes" + "io" + "strings" + "time" + + "modelRT/logger" "modelRT/util" "github.com/gin-gonic/gin" @@ -21,3 +27,71 @@ func StartTrace() gin.HandlerFunc { c.Next() } } + +type bodyLogWriter struct { + gin.ResponseWriter + body *bytes.Buffer +} + +// 包装一下 gin.ResponseWriter,通过这种方式拦截写响应 +// 让gin写响应的时候先写到 bodyLogWriter 再写gin.ResponseWriter , +// 这样利用中间件里输出访问日志时就能拿到响应了 +// https://stackoverflow.com/questions/38501325/how-to-log-response-body-in-gin +func (w bodyLogWriter) Write(b []byte) (int, error) { + w.body.Write(b) + return w.ResponseWriter.Write(b) +} + +func LogAccess() gin.HandlerFunc { + return func(c *gin.Context) { + // 保存body + var reqBody []byte + contentType := c.GetHeader("Content-Type") + // multipart/form-data 文件上传请求, 不在日志里记录body + if !strings.Contains(contentType, "multipart/form-data") { + reqBody, _ = io.ReadAll(c.Request.Body) + c.Request.Body = io.NopCloser(bytes.NewReader(reqBody)) + + // var request map[string]interface{} + // if err := c.ShouldBindBodyWith(&request, binding.JSON); err != nil { + // c.JSON(400, gin.H{"error": err.Error()}) + // return + // } + } + start := time.Now() + blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer} + c.Writer = blw + + accessLog(c, "access_start", time.Since(start), reqBody, nil) + defer func() { + var responseLogging string + if c.Writer.Size() > 10*1024 { // 响应大于10KB 不记录 + responseLogging = "Response data size is too Large to log" + } else { + responseLogging = blw.body.String() + } + accessLog(c, "access_end", time.Since(start), reqBody, responseLogging) + }() + c.Next() + return + } +} + +func accessLog(c *gin.Context, accessType string, dur time.Duration, body []byte, dataOut interface{}) { + req := c.Request + bodyStr := string(body) + query := req.URL.RawQuery + path := req.URL.Path + // TODO: 实现Token认证后再把访问日志里也加上token记录 + // token := c.Request.Header.Get("token") + logger.New(c).Info("AccessLog", + "type", accessType, + "ip", c.ClientIP(), + //"token", token, + "method", req.Method, + "path", path, + "query", query, + "body", bodyStr, + "output", dataOut, + "time(ms)", int64(dur/time.Millisecond)) +} diff --git a/router/diagram.go b/router/diagram.go new file mode 100644 index 0000000..6bf9451 --- /dev/null +++ b/router/diagram.go @@ -0,0 +1,17 @@ +package router + +import ( + "modelRT/handler" + + "github.com/gin-gonic/gin" +) + +// RegisterRoutes define func of register diagram routes +func registerDiagramRoutes(rg *gin.RouterGroup) { + g := rg.Group("/diagram/") + // TODO add diagram middleware + g.GET("load", handler.CircuitDiagramLoadHandler) + g.POST("create", handler.CircuitDiagramCreateHandler) + g.POST("update", handler.CircuitDiagramUpdateHandler) + g.POST("delete", handler.CircuitDiagramDeleteHandler) +} diff --git a/router/router.go b/router/router.go new file mode 100644 index 0000000..3e9b3dc --- /dev/null +++ b/router/router.go @@ -0,0 +1,22 @@ +package router + +import ( + "time" + + "modelRT/middleware" + + "github.com/gin-gonic/gin" +) + +var limiter *middleware.Limiter + +func init() { + limiter = middleware.NewLimiter(10, 1*time.Minute) // 设置限流器,允许每分钟最多请求10次 +} + +func RegisterRoutes(engine *gin.Engine) { + // use global middlewares + engine.Use(middleware.StartTrace(), limiter.Middleware) + routeGroup := engine.Group("") + registerDiagramRoutes(routeGroup) +}