diff --git a/agent/agent.go b/agent/agent.go index c717844d..764bc23e 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -6,7 +6,6 @@ package agent import ( "log/slog" - "os" "strings" "sync" "time" @@ -69,11 +68,11 @@ func NewAgent(dataDir ...string) (agent *Agent, err error) { slog.Info("Data directory", "path", agent.dataDir) } - agent.memCalc, _ = GetEnv("MEM_CALC") + agent.memCalc, _ = utils.GetEnv("MEM_CALC") agent.sensorConfig = agent.newSensorConfig() // Parse disk usage cache duration (e.g., "15m", "1h") to avoid waking sleeping disks - if diskUsageCache, exists := GetEnv("DISK_USAGE_CACHE"); exists { + if diskUsageCache, exists := utils.GetEnv("DISK_USAGE_CACHE"); exists { if duration, err := time.ParseDuration(diskUsageCache); err == nil { agent.diskUsageCacheDuration = duration slog.Info("DISK_USAGE_CACHE", "duration", duration) @@ -83,7 +82,7 @@ func NewAgent(dataDir ...string) (agent *Agent, err error) { } // Set up slog with a log level determined by the LOG_LEVEL env var - if logLevelStr, exists := GetEnv("LOG_LEVEL"); exists { + if logLevelStr, exists := utils.GetEnv("LOG_LEVEL"); exists { switch strings.ToLower(logLevelStr) { case "debug": agent.debug = true @@ -104,7 +103,7 @@ func NewAgent(dataDir ...string) (agent *Agent, err error) { agent.refreshSystemDetails() // SMART_INTERVAL env var to update smart data at this interval - if smartIntervalEnv, exists := GetEnv("SMART_INTERVAL"); exists { + if smartIntervalEnv, exists := utils.GetEnv("SMART_INTERVAL"); exists { if duration, err := time.ParseDuration(smartIntervalEnv); err == nil && duration > 0 { agent.systemDetails.SmartInterval = duration slog.Info("SMART_INTERVAL", "duration", duration) @@ -149,15 +148,6 @@ func NewAgent(dataDir ...string) (agent *Agent, err error) { return agent, nil } -// GetEnv retrieves an environment variable with a "BESZEL_AGENT_" prefix, or falls back to the unprefixed key. -func GetEnv(key string) (value string, exists bool) { - if value, exists = os.LookupEnv("BESZEL_AGENT_" + key); exists { - return value, exists - } - // Fallback to the old unprefixed key - return os.LookupEnv(key) -} - func (a *Agent) gatherStats(options common.DataRequestOptions) *system.CombinedData { a.Lock() defer a.Unlock() diff --git a/agent/client.go b/agent/client.go index 08984696..c71a28b3 100644 --- a/agent/client.go +++ b/agent/client.go @@ -14,6 +14,7 @@ import ( "time" "github.com/henrygd/beszel" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/common" "github.com/fxamacker/cbor/v2" @@ -43,7 +44,7 @@ type WebSocketClient struct { // newWebSocketClient creates a new WebSocket client for the given agent. // It reads configuration from environment variables and validates the hub URL. func newWebSocketClient(agent *Agent) (client *WebSocketClient, err error) { - hubURLStr, exists := GetEnv("HUB_URL") + hubURLStr, exists := utils.GetEnv("HUB_URL") if !exists { return nil, errors.New("HUB_URL environment variable not set") } @@ -72,12 +73,12 @@ func newWebSocketClient(agent *Agent) (client *WebSocketClient, err error) { // If neither is set, it returns an error. func getToken() (string, error) { // get token from env var - token, _ := GetEnv("TOKEN") + token, _ := utils.GetEnv("TOKEN") if token != "" { return token, nil } // get token from file - tokenFile, _ := GetEnv("TOKEN_FILE") + tokenFile, _ := utils.GetEnv("TOKEN_FILE") if tokenFile == "" { return "", errors.New("must set TOKEN or TOKEN_FILE") } @@ -197,7 +198,7 @@ func (client *WebSocketClient) handleAuthChallenge(msg *common.HubRequest[cbor.R } if authRequest.NeedSysInfo { - response.Name, _ = GetEnv("SYSTEM_NAME") + response.Name, _ = utils.GetEnv("SYSTEM_NAME") response.Hostname = client.agent.systemDetails.Hostname serverAddr := client.agent.connectionManager.serverOptions.Addr _, response.Port, _ = net.SplitHostPort(serverAddr) diff --git a/agent/data_dir.go b/agent/data_dir.go index d96844b4..12bc623d 100644 --- a/agent/data_dir.go +++ b/agent/data_dir.go @@ -6,6 +6,8 @@ import ( "os" "path/filepath" "runtime" + + "github.com/henrygd/beszel/agent/utils" ) // GetDataDir returns the path to the data directory for the agent and an error @@ -16,7 +18,7 @@ func GetDataDir(dataDirs ...string) (string, error) { return testDataDirs(dataDirs) } - dataDir, _ := GetEnv("DATA_DIR") + dataDir, _ := utils.GetEnv("DATA_DIR") if dataDir != "" { dataDirs = append(dataDirs, dataDir) } diff --git a/agent/disk.go b/agent/disk.go index 940b11d8..4f6e5502 100644 --- a/agent/disk.go +++ b/agent/disk.go @@ -38,7 +38,7 @@ func isDockerSpecialMountpoint(mountpoint string) bool { // Sets up the filesystems to monitor for disk usage and I/O. func (a *Agent) initializeDiskInfo() { - filesystem, _ := GetEnv("FILESYSTEM") + filesystem, _ := utils.GetEnv("FILESYSTEM") efPath := "/extra-filesystems" hasRoot := false isWindows := runtime.GOOS == "windows" @@ -142,7 +142,7 @@ func (a *Agent) initializeDiskInfo() { } // Add EXTRA_FILESYSTEMS env var values to fsStats - if extraFilesystems, exists := GetEnv("EXTRA_FILESYSTEMS"); exists { + if extraFilesystems, exists := utils.GetEnv("EXTRA_FILESYSTEMS"); exists { for fsEntry := range strings.SplitSeq(extraFilesystems, ",") { // Parse custom name from format: device__customname fs, customName := parseFilesystemEntry(fsEntry) diff --git a/agent/docker.go b/agent/docker.go index adb3baf9..1a04adf7 100644 --- a/agent/docker.go +++ b/agent/docker.go @@ -488,7 +488,7 @@ func (dm *dockerManager) deleteContainerStatsSync(id string) { // Creates a new http client for Docker or Podman API func newDockerManager() *dockerManager { - dockerHost, exists := GetEnv("DOCKER_HOST") + dockerHost, exists := utils.GetEnv("DOCKER_HOST") if exists { // return nil if set to empty string if dockerHost == "" { @@ -524,7 +524,7 @@ func newDockerManager() *dockerManager { // configurable timeout timeout := time.Millisecond * time.Duration(dockerTimeoutMs) - if t, set := GetEnv("DOCKER_TIMEOUT"); set { + if t, set := utils.GetEnv("DOCKER_TIMEOUT"); set { timeout, err = time.ParseDuration(t) if err != nil { slog.Error(err.Error()) @@ -541,7 +541,7 @@ func newDockerManager() *dockerManager { // Read container exclusion patterns from environment variable var excludeContainers []string - if excludeStr, set := GetEnv("EXCLUDE_CONTAINERS"); set && excludeStr != "" { + if excludeStr, set := utils.GetEnv("EXCLUDE_CONTAINERS"); set && excludeStr != "" { parts := strings.SplitSeq(excludeStr, ",") for part := range parts { trimmed := strings.TrimSpace(part) diff --git a/agent/gpu.go b/agent/gpu.go index 84b997da..8a0ffc1b 100644 --- a/agent/gpu.go +++ b/agent/gpu.go @@ -688,7 +688,7 @@ func (gm *GPUManager) resolveLegacyCollectorPriority(caps gpuCapabilities) []col priorities := make([]collectorSource, 0, 4) if caps.hasNvidiaSmi && !caps.hasTegrastats { - if nvml, _ := GetEnv("NVML"); nvml == "true" { + if nvml, _ := utils.GetEnv("NVML"); nvml == "true" { priorities = append(priorities, collectorSourceNVML, collectorSourceNvidiaSMI) } else { priorities = append(priorities, collectorSourceNvidiaSMI) @@ -696,7 +696,7 @@ func (gm *GPUManager) resolveLegacyCollectorPriority(caps gpuCapabilities) []col } if caps.hasRocmSmi { - if val, _ := GetEnv("AMD_SYSFS"); val == "true" { + if val, _ := utils.GetEnv("AMD_SYSFS"); val == "true" { priorities = append(priorities, collectorSourceAmdSysfs) } else { priorities = append(priorities, collectorSourceRocmSMI) @@ -729,7 +729,7 @@ func (gm *GPUManager) resolveLegacyCollectorPriority(caps gpuCapabilities) []col // NewGPUManager creates and initializes a new GPUManager func NewGPUManager() (*GPUManager, error) { - if skipGPU, _ := GetEnv("SKIP_GPU"); skipGPU == "true" { + if skipGPU, _ := utils.GetEnv("SKIP_GPU"); skipGPU == "true" { return nil, nil } var gm GPUManager @@ -746,7 +746,7 @@ func NewGPUManager() (*GPUManager, error) { } // if GPU_COLLECTOR is set, start user-defined collectors. - if collectorConfig, ok := GetEnv("GPU_COLLECTOR"); ok && strings.TrimSpace(collectorConfig) != "" { + if collectorConfig, ok := utils.GetEnv("GPU_COLLECTOR"); ok && strings.TrimSpace(collectorConfig) != "" { priorities := parseCollectorPriority(collectorConfig) if gm.startCollectorsByPriority(priorities, caps) == 0 { return nil, fmt.Errorf("no configured GPU collectors are available") diff --git a/agent/gpu_intel.go b/agent/gpu_intel.go index 1eaeb11d..fd73ed84 100644 --- a/agent/gpu_intel.go +++ b/agent/gpu_intel.go @@ -7,6 +7,7 @@ import ( "strconv" "strings" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/entities/system" ) @@ -52,7 +53,7 @@ func (gm *GPUManager) updateIntelFromStats(sample *intelGpuStats) bool { func (gm *GPUManager) collectIntelStats() (err error) { // Build command arguments, optionally selecting a device via -d args := []string{"-s", intelGpuStatsInterval, "-l"} - if dev, ok := GetEnv("INTEL_GPU_DEVICE"); ok && dev != "" { + if dev, ok := utils.GetEnv("INTEL_GPU_DEVICE"); ok && dev != "" { args = append(args, "-d", dev) } cmd := exec.Command(intelGpuStatsCmd, args...) diff --git a/agent/network.go b/agent/network.go index 2eb00bc4..1a66b2ad 100644 --- a/agent/network.go +++ b/agent/network.go @@ -95,7 +95,7 @@ func (a *Agent) initializeNetIoStats() { a.netInterfaces = make(map[string]struct{}, 0) // parse NICS env var for whitelist / blacklist - nicsEnvVal, nicsEnvExists := GetEnv("NICS") + nicsEnvVal, nicsEnvExists := utils.GetEnv("NICS") var nicCfg *NicConfig if nicsEnvExists { nicCfg = newNicConfig(nicsEnvVal) diff --git a/agent/sensors.go b/agent/sensors.go index 9cdc4c45..08cd29dd 100644 --- a/agent/sensors.go +++ b/agent/sensors.go @@ -27,9 +27,9 @@ type SensorConfig struct { } func (a *Agent) newSensorConfig() *SensorConfig { - primarySensor, _ := GetEnv("PRIMARY_SENSOR") - sysSensors, _ := GetEnv("SYS_SENSORS") - sensorsEnvVal, sensorsSet := GetEnv("SENSORS") + primarySensor, _ := utils.GetEnv("PRIMARY_SENSOR") + sysSensors, _ := utils.GetEnv("SYS_SENSORS") + sensorsEnvVal, sensorsSet := utils.GetEnv("SENSORS") skipCollection := sensorsSet && sensorsEnvVal == "" return a.newSensorConfigWithEnv(primarySensor, sysSensors, sensorsEnvVal, skipCollection) diff --git a/agent/server.go b/agent/server.go index bef024ff..5f42bbb4 100644 --- a/agent/server.go +++ b/agent/server.go @@ -12,6 +12,7 @@ import ( "time" "github.com/henrygd/beszel" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/common" "github.com/henrygd/beszel/internal/entities/system" @@ -36,7 +37,7 @@ var hubVersions map[string]semver.Version // and begins listening for connections. Returns an error if the server // is already running or if there's an issue starting the server. func (a *Agent) StartServer(opts ServerOptions) error { - if disableSSH, _ := GetEnv("DISABLE_SSH"); disableSSH == "true" { + if disableSSH, _ := utils.GetEnv("DISABLE_SSH"); disableSSH == "true" { return errors.New("SSH disabled") } if a.server != nil { @@ -238,11 +239,11 @@ func ParseKeys(input string) ([]gossh.PublicKey, error) { // and finally defaults to ":45876". func GetAddress(addr string) string { if addr == "" { - addr, _ = GetEnv("LISTEN") + addr, _ = utils.GetEnv("LISTEN") } if addr == "" { // Legacy PORT environment variable support - addr, _ = GetEnv("PORT") + addr, _ = utils.GetEnv("PORT") } if addr == "" { return ":45876" @@ -258,7 +259,7 @@ func GetAddress(addr string) string { // It checks the NETWORK environment variable first, then infers from // the address format: addresses starting with "/" are "unix", others are "tcp". func GetNetwork(addr string) string { - if network, ok := GetEnv("NETWORK"); ok && network != "" { + if network, ok := utils.GetEnv("NETWORK"); ok && network != "" { return network } if strings.HasPrefix(addr, "/") { diff --git a/agent/smart.go b/agent/smart.go index 09562d26..d3a131dd 100644 --- a/agent/smart.go +++ b/agent/smart.go @@ -18,6 +18,7 @@ import ( "sync" "time" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/entities/smart" ) @@ -156,7 +157,7 @@ func (sm *SmartManager) ScanDevices(force bool) error { currentDevices := sm.devicesSnapshot() var configuredDevices []*DeviceInfo - if configuredRaw, ok := GetEnv("SMART_DEVICES"); ok { + if configuredRaw, ok := utils.GetEnv("SMART_DEVICES"); ok { slog.Info("SMART_DEVICES", "value", configuredRaw) config := strings.TrimSpace(configuredRaw) if config == "" { @@ -260,7 +261,7 @@ func (sm *SmartManager) parseConfiguredDevices(config string) ([]*DeviceInfo, er } func (sm *SmartManager) refreshExcludedDevices() { - rawValue, _ := GetEnv("EXCLUDE_SMART") + rawValue, _ := utils.GetEnv("EXCLUDE_SMART") sm.excludedDevices = make(map[string]struct{}) for entry := range strings.SplitSeq(rawValue, ",") { diff --git a/agent/systemd.go b/agent/systemd.go index ac9b0ff2..3e37fef7 100644 --- a/agent/systemd.go +++ b/agent/systemd.go @@ -15,6 +15,7 @@ import ( "time" "github.com/coreos/go-systemd/v22/dbus" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/entities/systemd" ) @@ -49,7 +50,7 @@ func isSystemdAvailable() bool { // newSystemdManager creates a new systemdManager. func newSystemdManager() (*systemdManager, error) { - if skipSystemd, _ := GetEnv("SKIP_SYSTEMD"); skipSystemd == "true" { + if skipSystemd, _ := utils.GetEnv("SKIP_SYSTEMD"); skipSystemd == "true" { return nil, nil } @@ -294,7 +295,7 @@ func unescapeServiceName(name string) string { // otherwise defaults to "*service". func getServicePatterns() []string { patterns := []string{} - if envPatterns, _ := GetEnv("SERVICE_PATTERNS"); envPatterns != "" { + if envPatterns, _ := utils.GetEnv("SERVICE_PATTERNS"); envPatterns != "" { for pattern := range strings.SplitSeq(envPatterns, ",") { pattern = strings.TrimSpace(pattern) if pattern == "" { diff --git a/agent/utils/utils.go b/agent/utils/utils.go index 32cfcd8f..86b4567a 100644 --- a/agent/utils/utils.go +++ b/agent/utils/utils.go @@ -7,6 +7,14 @@ import ( "strings" ) +// GetEnv retrieves an environment variable with a "BESZEL_AGENT_" prefix, or falls back to the unprefixed key. +func GetEnv(key string) (value string, exists bool) { + if value, exists = os.LookupEnv("BESZEL_AGENT_" + key); exists { + return value, exists + } + return os.LookupEnv(key) +} + // BytesToMegabytes converts bytes to megabytes and rounds to two decimal places. func BytesToMegabytes(b float64) float64 { return TwoDecimals(b / 1048576) diff --git a/agent/utils/utils_test.go b/agent/utils/utils_test.go index eff27240..9231a6c4 100644 --- a/agent/utils/utils_test.go +++ b/agent/utils/utils_test.go @@ -128,3 +128,38 @@ func TestReadUintFile(t *testing.T) { assert.Equal(t, uint64(0), val) }) } + +func TestGetEnv(t *testing.T) { + key := "TEST_VAR" + prefixedKey := "BESZEL_AGENT_" + key + + t.Run("prefixed variable exists", func(t *testing.T) { + os.Setenv(prefixedKey, "prefixed_val") + os.Setenv(key, "unprefixed_val") + defer os.Unsetenv(prefixedKey) + defer os.Unsetenv(key) + + val, exists := GetEnv(key) + assert.True(t, exists) + assert.Equal(t, "prefixed_val", val) + }) + + t.Run("only unprefixed variable exists", func(t *testing.T) { + os.Unsetenv(prefixedKey) + os.Setenv(key, "unprefixed_val") + defer os.Unsetenv(key) + + val, exists := GetEnv(key) + assert.True(t, exists) + assert.Equal(t, "unprefixed_val", val) + }) + + t.Run("neither variable exists", func(t *testing.T) { + os.Unsetenv(prefixedKey) + os.Unsetenv(key) + + val, exists := GetEnv(key) + assert.False(t, exists) + assert.Empty(t, val) + }) +} diff --git a/internal/cmd/agent/agent.go b/internal/cmd/agent/agent.go index 425dd9f8..a27908db 100644 --- a/internal/cmd/agent/agent.go +++ b/internal/cmd/agent/agent.go @@ -9,6 +9,7 @@ import ( "github.com/henrygd/beszel" "github.com/henrygd/beszel/agent" "github.com/henrygd/beszel/agent/health" + "github.com/henrygd/beszel/agent/utils" "github.com/spf13/pflag" "golang.org/x/crypto/ssh" ) @@ -116,12 +117,12 @@ func (opts *cmdOptions) loadPublicKeys() ([]ssh.PublicKey, error) { } // Try environment variable - if key, ok := agent.GetEnv("KEY"); ok && key != "" { + if key, ok := utils.GetEnv("KEY"); ok && key != "" { return agent.ParseKeys(key) } // Try key file - keyFile, ok := agent.GetEnv("KEY_FILE") + keyFile, ok := utils.GetEnv("KEY_FILE") if !ok { return nil, fmt.Errorf("no key provided: must set -key flag, KEY env var, or KEY_FILE env var. Use 'beszel-agent help' for usage") } diff --git a/internal/hub/agent_connect_test.go b/internal/hub/agent_connect_test.go index 0266c7ed..131edb88 100644 --- a/internal/hub/agent_connect_test.go +++ b/internal/hub/agent_connect_test.go @@ -917,7 +917,7 @@ func TestAgentWebSocketIntegration(t *testing.T) { // Wait for connection result maxWait := 2 * time.Second - time.Sleep(20 * time.Millisecond) + time.Sleep(40 * time.Millisecond) checkInterval := 20 * time.Millisecond timeout := time.After(maxWait) ticker := time.Tick(checkInterval)