From 3598048e14608cb4c9093c1cb51e667a77854b86 Mon Sep 17 00:00:00 2001 From: samwaf Date: Mon, 19 May 2025 09:45:52 +0800 Subject: [PATCH] fix:create index #IC8E8S --- enums/db_enum.go | 8 ++++++++ main.go | 22 +++++++++++++++++---- utils/notifyhelper.go | 2 +- wafdb/localdb.go | 39 +++++++++++++++++++++++++++++++++---- waftask/task_createindex.go | 19 ++++++++++++++++++ waftask/task_health.go | 2 +- 6 files changed, 82 insertions(+), 10 deletions(-) create mode 100644 enums/db_enum.go diff --git a/enums/db_enum.go b/enums/db_enum.go new file mode 100644 index 0000000..c8352ba --- /dev/null +++ b/enums/db_enum.go @@ -0,0 +1,8 @@ +package enums + +const ( + DB_MAIN = "local.db" //主库数据库 + DB_LOG = "local_log.db" //日志数据库 + DB_STATS = "local_stats.db" //统计数据库 + +) diff --git a/main.go b/main.go index bc66b3c..fa84824 100644 --- a/main.go +++ b/main.go @@ -237,10 +237,24 @@ func (m *wafSystenService) run() { } //初始化本地数据库 - wafdb.InitCoreDb("") - wafdb.InitLogDb("") - wafdb.InitStatsDb("") - + isNewMainDb, err := wafdb.InitCoreDb("") + if err == nil { + if isNewMainDb { + waftask.TaskCreateIndexByDbName(enums.DB_MAIN) + } + } + isNewLogDb, err := wafdb.InitLogDb("") + if err == nil { + if isNewLogDb { + waftask.TaskCreateIndexByDbName(enums.DB_LOG) + } + } + isNewStatsDb, err := wafdb.InitStatsDb("") + if err == nil { + if isNewStatsDb { + waftask.TaskCreateIndexByDbName(enums.DB_STATS) + } + } wafssl.InitDnsProviderEnvInfo() //初始化队列引擎 wafqueue.InitDequeEngine() diff --git a/utils/notifyhelper.go b/utils/notifyhelper.go index abbb0b7..6ff9cba 100644 --- a/utils/notifyhelper.go +++ b/utils/notifyhelper.go @@ -42,7 +42,7 @@ func (receiver *NotifyHelper) SendRuleInfo(ruleMessageInfo innerbean.RuleMessage func (receiver *NotifyHelper) SendNoticeInfo(operatorMessageInfo innerbean.OperatorMessageInfo) { if global.GCACHE_WECHAT_ACCESS == "" { - zlog.Error("未初始化wechat") + zlog.Debug("未初始化wechat") return } operatorMessageInfo.Server = global.GWAF_CUSTOM_SERVER_NAME diff --git a/wafdb/localdb.go b/wafdb/localdb.go index 8a5cd85..aeeca9a 100644 --- a/wafdb/localdb.go +++ b/wafdb/localdb.go @@ -26,7 +26,7 @@ import ( "gorm.io/gorm" ) -func InitCoreDb(currentDir string) { +func InitCoreDb(currentDir string) (bool, error) { if currentDir == "" { currentDir = utils.GetCurrentDir() } @@ -34,11 +34,17 @@ func InitCoreDb(currentDir string) { if _, err := os.Stat(currentDir + "/data/"); os.IsNotExist(err) { if err := os.MkdirAll(currentDir+"/data/", os.ModePerm); err != nil { zlog.Error("创建data目录失败:", err) - return + return false, err } } if global.GWAF_LOCAL_DB == nil { path := currentDir + "/data/local.db" + // 检查数据库文件是否存在 + isNewDb := false + if _, err := os.Stat(path); os.IsNotExist(err) { + isNewDb = true + zlog.Debug("本地主数据库文件不存在，将创建新数据库") + } // 检查文件是否存在 if _, err := os.Stat(path); !os.IsNotExist(err) { // 文件存在的逻辑，使用工具函数进行备份 @@ -153,15 +159,26 @@ func InitCoreDb(currentDir string) { db.Where("user_code = ? and rule_status = 999", global.GWAF_USER_CODE).Delete(model.Rules{}) pathCoreSql(db) + return isNewDb, nil + } else { + return false, nil } } -func InitLogDb(currentDir string) { +func InitLogDb(currentDir string) (bool, error) { if currentDir == "" { currentDir = utils.GetCurrentDir() } if global.GWAF_LOCAL_LOG_DB == nil { path := currentDir + "/data/local_log.db" + + // 检查数据库文件是否存在 + isNewDb := false + if _, err := os.Stat(path); os.IsNotExist(err) { + isNewDb = true + zlog.Debug("本地日志数据库文件不存在，将创建新数据库") + } + key := url.QueryEscape(global.GWAF_PWD_LOGDB) dns := fmt.Sprintf("%s?_db_key=%s", path, key) db, err := gorm.Open(sqlite.Open(dns), &gorm.Config{}) @@ -216,6 +233,10 @@ func InitLogDb(currentDir string) { } global.GWAF_LOCAL_DB.Create(sharDbBean) } + + return isNewDb, nil + } else { + return false, nil } } @@ -263,12 +284,18 @@ func InitManaulLogDb(currentDir string, custFileName string) { } } -func InitStatsDb(currentDir string) { +func InitStatsDb(currentDir string) (bool, error) { if currentDir == "" { currentDir = utils.GetCurrentDir() } if global.GWAF_LOCAL_STATS_DB == nil { path := currentDir + "/data/local_stats.db" + // 检查数据库文件是否存在 + isNewDb := false + if _, err := os.Stat(path); os.IsNotExist(err) { + isNewDb = true + zlog.Debug("本地统计数据库文件不存在，将创建新数据库") + } key := url.QueryEscape(global.GWAF_PWD_STATDB) dns := fmt.Sprintf("%s?_db_key=%s", path, key) db, err := gorm.Open(sqlite.Open(dns), &gorm.Config{}) @@ -300,6 +327,10 @@ func InitStatsDb(currentDir string) { global.GWAF_LOCAL_STATS_DB.Callback().Query().Before("gorm:update").Register("tenant_plugin:before_update", before_update) pathStatsSql(db) + + return isNewDb, nil + } else { + return false, nil } } diff --git a/waftask/task_createindex.go b/waftask/task_createindex.go index e3e4d19..dcb865e 100644 --- a/waftask/task_createindex.go +++ b/waftask/task_createindex.go @@ -2,6 +2,7 @@ package waftask import ( "SamWaf/common/zlog" + "SamWaf/enums" "SamWaf/global" "time" ) @@ -17,6 +18,24 @@ func TaskCreateIndex() { createStatDbIndex() } +// TaskCreateIndexByDbName 创建索引通过数据库名称 +func TaskCreateIndexByDbName(dbName string) { + + //主库索引创建 + if dbName == enums.DB_MAIN { + createMainDbIndex() + } + + //日志库索引创建 + if dbName == enums.DB_LOG { + createLogDbIndex() + } + //统计库索引创建 + if dbName == enums.DB_STATS { + createStatDbIndex() + } +} + func createMainDbIndex() { db := global.GWAF_LOCAL_DB if db == nil { diff --git a/waftask/task_health.go b/waftask/task_health.go index f2a5c41..f003068 100644 --- a/waftask/task_health.go +++ b/waftask/task_health.go @@ -43,7 +43,7 @@ func TaskHealth() { go func() { for { if global.GWAF_SHUTDOWN_SIGNAL { - zlog.Info("TaskHealth - Shutdown ") + zlog.Debug("TaskHealth - Shutdown ") cancel() return }