refactor: streamline key retrieval logic in agent.go

beszel/cmd/agent/agent.go

@@ -7,7 +7,6 @@ import (
 	"log"
 	"os"
 	"strings"
-	"io/ioutil"
 )
 
 func main() {
@@ -22,19 +21,19 @@ func main() {
 		os.Exit(0)
 	}
 
-	var pubKey []byte
+	// Try to get the key from the KEY environment variable.
-	if pubKeyEnv, exists := os.LookupEnv("KEY"); exists {
+	pubKey := []byte(os.Getenv("KEY"))
-		pubKey = []byte(pubKeyEnv)
+
-	} else {
+	// If KEY is not set, try to read the key from the file specified by KEY_FILE.
-		keyFile := os.Getenv("KEY_FILE")
+	if len(pubKey) == 0 {
-		if keyFile != "" {
+		keyFile, exists := os.LookupEnv("KEY_FILE")
-			if keyData, err := ioutil.ReadFile(keyFile); err == nil {
+		if !exists {
-				pubKey = keyData
+			log.Fatal("Must set KEY or KEY_FILE environment variable")
-			} else {
+		}
-				log.Fatalf("Failed to read key from file '%s': %v", keyFile, err)
+		var err error
-			}
+		pubKey, err = os.ReadFile(keyFile)
-		} else {
+		if err != nil {
-			log.Fatal("KEY environment variable is not set, and KEY_FILE environment variable is not set")
+			log.Fatal(err)
 		}
 	}