support blacklisting and wildcard matching in SENSORS env var (#650)

- Moved sensor related code to sensors.go
- Added SensorConfig struct
- Added newSensorConfig
- Added tests

beszel/internal/agent/agent.go

@@ -4,42 +4,36 @@ package agent
 import (
 	"beszel"
 	"beszel/internal/entities/system"
-	"context"
 	"log/slog"
 	"os"
 	"strings"
 	"sync"
 	"time"
-
-	"github.com/shirou/gopsutil/v4/common"
 )
 
 type Agent struct {
-	sync.Mutex                                  // Used to lock agent while collecting data
+	sync.Mutex                               // Used to lock agent while collecting data
-	debug            bool                       // true if LOG_LEVEL is set to debug
+	debug         bool                       // true if LOG_LEVEL is set to debug
-	zfs              bool                       // true if system has arcstats
+	zfs           bool                       // true if system has arcstats
-	memCalc          string                     // Memory calculation formula
+	memCalc       string                     // Memory calculation formula
-	fsNames          []string                   // List of filesystem device names being monitored
+	fsNames       []string                   // List of filesystem device names being monitored
-	fsStats          map[string]*system.FsStats // Keeps track of disk stats for each filesystem
+	fsStats       map[string]*system.FsStats // Keeps track of disk stats for each filesystem
-	netInterfaces    map[string]struct{}        // Stores all valid network interfaces
+	netInterfaces map[string]struct{}        // Stores all valid network interfaces
-	netIoStats       system.NetIoStats          // Keeps track of bandwidth usage
+	netIoStats    system.NetIoStats          // Keeps track of bandwidth usage
-	dockerManager    *dockerManager             // Manages Docker API requests
+	dockerManager *dockerManager             // Manages Docker API requests
-	sensorsContext   context.Context            // Sensors context to override sys location
+	sensorConfig  *SensorConfig              // Sensors config
-	sensorsWhitelist map[string]struct{}        // List of sensors to monitor
+	systemInfo    system.Info                // Host system info
-	primarySensor    string                     // Value of PRIMARY_SENSOR env var
+	gpuManager    *GPUManager                // Manages GPU data
-	systemInfo       system.Info                // Host system info
+	cache         *SessionCache              // Cache for system stats based on primary session ID
-	gpuManager       *GPUManager                // Manages GPU data
-	cache            *SessionCache              // Cache for system stats based on primary session ID
 }
 
 func NewAgent() *Agent {
 	agent := &Agent{
-		sensorsContext: context.Background(),
+		fsStats: make(map[string]*system.FsStats),
-		fsStats:        make(map[string]*system.FsStats),
+		cache:   NewSessionCache(69 * time.Second),
-		cache:          NewSessionCache(69 * time.Second),
 	}
 	agent.memCalc, _ = GetEnv("MEM_CALC")
-	agent.primarySensor, _ = GetEnv("PRIMARY_SENSOR")
+	agent.sensorConfig = agent.newSensorConfig()
 	// Set up slog with a log level determined by the LOG_LEVEL env var
 	if logLevelStr, exists := GetEnv("LOG_LEVEL"); exists {
 		switch strings.ToLower(logLevelStr) {
@@ -55,24 +49,6 @@ func NewAgent() *Agent {
 
 	slog.Debug(beszel.Version)
 
-	// Set sensors context (allows overriding sys location for sensors)
-	if sysSensors, exists := GetEnv("SYS_SENSORS"); exists {
-		slog.Info("SYS_SENSORS", "path", sysSensors)
-		agent.sensorsContext = context.WithValue(agent.sensorsContext,
-			common.EnvKey, common.EnvMap{common.HostSysEnvKey: sysSensors},
-		)
-	}
-
-	// Set sensors whitelist
-	if sensors, exists := GetEnv("SENSORS"); exists {
-		agent.sensorsWhitelist = make(map[string]struct{})
-		for sensor := range strings.SplitSeq(sensors, ",") {
-			if sensor != "" {
-				agent.sensorsWhitelist[sensor] = struct{}{}
-			}
-		}
-	}
-
 	// initialize system info / docker manager
 	agent.initializeSystemInfo()
 	agent.initializeDiskInfo()

beszel/internal/agent/sensors.go

@@ -0,0 +1,142 @@
+package agent
+
+import (
+	"beszel/internal/entities/system"
+	"context"
+	"log/slog"
+	"path"
+	"strconv"
+	"strings"
+
+	"github.com/shirou/gopsutil/v4/common"
+	"github.com/shirou/gopsutil/v4/sensors"
+)
+
+type SensorConfig struct {
+	context       context.Context
+	sensors       map[string]struct{}
+	primarySensor string
+	isBlacklist   bool
+	hasWildcards  bool
+}
+
+func (a *Agent) newSensorConfig() *SensorConfig {
+	primarySensor, _ := GetEnv("PRIMARY_SENSOR")
+	sysSensors, _ := GetEnv("SYS_SENSORS")
+	sensors, _ := GetEnv("SENSORS")
+
+	return a.newSensorConfigWithEnv(primarySensor, sysSensors, sensors)
+}
+
+// newSensorConfigWithEnv creates a SensorConfig with the provided environment variables
+func (a *Agent) newSensorConfigWithEnv(primarySensor, sysSensors, sensors string) *SensorConfig {
+	config := &SensorConfig{
+		context:       context.Background(),
+		primarySensor: primarySensor,
+	}
+
+	// Set sensors context (allows overriding sys location for sensors)
+	if sysSensors != "" {
+		slog.Info("SYS_SENSORS", "path", sysSensors)
+		config.context = context.WithValue(config.context,
+			common.EnvKey, common.EnvMap{common.HostSysEnvKey: sysSensors},
+		)
+	}
+
+	// Set sensors whitelist
+	if sensors != "" {
+		// handle blacklist
+		if strings.HasPrefix(sensors, "-") {
+			config.isBlacklist = true
+			sensors = sensors[1:]
+		}
+
+		config.sensors = make(map[string]struct{})
+		for sensor := range strings.SplitSeq(sensors, ",") {
+			sensor = strings.TrimSpace(sensor)
+			if sensor != "" {
+				config.sensors[sensor] = struct{}{}
+				if strings.Contains(sensor, "*") {
+					config.hasWildcards = true
+				}
+			}
+		}
+	}
+
+	return config
+}
+
+// updateTemperatures updates the agent with the latest sensor temperatures
+func (a *Agent) updateTemperatures(systemStats *system.Stats) {
+	// skip if sensors whitelist is set to empty string
+	if a.sensorConfig.sensors != nil && len(a.sensorConfig.sensors) == 0 {
+		slog.Debug("Skipping temperature collection")
+		return
+	}
+
+	// reset high temp
+	a.systemInfo.DashboardTemp = 0
+
+	// get sensor data
+	temps, _ := sensors.TemperaturesWithContext(a.sensorConfig.context)
+	slog.Debug("Temperature", "sensors", temps)
+
+	// return if no sensors
+	if len(temps) == 0 {
+		return
+	}
+
+	systemStats.Temperatures = make(map[string]float64, len(temps))
+	for i, sensor := range temps {
+		// skip if temperature is unreasonable
+		if sensor.Temperature <= 0 || sensor.Temperature >= 200 {
+			continue
+		}
+		sensorName := sensor.SensorKey
+		if _, ok := systemStats.Temperatures[sensorName]; ok {
+			// if key already exists, append int to key
+			sensorName = sensorName + "_" + strconv.Itoa(i)
+		}
+		// skip if not in whitelist or blacklist
+		if !isValidSensor(sensorName, a.sensorConfig) {
+			continue
+		}
+		// set dashboard temperature
+		if a.sensorConfig.primarySensor == "" {
+			a.systemInfo.DashboardTemp = max(a.systemInfo.DashboardTemp, sensor.Temperature)
+		} else if a.sensorConfig.primarySensor == sensorName {
+			a.systemInfo.DashboardTemp = sensor.Temperature
+		}
+		systemStats.Temperatures[sensorName] = twoDecimals(sensor.Temperature)
+	}
+}
+
+// isValidSensor checks if a sensor is valid based on the sensor name and the sensor config
+func isValidSensor(sensorName string, config *SensorConfig) bool {
+	// If no sensors configuration, everything is valid
+	if config.sensors == nil {
+		return true
+	}
+
+	// Exact match - return true if whitelist, false if blacklist
+	if _, exactMatch := config.sensors[sensorName]; exactMatch {
+		return !config.isBlacklist
+	}
+
+	// If no wildcards, return false if blacklist, true if whitelist
+	if !config.hasWildcards {
+		return config.isBlacklist
+	}
+
+	// Check for wildcard patterns
+	for pattern := range config.sensors {
+		if !strings.Contains(pattern, "*") {
+			continue
+		}
+		if match, _ := path.Match(pattern, sensorName); match {
+			return !config.isBlacklist
+		}
+	}
+
+	return config.isBlacklist
+}

beszel/internal/agent/sensors_test.go

@@ -0,0 +1,338 @@
+//go:build testing
+// +build testing
+
+package agent
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/shirou/gopsutil/v4/common"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsValidSensor(t *testing.T) {
+	tests := []struct {
+		name          string
+		sensorName    string
+		config        *SensorConfig
+		expectedValid bool
+	}{
+		{
+			name:       "Whitelist - sensor in list",
+			sensorName: "cpu_temp",
+			config: &SensorConfig{
+				sensors:     map[string]struct{}{"cpu_temp": {}},
+				isBlacklist: false,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Whitelist - sensor not in list",
+			sensorName: "gpu_temp",
+			config: &SensorConfig{
+				sensors:     map[string]struct{}{"cpu_temp": {}},
+				isBlacklist: false,
+			},
+			expectedValid: false,
+		},
+		{
+			name:       "Blacklist - sensor in list",
+			sensorName: "cpu_temp",
+			config: &SensorConfig{
+				sensors:     map[string]struct{}{"cpu_temp": {}},
+				isBlacklist: true,
+			},
+			expectedValid: false,
+		},
+		{
+			name:       "Blacklist - sensor not in list",
+			sensorName: "gpu_temp",
+			config: &SensorConfig{
+				sensors:     map[string]struct{}{"cpu_temp": {}},
+				isBlacklist: true,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Whitelist with wildcard - matching pattern",
+			sensorName: "core_0_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"core_*_temp": {}},
+				isBlacklist:  false,
+				hasWildcards: true,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Whitelist with wildcard - non-matching pattern",
+			sensorName: "gpu_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"core_*_temp": {}},
+				isBlacklist:  false,
+				hasWildcards: true,
+			},
+			expectedValid: false,
+		},
+		{
+			name:       "Blacklist with wildcard - matching pattern",
+			sensorName: "core_0_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"core_*_temp": {}},
+				isBlacklist:  true,
+				hasWildcards: true,
+			},
+			expectedValid: false,
+		},
+		{
+			name:       "Blacklist with wildcard - non-matching pattern",
+			sensorName: "gpu_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"core_*_temp": {}},
+				isBlacklist:  true,
+				hasWildcards: true,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Nil sensor config",
+			sensorName: "any_temp",
+			config: &SensorConfig{
+				sensors: nil,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Mixed patterns in whitelist - exact match",
+			sensorName: "cpu_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}},
+				isBlacklist:  false,
+				hasWildcards: true,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Mixed patterns in whitelist - wildcard match",
+			sensorName: "core_1_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}},
+				isBlacklist:  false,
+				hasWildcards: true,
+			},
+			expectedValid: true,
+		},
+		{
+			name:       "Mixed patterns in blacklist - exact match",
+			sensorName: "cpu_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}},
+				isBlacklist:  true,
+				hasWildcards: true,
+			},
+			expectedValid: false,
+		},
+		{
+			name:       "Mixed patterns in blacklist - wildcard match",
+			sensorName: "core_1_temp",
+			config: &SensorConfig{
+				sensors:      map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}},
+				isBlacklist:  true,
+				hasWildcards: true,
+			},
+			expectedValid: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isValidSensor(tt.sensorName, tt.config)
+			assert.Equal(t, tt.expectedValid, result, "isValidSensor(%q, config) returned unexpected result", tt.sensorName)
+		})
+	}
+}
+
+func TestNewSensorConfigWithEnv(t *testing.T) {
+	agent := &Agent{}
+
+	tests := []struct {
+		name           string
+		primarySensor  string
+		sysSensors     string
+		sensors        string
+		expectedConfig *SensorConfig
+	}{
+		{
+			name:          "Empty configuration",
+			primarySensor: "",
+			sysSensors:    "",
+			sensors:       "",
+			expectedConfig: &SensorConfig{
+				context:       context.Background(),
+				primarySensor: "",
+				sensors:       nil,
+				isBlacklist:   false,
+				hasWildcards:  false,
+			},
+		},
+		{
+			name:          "Primary sensor only",
+			primarySensor: "cpu_temp",
+			sysSensors:    "",
+			sensors:       "",
+			expectedConfig: &SensorConfig{
+				context:       context.Background(),
+				primarySensor: "cpu_temp",
+				sensors:       nil,
+				isBlacklist:   false,
+				hasWildcards:  false,
+			},
+		},
+		{
+			name:          "Whitelist sensors",
+			primarySensor: "cpu_temp",
+			sysSensors:    "",
+			sensors:       "cpu_temp,gpu_temp",
+			expectedConfig: &SensorConfig{
+				context:       context.Background(),
+				primarySensor: "cpu_temp",
+				sensors: map[string]struct{}{
+					"cpu_temp": {},
+					"gpu_temp": {},
+				},
+				isBlacklist:  false,
+				hasWildcards: false,
+			},
+		},
+		{
+			name:          "Blacklist sensors",
+			primarySensor: "cpu_temp",
+			sysSensors:    "",
+			sensors:       "-cpu_temp,gpu_temp",
+			expectedConfig: &SensorConfig{
+				context:       context.Background(),
+				primarySensor: "cpu_temp",
+				sensors: map[string]struct{}{
+					"cpu_temp": {},
+					"gpu_temp": {},
+				},
+				isBlacklist:  true,
+				hasWildcards: false,
+			},
+		},
+		{
+			name:          "Sensors with wildcard",
+			primarySensor: "cpu_temp",
+			sysSensors:    "",
+			sensors:       "cpu_*,gpu_temp",
+			expectedConfig: &SensorConfig{
+				context:       context.Background(),
+				primarySensor: "cpu_temp",
+				sensors: map[string]struct{}{
+					"cpu_*":    {},
+					"gpu_temp": {},
+				},
+				isBlacklist:  false,
+				hasWildcards: true,
+			},
+		},
+		{
+			name:          "With SYS_SENSORS path",
+			primarySensor: "cpu_temp",
+			sysSensors:    "/custom/path",
+			sensors:       "cpu_temp",
+			expectedConfig: &SensorConfig{
+				primarySensor: "cpu_temp",
+				sensors: map[string]struct{}{
+					"cpu_temp": {},
+				},
+				isBlacklist:  false,
+				hasWildcards: false,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := agent.newSensorConfigWithEnv(tt.primarySensor, tt.sysSensors, tt.sensors)
+
+			// Check primary sensor
+			assert.Equal(t, tt.expectedConfig.primarySensor, result.primarySensor)
+
+			// Check sensor map
+			if tt.expectedConfig.sensors == nil {
+				assert.Nil(t, result.sensors)
+			} else {
+				assert.Equal(t, len(tt.expectedConfig.sensors), len(result.sensors))
+				for sensor := range tt.expectedConfig.sensors {
+					_, exists := result.sensors[sensor]
+					assert.True(t, exists, "Sensor %s should exist in the result", sensor)
+				}
+			}
+
+			// Check flags
+			assert.Equal(t, tt.expectedConfig.isBlacklist, result.isBlacklist)
+			assert.Equal(t, tt.expectedConfig.hasWildcards, result.hasWildcards)
+
+			// Check context
+			if tt.sysSensors != "" {
+				// Verify context contains correct values
+				envMap, ok := result.context.Value(common.EnvKey).(common.EnvMap)
+				require.True(t, ok, "Context should contain EnvMap")
+				sysPath, ok := envMap[common.HostSysEnvKey]
+				require.True(t, ok, "EnvMap should contain HostSysEnvKey")
+				assert.Equal(t, tt.sysSensors, sysPath)
+			}
+		})
+	}
+}
+
+func TestNewSensorConfig(t *testing.T) {
+	// Save original environment variables
+	originalPrimary, hasPrimary := os.LookupEnv("BESZEL_AGENT_PRIMARY_SENSOR")
+	originalSys, hasSys := os.LookupEnv("BESZEL_AGENT_SYS_SENSORS")
+	originalSensors, hasSensors := os.LookupEnv("BESZEL_AGENT_SENSORS")
+
+	// Restore environment variables after the test
+	defer func() {
+		// Clean up test environment variables
+		os.Unsetenv("BESZEL_AGENT_PRIMARY_SENSOR")
+		os.Unsetenv("BESZEL_AGENT_SYS_SENSORS")
+		os.Unsetenv("BESZEL_AGENT_SENSORS")
+
+		// Restore original values if they existed
+		if hasPrimary {
+			os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", originalPrimary)
+		}
+		if hasSys {
+			os.Setenv("BESZEL_AGENT_SYS_SENSORS", originalSys)
+		}
+		if hasSensors {
+			os.Setenv("BESZEL_AGENT_SENSORS", originalSensors)
+		}
+	}()
+
+	// Set test environment variables
+	os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", "test_primary")
+	os.Setenv("BESZEL_AGENT_SYS_SENSORS", "/test/path")
+	os.Setenv("BESZEL_AGENT_SENSORS", "test_sensor1,test_*,test_sensor3")
+
+	agent := &Agent{}
+	result := agent.newSensorConfig()
+
+	// Verify results
+	assert.Equal(t, "test_primary", result.primarySensor)
+	assert.NotNil(t, result.sensors)
+	assert.Equal(t, 3, len(result.sensors))
+	assert.True(t, result.hasWildcards)
+	assert.False(t, result.isBlacklist)
+
+	// Check that sys sensors path is in context
+	envMap, ok := result.context.Value(common.EnvKey).(common.EnvMap)
+	require.True(t, ok, "Context should contain EnvMap")
+	sysPath, ok := envMap[common.HostSysEnvKey]
+	require.True(t, ok, "EnvMap should contain HostSysEnvKey")
+	assert.Equal(t, "/test/path", sysPath)
+}

