From dcf0abb0a787d194f8edeefbaeabcc03ba0020d2 Mon Sep 17 00:00:00 2001 From: "qcqcqc@wsl" <1220204124@zust.edu.cn> Date: Thu, 10 Apr 2025 17:12:45 +0800 Subject: [PATCH] init --- README.md | 0 cmd/manager/main.go | 44 ++++++++ config/execve_rules.json | 21 ++++ go.mod | 5 + go.sum | 2 + internal/constants/constants.go | 13 +++ internal/logger/logger.go | 83 ++++++++++++++ internal/manager/config_manager.go | 167 +++++++++++++++++++++++++++++ internal/models/config.go | 33 ++++++ internal/shm/shm.go | 111 +++++++++++++++++++ internal/utils/config_loader.go | 88 +++++++++++++++ 11 files changed, 567 insertions(+) create mode 100644 README.md create mode 100644 cmd/manager/main.go create mode 100644 config/execve_rules.json create mode 100644 go.mod create mode 100644 go.sum create mode 100644 internal/constants/constants.go create mode 100644 internal/logger/logger.go create mode 100644 internal/manager/config_manager.go create mode 100644 internal/models/config.go create mode 100644 internal/shm/shm.go create mode 100644 internal/utils/config_loader.go diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/cmd/manager/main.go b/cmd/manager/main.go new file mode 100644 index 0000000..355ff1c --- /dev/null +++ b/cmd/manager/main.go @@ -0,0 +1,44 @@ +package main + +import ( + "os" + "os/signal" + "syscall" + "time" + + "config-loader/internal/constants" + "config-loader/internal/logger" + "config-loader/internal/manager" +) + +func main() { + // Create config manager instance + configManager := manager.NewConfigManager(constants.ConfigFile, 5*time.Second) + + // Initialize the manager + if err := configManager.Initialize(); err != nil { + logger.ErrorLogger.Printf("Failed to initialize config manager: %v", err) + os.Exit(1) + } + + // Start watching for changes + configManager.StartWatching() + + // Set up signal handling + signalChan := make(chan os.Signal, 1) + signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM) + + logger.InfoLogger.Println("Configuration manager started. Press Ctrl+C to exit.") + + // Wait for shutdown signal + <-signalChan + logger.InfoLogger.Println("Shutdown signal received") + + // Clean up + configManager.Stop() + if err := configManager.ClearSharedMemory(); err != nil { + logger.ErrorLogger.Printf("Failed to clear shared memory: %v", err) + } + + logger.InfoLogger.Println("Configuration manager shutdown complete") +} diff --git a/config/execve_rules.json b/config/execve_rules.json new file mode 100644 index 0000000..d299bc0 --- /dev/null +++ b/config/execve_rules.json @@ -0,0 +1,21 @@ +{ + "enabled": true, + "rules": [ + { + "cmd": "nvidia-smi", + "type": "warn", + "msg": "在沐曦环境下请执行mx-smi" + }, + { + "cmd": "rm", + "type": "error", + "msg": "Error: rm command is forbidden" + }, + { + "cmd": "pip", + "type": "warn", + "msg": "使用pip安装torch时，请注意使用厂家支持版本", + "args": ["install", "torch"] + } + ] +} \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..ccf9537 --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module config-loader + +go 1.21 + +require golang.org/x/sys v0.15.0 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..063d2d3 --- /dev/null +++ b/go.sum @@ -0,0 +1,2 @@ +golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= diff --git a/internal/constants/constants.go b/internal/constants/constants.go new file mode 100644 index 0000000..f848199 --- /dev/null +++ b/internal/constants/constants.go @@ -0,0 +1,13 @@ +package constants + +const ( + MaxRules = 128 // 最大规则数量 + MaxArgs = 10 // 每条规则的最大参数数量 + MaxRuleCmdLength = 256 // 规则命令的最大长度 + MaxRuleTypeLength = 32 // 规则类型的最大长度 + MaxRuleMsgLength = 1024 // 规则消息的最大长度 + MaxArgLength = 256 // 单个参数的最大长度 + ShmKey = 0x78945 // 共享内存的键值 + ShmSize = 512 * 1024 // 共享内存的大小（字节） + ConfigFile = "./config/execve_rules.json" // 配置文件路径 +) diff --git a/internal/logger/logger.go b/internal/logger/logger.go new file mode 100644 index 0000000..b2d1de3 --- /dev/null +++ b/internal/logger/logger.go @@ -0,0 +1,83 @@ +package logger + +import ( + "fmt" + "log" + "os" + "path/filepath" + "runtime" + "time" +) + +const ( + colorRed = "\033[31m" + colorGreen = "\033[32m" + colorYellow = "\033[33m" + colorBlue = "\033[34m" + colorPurple = "\033[35m" + colorCyan = "\033[36m" + colorGray = "\033[37m" + colorReset = "\033[0m" + colorBold = "\033[1m" +) + +var ( + InfoLogger *log.Logger + DebugLogger *log.Logger + ErrorLogger *log.Logger + WarningLogger *log.Logger +) + +type LogWriter struct { + prefix string + color string +} + +func (w *LogWriter) Write(p []byte) (n int, err error) { + // 获取调用信息 + _, file, line, ok := runtime.Caller(3) // 调整调用深度 + if !ok { + file = "???" + line = 0 + } + + // 获取短文件名 + filename := filepath.Base(file) + + // 当前时间 + now := time.Now().Format("2006-01-02 15:04:05.000") + + // 构建日志前缀 + logPrefix := fmt.Sprintf("%s%s%s [%s%-7s%s] %s%s:%d%s | ", + colorGray, now, colorReset, + w.color, w.prefix, colorReset, + colorBold, filename, line, colorReset, + ) + + // 写入完整日志 + return os.Stdout.Write([]byte(logPrefix + string(p))) +} + +func init() { + InfoLogger = log.New(&LogWriter{prefix: "INFO", color: colorGreen}, "", 0) + DebugLogger = log.New(&LogWriter{prefix: "DEBUG", color: colorBlue}, "", 0) + ErrorLogger = log.New(&LogWriter{prefix: "ERROR", color: colorRed}, "", 0) + WarningLogger = log.New(&LogWriter{prefix: "WARN", color: colorYellow}, "", 0) +} + +// 辅助函数封装 +func Info(format string, v ...interface{}) { + InfoLogger.Printf(format, v...) +} + +func Debug(format string, v ...interface{}) { + DebugLogger.Printf(format, v...) +} + +func Error(format string, v ...interface{}) { + ErrorLogger.Printf(format, v...) +} + +func Warn(format string, v ...interface{}) { + WarningLogger.Printf(format, v...) +} diff --git a/internal/manager/config_manager.go b/internal/manager/config_manager.go new file mode 100644 index 0000000..8708ce1 --- /dev/null +++ b/internal/manager/config_manager.go @@ -0,0 +1,167 @@ +package manager + +import ( + "fmt" + "os" + "sync" + "time" + + "config-loader/internal/logger" + "config-loader/internal/models" + "config-loader/internal/shm" + "config-loader/internal/utils" +) + +// ConfigManager 负责配置的加载和同步 +type ConfigManager struct { + configFile string // 配置文件路径 + lastModified time.Time // 配置文件的最后修改时间 + watchInterval time.Duration // 配置文件监控的时间间隔 + stopChan chan struct{} // 停止信号通道 + wg sync.WaitGroup // 等待组，用于管理 goroutine + mu sync.RWMutex // 读写锁，保护 currentConfig 的并发访问 + currentConfig *models.ConfigData // 当前的配置数据 +} + +// NewConfigManager 创建一个新的 ConfigManager 实例 +func NewConfigManager(configFile string, watchInterval time.Duration) *ConfigManager { + if watchInterval == 0 { + watchInterval = 5 * time.Second // 默认监控间隔为 5 秒 + } + return &ConfigManager{ + configFile: configFile, + watchInterval: watchInterval, + stopChan: make(chan struct{}), + } +} + +// Initialize 加载初始配置并开始监控配置文件的变化 +func (cm *ConfigManager) Initialize() error { + logger.InfoLogger.Println("Initializing ConfigManager") + + // 尝试从共享内存中读取现有配置 + if config, err := shm.ReadConfigFromSharedMemory(); err == nil { + cm.currentConfig = config + logger.InfoLogger.Printf("Loaded existing configuration from shared memory. Enabled: %v, Rule Count: %d", + config.Enabled, config.RuleCount) + if config.RuleCount == 0 { + cm.ForceSync() // 如果规则数量为 0，强制同步配置 + } + } else { + logger.WarningLogger.Printf("Failed to read from shared memory: %v", err) + // 如果从共享内存读取失败，则从文件加载 + if err := cm.syncFromFile(); err != nil { + return fmt.Errorf("initial configuration load failed: %w", err) + } + } + + return nil +} + +// StartWatching 开始监控配置文件的变化 +func (cm *ConfigManager) StartWatching() { + cm.wg.Add(1) + go cm.watchConfigFile() + logger.InfoLogger.Println("Started watching configuration file for changes") +} + +// Stop 优雅地停止配置管理器 +func (cm *ConfigManager) Stop() { + logger.InfoLogger.Println("Stopping ConfigManager") + close(cm.stopChan) + cm.wg.Wait() +} + +// GetCurrentConfig 返回当前的配置（线程安全） +func (cm *ConfigManager) GetCurrentConfig() *models.ConfigData { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.currentConfig +} + +// ForceSync 强制从文件重新加载配置 +func (cm *ConfigManager) ForceSync() error { + logger.InfoLogger.Println("Force syncing configuration from file") + return cm.syncFromFile() +} + +// ClearSharedMemory 清除共享内存中的配置 +func (cm *ConfigManager) ClearSharedMemory() error { + logger.InfoLogger.Println("Clearing shared memory configuration") + // 这里可以实现实际的清除逻辑 + // 当前仅写入一个空配置 + emptyConfig := &models.ConfigData{ + Enabled: false, + RuleCount: 0, + } + return shm.WriteConfigToSharedMemory(emptyConfig) +} + +// 内部方法 + +// syncFromFile 从文件加载配置 +func (cm *ConfigManager) syncFromFile() error { + cm.mu.Lock() + defer cm.mu.Unlock() + + logger.DebugLogger.Printf("Loading configuration from file: %s", cm.configFile) + + // 从文件加载 JSON 配置 + jsonConfig, err := utils.LoadConfigFromFile(cm.configFile) + if err != nil { + return fmt.Errorf("failed to load config file: %w", err) + } + + // 转换为共享内存格式并写入共享内存 + cConfig := utils.ConvertToCConfig(jsonConfig) + if err := shm.WriteConfigToSharedMemory(cConfig); err != nil { + return fmt.Errorf("failed to write to shared memory: %w", err) + } + + cm.currentConfig = cConfig + if fileInfo, err := os.Stat(cm.configFile); err == nil { + cm.lastModified = fileInfo.ModTime() // 更新最后修改时间 + } + + return nil +} + +// watchConfigFile 监控配置文件的变化 +func (cm *ConfigManager) watchConfigFile() { + defer cm.wg.Done() + + ticker := time.NewTicker(cm.watchInterval) + defer ticker.Stop() + + for { + select { + case <-cm.stopChan: + logger.InfoLogger.Println("Stopping configuration file watcher") + return + case <-ticker.C: + if err := cm.checkAndReloadConfig(); err != nil { + logger.ErrorLogger.Printf("Error checking/reloading config: %v", err) + } + } + } +} + +// checkAndReloadConfig 检查并重新加载配置文件 +func (cm *ConfigManager) checkAndReloadConfig() error { + fileInfo, err := os.Stat(cm.configFile) + if err != nil { + return fmt.Errorf("failed to stat config file: %w", err) + } + + cm.mu.RLock() + lastMod := cm.lastModified + cm.mu.RUnlock() + + // 如果文件修改时间更新，则重新加载配置 + if fileInfo.ModTime().After(lastMod) { + logger.InfoLogger.Println("Configuration file changed, reloading") + return cm.syncFromFile() + } + + return nil +} diff --git a/internal/models/config.go b/internal/models/config.go new file mode 100644 index 0000000..eafc709 --- /dev/null +++ b/internal/models/config.go @@ -0,0 +1,33 @@ +package models + +import "config-loader/internal/constants" + +// Rule 映射 C 结构体 +type Rule struct { + Cmd [constants.MaxRuleCmdLength]byte // 命令 + Type [constants.MaxRuleTypeLength]byte // 类型 + Msg [constants.MaxRuleMsgLength]byte // 消息 + Args [constants.MaxArgs][constants.MaxArgLength]byte // 参数 + ArgCount int32 // 参数数量 +} + +// ConfigData 映射 C 结构体 +type ConfigData struct { + Enabled bool // 是否启用 + Rules [constants.MaxRules]Rule // 规则列表 + RuleCount int32 // 规则数量 +} + +// JSONRule 表示 JSON 配置中的一条规则 +type JSONRule struct { + Cmd string `json:"cmd"` // 命令 + Type string `json:"type"` // 类型 + Msg string `json:"msg"` // 消息 + Args []string `json:"args"` // 参数列表 +} + +// JSONConfig 表示整个 JSON 配置 +type JSONConfig struct { + Enabled bool `json:"enabled"` // 是否启用 + Rules []JSONRule `json:"rules"` // 规则列表 +} diff --git a/internal/shm/shm.go b/internal/shm/shm.go new file mode 100644 index 0000000..7d95aa3 --- /dev/null +++ b/internal/shm/shm.go @@ -0,0 +1,111 @@ +package shm + +import ( + "fmt" + "unsafe" + + "config-loader/internal/constants" + "config-loader/internal/logger" + "config-loader/internal/models" + + "golang.org/x/sys/unix" +) + +// shmget 调用系统调用 SYS_SHMGET 创建或获取共享内存段 +// key: 共享内存的键值 +// size: 共享内存的大小 +// shmflg: 标志位，用于指定权限和行为 +func shmget(key int, size int, shmflg int) (int, error) { + id, _, errno := unix.Syscall(unix.SYS_SHMGET, uintptr(key), uintptr(size), uintptr(shmflg)) + if errno != 0 { + return 0, errno + } + return int(id), nil +} + +// shmat 调用系统调用 SYS_SHMAT 将共享内存段附加到进程的地址空间 +// shmid: 共享内存段的 ID +// shmaddr: 指定的附加地址（通常为 0 表示系统自动选择） +// shmflg: 标志位，用于指定附加模式 +func shmat(shmid int, shmaddr uintptr, shmflg int) (unsafe.Pointer, error) { + addr, _, errno := unix.Syscall(unix.SYS_SHMAT, uintptr(shmid), shmaddr, uintptr(shmflg)) + if errno != 0 { + return nil, errno + } + return unsafe.Pointer(addr), nil +} + +// shmdt 调用系统调用 SYS_SHMDT 将共享内存段从进程的地址空间分离 +// shmaddr: 共享内存段的地址 +func shmdt(shmaddr unsafe.Pointer) error { + _, _, errno := unix.Syscall(unix.SYS_SHMDT, uintptr(shmaddr), 0, 0) + if errno != 0 { + return errno + } + return nil +} + +// WriteConfigToSharedMemory 将配置数据写入共享内存 +// config: 要写入的配置数据 +func WriteConfigToSharedMemory(config *models.ConfigData) error { + logger.DebugLogger.Println("Starting to write configuration to shared memory") + + // 创建或获取共享内存段 + shmID, err := shmget(constants.ShmKey, constants.ShmSize, unix.IPC_CREAT|0666) + if err != nil { + logger.ErrorLogger.Printf("shmget failed: %v", err) + return fmt.Errorf("shmget failed: %w", err) + } + + // 将共享内存段附加到进程地址空间 + shmPtr, err := shmat(shmID, 0, 0) + if err != nil { + logger.ErrorLogger.Printf("shmat failed: %v", err) + return fmt.Errorf("shmat failed: %w", err) + } + defer shmdt(shmPtr) // 确保在函数结束时分离共享内存 + + // 检查配置数据大小是否超过共享内存大小 + configSize := unsafe.Sizeof(*config) + if configSize > constants.ShmSize { + return fmt.Errorf("configuration size (%d) exceeds shared memory size (%d)", configSize, constants.ShmSize) + } + + // 将配置数据复制到共享内存 + data := (*[constants.ShmSize]byte)(unsafe.Pointer(config))[:] + copy((*[constants.ShmSize]byte)(shmPtr)[:], data) + + logger.InfoLogger.Println("Configuration successfully written to shared memory") + return nil +} + +// ReadConfigFromSharedMemory 从共享内存读取配置数据 +func ReadConfigFromSharedMemory() (*models.ConfigData, error) { + logger.DebugLogger.Println("Starting to read configuration from shared memory") + + // 获取共享内存段 + shmID, err := shmget(constants.ShmKey, constants.ShmSize, 0) + if err != nil { + return nil, fmt.Errorf("shmget failed: %w", err) + } + + // 将共享内存段附加到进程地址空间（只读模式） + shmPtr, err := shmat(shmID, 0, unix.SHM_RDONLY) + if err != nil { + return nil, fmt.Errorf("shmat failed: %w", err) + } + defer shmdt(shmPtr) // 确保在函数结束时分离共享内存 + + // 创建一个新的配置数据对象 + config := new(models.ConfigData) + configSize := unsafe.Sizeof(*config) + if configSize > constants.ShmSize { + return nil, fmt.Errorf("configuration size (%d) exceeds shared memory size (%d)", configSize, constants.ShmSize) + } + + // 从共享内存复制数据到配置对象 + copy((*[constants.ShmSize]byte)(unsafe.Pointer(config))[:], (*[constants.ShmSize]byte)(shmPtr)[:]) + + logger.InfoLogger.Println("Configuration successfully read from shared memory") + return config, nil +} diff --git a/internal/utils/config_loader.go b/internal/utils/config_loader.go new file mode 100644 index 0000000..b7784c7 --- /dev/null +++ b/internal/utils/config_loader.go @@ -0,0 +1,88 @@ +package utils + +import ( + "encoding/json" + "os" + + "config-loader/internal/constants" + "config-loader/internal/logger" + "config-loader/internal/models" +) + +func LoadConfigFromFile(filename string) (*models.JSONConfig, error) { + logger.DebugLogger.Printf("Loading config from file: %s", filename) + + data, err := os.ReadFile(filename) + if err != nil { + logger.ErrorLogger.Printf("Failed to read config file: %v", err) + return nil, err + } + + var config models.JSONConfig + err = json.Unmarshal(data, &config) + if err != nil { + logger.ErrorLogger.Printf("Failed to unmarshal config: %v", err) + return nil, err + } + + // 添加跳过规则 + skipRules := []models.JSONRule{ + { + Cmd: "dirname", + Type: "skip", + }, + { + Cmd: "lesspipe", + Type: "skip", + }, + { + Cmd: "command-not-found", + Type: "skip", + }, + { + Cmd: "dircolors", + Type: "skip", + }, + { + Cmd: "basename", + Type: "skip", + }, + } + + // 将规则添加到配置中 + config.Rules = append(config.Rules, skipRules...) + + logger.InfoLogger.Println("Configuration successfully loaded from file with skip rules added") + return &config, nil +} + +func ConvertToCConfig(jsonConfig *models.JSONConfig) *models.ConfigData { + logger.DebugLogger.Println("Converting JSON config to C config") + + cConfig := &models.ConfigData{ + Enabled: jsonConfig.Enabled, + RuleCount: int32(len(jsonConfig.Rules)), + } + + for i, rule := range jsonConfig.Rules { + if i >= constants.MaxRules { + logger.WarningLogger.Printf("Exceeded maximum rules limit (%d)", constants.MaxRules) + break + } + copy(cConfig.Rules[i].Cmd[:], rule.Cmd) + copy(cConfig.Rules[i].Type[:], rule.Type) + copy(cConfig.Rules[i].Msg[:], rule.Msg) + cConfig.Rules[i].ArgCount = int32(len(rule.Args)) + + for j, arg := range rule.Args { + if j >= constants.MaxArgs { + logger.WarningLogger.Printf("Exceeded maximum arguments limit (%d) for rule %d", constants.MaxArgs, i) + break + } + copy(cConfig.Rules[i].Args[j][:], arg) + } + } + + logger.InfoLogger.Println("Successfully converted JSON config to C config") + return cConfig +}