diff --git a/agent/client.go b/agent/client.go index 40e53df6..2436b294 100644 --- a/agent/client.go +++ b/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. diff --git a/agent/client_test.go b/agent/client_test.go index e9c2bd3f..1b987eef 100644 --- a/agent/client_test.go +++ b/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"