diff --git a/main.go b/main.go index 94bb800..b47f55e 100644 --- a/main.go +++ b/main.go @@ -36,6 +36,7 @@ import ( "gorm.io/gorm" "io/ioutil" "log" + "net" "net/http" "os" "os/exec" @@ -111,6 +112,23 @@ func NeverExit(name string, f func()) { // run 是服务的主要逻辑 func (m *wafSystenService) run() { + // 先尝试监听端口,检查是否被占用 + listener, err := net.Listen("tcp", ":"+strconv.Itoa(global.GWAF_LOCAL_SERVER_PORT)) + defer func() { + if listener != nil { + err := listener.Close() + if err != nil { + return + } + } + }() + if err != nil { + errMsg := fmt.Sprintf("管理界面端口 %d 已被占用,请检查并修改配置(conf/config.yml local_port字段)或关闭占用该端口的程序: %s", + global.GWAF_LOCAL_SERVER_PORT, err.Error()) + zlog.Error(errMsg) + panic(errMsg) + return + } // 获取当前执行文件的路径 executablePath, err := os.Executable() if err != nil { diff --git a/wafmangeweb/localserver.go b/wafmangeweb/localserver.go index 6883aa4..4cd903c 100644 --- a/wafmangeweb/localserver.go +++ b/wafmangeweb/localserver.go @@ -8,6 +8,7 @@ import ( "SamWaf/wafmangeweb/static" "context" "errors" + "fmt" "github.com/gin-contrib/pprof" "github.com/gin-gonic/gin" "io" @@ -117,7 +118,9 @@ func (web *WafWebManager) cors() gin.HandlerFunc { c.Next() } } -func (web *WafWebManager) StartLocalServer() { + +// StartLocalServer 启动本地管理服务器 +func (web *WafWebManager) StartLocalServer() error { if global.GWAF_RELEASE == "true" { gin.SetMode(gin.ReleaseMode) gin.DefaultWriter = io.Discard @@ -131,10 +134,16 @@ func (web *WafWebManager) StartLocalServer() { Addr: ":" + strconv.Itoa(global.GWAF_LOCAL_SERVER_PORT), Handler: r, } - if err := web.HttpServer.ListenAndServe(); err != nil && errors.Is(err, http.ErrServerClosed) { - zlog.Error(web.LogName, "use static asset", err.Error()) + + // 正式启动服务 + if err := web.HttpServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + errMsg := fmt.Sprintf("启动管理界面失败: %s", err.Error()) + zlog.Error(web.LogName, errMsg) + return err } + zlog.Info(web.LogName, "本地 port:", global.GWAF_LOCAL_SERVER_PORT) + return nil } /* @@ -146,10 +155,18 @@ func (web *WafWebManager) CloseLocalServer() { // The context is used to inform the server it has 5 seconds to finish // the request it is currently handling ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if err := web.HttpServer.Shutdown(ctx); err != nil { - zlog.Error(web.LogName, "Server forced to shutdown:", err.Error()) + defer func() { + defer cancel() + }() + + if web != nil && web.HttpServer != nil { + if err := web.HttpServer.Shutdown(ctx); err != nil { + zlog.Error(web.LogName, "Server forced to shutdown:", err.Error()) + } + zlog.Info(web.LogName, "local Server exiting") + } else { + zlog.Info("local Server exiting") } - zlog.Info(web.LogName, "local Server exiting") + } diff --git a/wafsafeclear/wafsafeclear.go b/wafsafeclear/wafsafeclear.go index c0cc13f..813a2ce 100644 --- a/wafsafeclear/wafsafeclear.go +++ b/wafsafeclear/wafsafeclear.go @@ -6,48 +6,57 @@ import ( ) func SafeClear() { - - sqlDB, err := global.GWAF_LOCAL_DB.DB() - if err != nil { - zlog.Error("清理异常", err) - } else { - - // 关闭数据库连接 - if err := sqlDB.Close(); err != nil { - zlog.Error("清理异常关闭错误", err) - } - } - sqlDB, err = global.GWAF_LOCAL_LOG_DB.DB() - if err != nil { - zlog.Error("清理log退出异常", err) - } else { - - // 关闭数据库连接 - if err := sqlDB.Close(); err != nil { - zlog.Error("清理log异常关闭错误", err) - } - } - sqlDB, err = global.GWAF_LOCAL_STATS_DB.DB() - if err != nil { - zlog.Error("清理stat退出异常", err) - } else { - - // 关闭数据库连接 - if err := sqlDB.Close(); err != nil { - zlog.Error("清理stat异常关闭错误", err) - } - } - for _, value := range global.GDATA_CURRENT_LOG_DB_MAP { - sqlDB, err = value.DB() + if global.GWAF_LOCAL_DB != nil { + sqlDB, err := global.GWAF_LOCAL_DB.DB() if err != nil { - zlog.Error("清理异常错误存档", err) + zlog.Error("清理异常", err) } else { // 关闭数据库连接 if err := sqlDB.Close(); err != nil { - zlog.Error("清理异常错误存档关闭错误", err) + zlog.Error("清理异常关闭错误", err) } } } + + if global.GWAF_LOCAL_LOG_DB != nil { + sqlDB, err := global.GWAF_LOCAL_LOG_DB.DB() + if err != nil { + zlog.Error("清理log退出异常", err) + } else { + + // 关闭数据库连接 + if err := sqlDB.Close(); err != nil { + zlog.Error("清理log异常关闭错误", err) + } + } + } + if global.GWAF_LOCAL_STATS_DB != nil { + sqlDB, err := global.GWAF_LOCAL_STATS_DB.DB() + if err != nil { + zlog.Error("清理stat退出异常", err) + } else { + + // 关闭数据库连接 + if err := sqlDB.Close(); err != nil { + zlog.Error("清理stat异常关闭错误", err) + } + } + } + if global.GDATA_CURRENT_LOG_DB_MAP != nil { + for _, value := range global.GDATA_CURRENT_LOG_DB_MAP { + sqlDB, err := value.DB() + if err != nil { + zlog.Error("清理异常错误存档", err) + } else { + + // 关闭数据库连接 + if err := sqlDB.Close(); err != nil { + zlog.Error("清理异常错误存档关闭错误", err) + } + } + } + } + zlog.Info("退出清理完成") } diff --git a/waftask/core_scheduler.go b/waftask/core_scheduler.go index d23dd38..a9c183d 100644 --- a/waftask/core_scheduler.go +++ b/waftask/core_scheduler.go @@ -63,7 +63,9 @@ func (ts *TaskScheduler) Start() { ts.Scheduler.StartAsync() } func (ts *TaskScheduler) Stop() { - ts.Scheduler.Stop() + if ts != nil { + ts.Scheduler.Stop() + } } func (ts *TaskScheduler) RunManual(taskMethod string) {