beszel/internal/agent/system.go

@@ -16,7 +16,6 @@ import (
 	"github.com/shirou/gopsutil/v4/host"
 	"github.com/shirou/gopsutil/v4/mem"
 	psutilNet "github.com/shirou/gopsutil/v4/net"
-	"github.com/shirou/gopsutil/v4/sensors"
 )
 
 // Sets initial / non-changing values about the host system
@@ -203,7 +202,7 @@ func (a *Agent) getSystemStats() system.Stats {
 			for _, gpu := range gpuData {
 				if gpu.Temperature > 0 {
 					systemStats.Temperatures[gpu.Name] = gpu.Temperature
-					if a.primarySensor == gpu.Name {
+					if a.sensorConfig.primarySensor == gpu.Name {
 						a.systemInfo.DashboardTemp = gpu.Temperature
 					}
 				}
@@ -224,52 +223,6 @@ func (a *Agent) getSystemStats() system.Stats {
 	return systemStats
 }
 
-func (a *Agent) updateTemperatures(systemStats *system.Stats) {
-	// skip if sensors whitelist is set to empty string
-	if a.sensorsWhitelist != nil && len(a.sensorsWhitelist) == 0 {
-		slog.Debug("Skipping temperature collection")
-		return
-	}
-
-	// reset high temp
-	a.systemInfo.DashboardTemp = 0
-
-	// get sensor data
-	temps, _ := sensors.TemperaturesWithContext(a.sensorsContext)
-	slog.Debug("Temperature", "sensors", temps)
-
-	// return if no sensors
-	if len(temps) == 0 {
-		return
-	}
-
-	systemStats.Temperatures = make(map[string]float64, len(temps))
-	for i, sensor := range temps {
-		// skip if temperature is unreasonable
-		if sensor.Temperature <= 0 || sensor.Temperature >= 200 {
-			continue
-		}
-		sensorName := sensor.SensorKey
-		if _, ok := systemStats.Temperatures[sensorName]; ok {
-			// if key already exists, append int to key
-			sensorName = sensorName + "_" + strconv.Itoa(i)
-		}
-		// skip if not in whitelist
-		if a.sensorsWhitelist != nil {
-			if _, nameInWhitelist := a.sensorsWhitelist[sensorName]; !nameInWhitelist {
-				continue
-			}
-		}
-		// set dashboard temperature
-		if a.primarySensor == "" {
-			a.systemInfo.DashboardTemp = max(a.systemInfo.DashboardTemp, sensor.Temperature)
-		} else if a.primarySensor == sensorName {
-			a.systemInfo.DashboardTemp = sensor.Temperature
-		}
-		systemStats.Temperatures[sensorName] = twoDecimals(sensor.Temperature)
-	}
-}
-
 // Returns the size of the ZFS ARC memory cache in bytes
 func getARCSize() (uint64, error) {
 	file, err := os.Open("/proc/spl/kstat/zfs/arcstats")