ebpf/flags/port_whitelist.go

package flags

import (
	"strconv"
	"strings"

	"github.com/coroot/coroot-node-agent/ebpftracer/l7"
)

// 端口白名单数据结构
type PortWhitelist struct {
	Ports map[uint16]bool
}

// 全局端口白名单实例
var (
	MysqlPorts   *PortWhitelist
	MariadbPorts *PortWhitelist
	TidbPorts    *PortWhitelist
)

// 创建端口白名单
func NewPortWhitelist(portList string) *PortWhitelist {
	whitelist := &PortWhitelist{
		Ports: make(map[uint16]bool),
	}

	if portList == "" {
		return whitelist
	}

	ports := strings.Split(portList, ",")
	for _, portStr := range ports {
		portStr = strings.TrimSpace(portStr)
		if port, err := strconv.ParseUint(portStr, 10, 16); err == nil {
			whitelist.Ports[uint16(port)] = true
		}
	}

	return whitelist
}

// 检查端口是否在白名单中
func (pw *PortWhitelist) Contains(port uint16) bool {
	if pw == nil {
		return false
	}
	return pw.Ports[port]
}

// 获取协议类型（基于端口白名单）
func GetProtocolByPort(port uint16) l7.Protocol {
	if MysqlPorts != nil && MysqlPorts.Contains(port) {
		return l7.ProtocolMysql
	}
	if MariadbPorts != nil && MariadbPorts.Contains(port) {
		return l7.ProtocolMariaDB
	}
	if TidbPorts != nil && TidbPorts.Contains(port) {
		return l7.ProtocolTiDB
	}
	// 如果端口不在白名单中，返回默认协议
	if *MysqlDefault == "mariadb" {
		return l7.ProtocolMariaDB
	}
	if *MysqlDefault == "tidb" {
		return l7.ProtocolTiDB
	}

	return l7.ProtocolMysql
}

// 初始化端口白名单
func InitPortWhitelists() {
	MysqlPorts = NewPortWhitelist(*MysqlPortWhitelist)
	MariadbPorts = NewPortWhitelist(*MariadbPortWhitelist)
	TidbPorts = NewPortWhitelist(*TidbPortWhitelist)
}