fix(agent): read TOKEN_FILE like KEY_FILE instead of sending the whole file (#2276)

agent/client.go

@@ -87,7 +87,27 @@ func getToken() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	return strings.TrimSpace(string(tokenBytes)), nil
+	return parseTokenFile(string(tokenBytes), tokenFile)
+}
+
+// parseTokenFile reads a single token from TOKEN_FILE.
+// Blank lines and comments are ignored. Multiple tokens are rejected because
+// the agent supports only one outbound hub connection.
+func parseTokenFile(contents, path string) (string, error) {
+	var token string
+	for line := range strings.Lines(contents) {
+		line = strings.TrimSpace(line)
+		if len(line) == 0 || strings.HasPrefix(line, "#") {
+			continue
+		}
+		if token != "" {
+			return "", fmt.Errorf("%s must contain a single token", path)
+		}
+		token = line
+	}
+	// An empty file keeps returning an empty token, as before: the caller decides
+	// what to do about it.
+	return token, nil
 }
 
 // getOptions returns the WebSocket client options, creating them if necessary.

agent/client_test.go

@@ -6,6 +6,7 @@ import (
 	"crypto/ed25519"
 	"net/url"
 	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
@@ -409,6 +410,41 @@ func TestGetToken(t *testing.T) {
 		assert.Equal(t, expectedToken, token)
 	})
 
+	t.Run("TOKEN_FILE with surrounding blank lines and comments", func(t *testing.T) {
+		expectedToken := "test-token-with-noise"
+		tokenFile := filepath.Join(t.TempDir(), "token")
+		require.NoError(t, os.WriteFile(tokenFile, []byte("# hub token\n\n"+expectedToken+"\n\n"), 0o600))
+
+		t.Setenv("TOKEN_FILE", tokenFile)
+
+		token, err := getToken()
+		assert.NoError(t, err)
+		assert.Equal(t, expectedToken, token)
+	})
+
+	t.Run("TOKEN_FILE with multiple tokens is rejected", func(t *testing.T) {
+		tokenFile := filepath.Join(t.TempDir(), "token")
+		require.NoError(t, os.WriteFile(tokenFile, []byte("11111111-1111-1111-1111-111111111111\n22222222-2222-2222-2222-222222222222\n"), 0o600))
+
+		t.Setenv("TOKEN_FILE", tokenFile)
+
+		token, err := getToken()
+		require.Error(t, err)
+		assert.Empty(t, token)
+		assert.Contains(t, err.Error(), "must contain a single token")
+	})
+
+	t.Run("TOKEN_FILE holding only comments behaves like an empty file", func(t *testing.T) {
+		tokenFile := filepath.Join(t.TempDir(), "token")
+		require.NoError(t, os.WriteFile(tokenFile, []byte("\n# only a comment\n"), 0o600))
+
+		t.Setenv("TOKEN_FILE", tokenFile)
+
+		token, err := getToken()
+		assert.NoError(t, err)
+		assert.Equal(t, "", token)
+	})
+
 	t.Run("token from BESZEL_AGENT_TOKEN_FILE", func(t *testing.T) {
 		// Create a temporary token file
 		expectedToken := "test-token-from-beszel-file"