ebpf/containers/apm_white_list_v2.go

package containers

import (
	"encoding/json"
	"fmt"
	"github.com/coroot/coroot-node-agent/common"
	"github.com/coroot/coroot-node-agent/utils"
	. "github.com/coroot/coroot-node-agent/utils/modelse"
	log "github.com/sirupsen/logrus"
	"os"
	"path"
	"strings"
)

func (r *Registry) getWhiteListAll() []WhiteSettingInfo {
	return r.whiteListRules["all"]
}

func (r *Registry) getWhiteListFromEnv() (WhiteDataV2, error) {
	whitelistStr := os.Getenv("Whitelist")
	if whitelistStr == "" {
		return WhiteDataV2{}, fmt.Errorf("Whitelist environment variable is not set")
	}

	// 按逗号分割字符串
	appRules := strings.Split(whitelistStr, ",")
	var settingList []WhiteSettingInfo

	for _, appRule := range appRules {
		// 去除前后空格
		appRule = strings.TrimSpace(appRule)
		if appRule == "" {
			continue
		}

		// 按冒号分割 appname:rule
		parts := strings.Split(appRule, ":")
		if len(parts) != 2 {
			log.Warnf("Invalid whitelist format: %s, expected 'appname:rule'", appRule)
			continue
		}

		appName := strings.TrimSpace(parts[0])
		rule := strings.TrimSpace(parts[1])

		if appName == "" || rule == "" {
			log.Warnf("Empty appname or rule in: %s", appRule)
			continue
		}

		// 创建 WhiteSettingInfo 结构
		setting := WhiteSettingInfo{
			AppName:    appName,
			ProcessKey: rule,
			Filters:    rule, // 根据现有代码逻辑，Filters 设置为 ProcessKey
		}

		settingList = append(settingList, setting)
	}

	return WhiteDataV2{
		SettingList:     settingList,
		LastUpdatedTime: 0, // 环境变量配置没有时间戳
	}, nil
}

func (r *Registry) setWhiteListV2(whiteData WhiteDataV2) {
	whiteListMap := make(WhiteListMap)
	code := "all"
	for _, setting := range whiteData.SettingList {
		setting.Filters = setting.ProcessKey
		whiteListMap[code] = append(whiteListMap[code], setting)
	}
	r.whiteListRules = whiteListMap
}

func (r *Registry) pullWhiteListV2() (bool, error) {
	if r.isFusing {
		return false, nil
	}
	if common.IsOpenFilter() {
		return false, nil
	}

	// 优先尝试从环境变量获取白名单配置
	whiteData, err := r.getWhiteListFromEnv()
	if err == nil {
		// 环境变量配置成功，使用环境变量配置
		log.Infof("Using whitelist configuration from environment variable")
		
		// 将白名单配置转换为 JSON 字符串并输出日志
		whiteDataJSON, jsonErr := json.Marshal(whiteData)
		if jsonErr != nil {
			log.Errorf("Failed to marshal whitelist data to JSON: %v", jsonErr)
		} else {
			log.Infof("Whitelist configuration from environment variable: %s", string(whiteDataJSON))
		}
		
		r.whiteLastUpdatedTime = whiteData.LastUpdatedTime
		r.setWhiteListV2(whiteData)
		saveRule(whiteData)
		return true, nil
	}

	// 环境变量配置失败，回退到原有的 WhiteListV2 函数
	log.Infof("Environment variable whitelist not available, falling back to server API: %v", err)
	
	nodeInfo := r.nodeInfo.GetNodeInfo()
	if nodeInfo == nil {
		return false, fmt.Errorf("could not find node info")
	}

	whiteListReq := WhiteListReq{
		HostId:    nodeInfo.HostID,
		AccountId: nodeInfo.AccountID,
		WhiteType: 2,
	}
	whiteData, err = r.connServer.WhiteListV2(whiteListReq)

	if err != nil {
		log.Errorf("report WhiteList info error is %v.", err)
		return false, err
	}

	//fmt.Println(r.whiteLastUpdatedTime)
	//fmt.Println(whiteData.LastUpdatedTime)
	// 后端切换主机与全局规则时,时间不会变化
	// 不用更新
	//if r.whiteLastUpdatedTime == whiteData.LastUpdatedTime {
	//	return false, nil
	//}
	//
	//// 更新时间
	r.whiteLastUpdatedTime = whiteData.LastUpdatedTime
	r.setWhiteListV2(whiteData)
	saveRule(whiteData)
	return true, nil
}

func saveRule(runtimeApps WhiteDataV2) {
	appStr, _ := json.Marshal(runtimeApps)
	dumpPath := path.Join(utils.GetDefaultRuntimePath(), "memdump")
	err := os.MkdirAll(dumpPath, 0755)
	if err != nil {
		log.Error(err)
	}
	fileName := path.Join(dumpPath, "rule.snap")
	err = os.WriteFile(fileName, appStr, 0644)
	if err != nil {
		log.Error(err)
	}
}