!7 fix:create index

Merge pull request !7 from samwaf/fix_IC8E8S

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" //统计数据库
+
+)

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()

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

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
 	}
 }
 

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 {

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
 			}