bee/audit/internal/collector/cpu_telemetry.go

package collector

import (
	"bee/audit/internal/schema"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var (
	cpuSysBaseDir = "/sys/devices/system/cpu"
	socketIndexRe = regexp.MustCompile(`(?i)(?:package id|socket|cpu)\s*([0-9]+)`)
)

func enrichCPUsWithTelemetry(cpus []schema.HardwareCPU, doc sensorsDoc) []schema.HardwareCPU {
	if len(cpus) == 0 {
		return cpus
	}

	tempBySocket := cpuTempsFromSensors(doc, len(cpus))
	powerBySocket := cpuPowerFromSensors(doc, len(cpus))
	throttleBySocket := cpuThrottleBySocket()

	for i := range cpus {
		socket := 0
		if cpus[i].Socket != nil {
			socket = *cpus[i].Socket
		}
		if value, ok := tempBySocket[socket]; ok {
			cpus[i].TemperatureC = &value
		}
		if value, ok := powerBySocket[socket]; ok {
			cpus[i].PowerW = &value
		}
		if value, ok := throttleBySocket[socket]; ok {
			cpus[i].Throttled = &value
		}
	}

	return cpus
}

func cpuTempsFromSensors(doc sensorsDoc, cpuCount int) map[int]float64 {
	out := map[int]float64{}
	if len(doc) == 0 {
		return out
	}
	var fallback []float64
	for chip, features := range doc {
		for featureName, raw := range features {
			feature, ok := raw.(map[string]any)
			if !ok {
				continue
			}
			if classifySensorFeature(feature) != "temp" {
				continue
			}
			temp, ok := firstFeatureFloat(feature, "_input")
			if !ok {
				continue
			}
			if socket, ok := detectCPUSocket(chip, featureName); ok {
				if _, exists := out[socket]; !exists {
					out[socket] = temp
				}
				continue
			}
			if isLikelyCPUTemp(chip, featureName) {
				fallback = append(fallback, temp)
			}
		}
	}
	if len(out) == 0 && cpuCount == 1 && len(fallback) > 0 {
		out[0] = fallback[0]
	}
	return out
}

func cpuPowerFromSensors(doc sensorsDoc, cpuCount int) map[int]float64 {
	out := map[int]float64{}
	if len(doc) == 0 {
		return out
	}
	var fallback []float64
	for chip, features := range doc {
		for featureName, raw := range features {
			feature, ok := raw.(map[string]any)
			if !ok {
				continue
			}
			if classifySensorFeature(feature) != "power" {
				continue
			}
			power, ok := firstFeatureFloatWithContains(feature, []string{"power"})
			if !ok {
				continue
			}
			if socket, ok := detectCPUSocket(chip, featureName); ok {
				if _, exists := out[socket]; !exists {
					out[socket] = power
				}
				continue
			}
			if isLikelyCPUPower(chip, featureName) {
				fallback = append(fallback, power)
			}
		}
	}
	if len(out) == 0 && cpuCount == 1 && len(fallback) > 0 {
		out[0] = fallback[0]
	}
	return out
}

func detectCPUSocket(parts ...string) (int, bool) {
	for _, part := range parts {
		matches := socketIndexRe.FindStringSubmatch(strings.ToLower(part))
		if len(matches) == 2 {
			value, err := strconv.Atoi(matches[1])
			if err == nil {
				return value, true
			}
		}
	}
	return 0, false
}

func isLikelyCPUTemp(chip, feature string) bool {
	value := strings.ToLower(chip + " " + feature)
	return strings.Contains(value, "coretemp") ||
		strings.Contains(value, "k10temp") ||
		strings.Contains(value, "package id") ||
		strings.Contains(value, "tdie") ||
		strings.Contains(value, "tctl") ||
		strings.Contains(value, "cpu temp")
}

func isLikelyCPUPower(chip, feature string) bool {
	value := strings.ToLower(chip + " " + feature)
	return strings.Contains(value, "intel-rapl") ||
		strings.Contains(value, "package id") ||
		strings.Contains(value, "package-") ||
		strings.Contains(value, "cpu power")
}

func cpuThrottleBySocket() map[int]bool {
	out := map[int]bool{}
	cpuDirs, err := filepath.Glob(filepath.Join(cpuSysBaseDir, "cpu[0-9]*"))
	if err != nil {
		return out
	}
	sort.Strings(cpuDirs)
	for _, cpuDir := range cpuDirs {
		socket, ok := readSocketIndex(cpuDir)
		if !ok {
			continue
		}
		if cpuPackageThrottled(cpuDir) {
			out[socket] = true
		}
	}
	return out
}

func readSocketIndex(cpuDir string) (int, bool) {
	raw, err := os.ReadFile(filepath.Join(cpuDir, "topology", "physical_package_id"))
	if err != nil {
		return 0, false
	}
	value, err := strconv.Atoi(strings.TrimSpace(string(raw)))
	if err != nil || value < 0 {
		return 0, false
	}
	return value, true
}

func cpuPackageThrottled(cpuDir string) bool {
	paths := []string{
		filepath.Join(cpuDir, "thermal_throttle", "package_throttle_count"),
		filepath.Join(cpuDir, "thermal_throttle", "core_throttle_count"),
	}
	for _, path := range paths {
		raw, err := os.ReadFile(path)
		if err != nil {
			continue
		}
		value, err := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64)
		if err == nil && value > 0 {
			return true
		}
	}
	return false
}