From d324ccc0977d0347e4f94529490fefbc72e54900 Mon Sep 17 00:00:00 2001 From: HTTP/2 410 Gone <24599209+Mr-Sheep@users.noreply.github.com> Date: Sun, 11 Oct 2026 08:22:37 -0700 Subject: [PATCH] feat(agent): support custom WebSocket headers with `HUB_HEADERS` (#2585) Co-authored-by: henrygd --- agent/client.go | 51 +++++++++++++++ agent/client_test.go | 124 ++++++++++++++++++++++++++++++++++++ agent/connection_manager.go | 2 +- 3 files changed, 176 insertions(+), 1 deletion(-) diff --git a/agent/client.go b/agent/client.go index c90f2c48d..ca42978e9 100644 --- a/agent/client.go +++ b/agent/client.go @@ -23,6 +23,7 @@ import ( "github.com/fxamacker/cbor/v2" "github.com/lxzan/gws" "golang.org/x/crypto/ssh" + "golang.org/x/net/http/httpguts" "golang.org/x/net/proxy" ) @@ -37,6 +38,9 @@ const ( // where the hub dials the agent and no outbound WebSocket client is expected. var errNoHubURL = errors.New("HUB_URL environment variable not set") +// errInvalidHubHeaders returned when there is malformed HUB_HEADERS value. +var errInvalidHubHeaders = errors.New("invalid HUB_HEADERS") + type caCertFileError struct { err error } @@ -64,6 +68,7 @@ type WebSocketClient struct { lastConnectAttempt time.Time // Timestamp of last connection attempt hubVerified bool // Whether the hub has been cryptographically verified tlsConfig *tls.Config // Optional TLS configuration with custom CA certificates + hubHeaders http.Header // Extra request headers from HUB_HEADERS } // newWebSocketClient creates a new WebSocket client for the given agent. @@ -89,6 +94,10 @@ func newWebSocketClient(agent *Agent) (client *WebSocketClient, err error) { if err != nil { return nil, err } + client.hubHeaders, err = getHubHeaders() + if err != nil { + return nil, err + } client.agent = agent client.hubRequest = &common.HubRequest[cbor.RawMessage]{} @@ -163,6 +172,38 @@ func getTLSConfig() (*tls.Config, error) { return &tls.Config{RootCAs: rootCAs}, nil } +// getHubHeaders parses HUB_HEADERS into extra headers sent with +// the WebSocket handshake. +// +// The value is a comma separated list of "Name: value" pairs: +// +// HUB_HEADERS="CF-Access-Client-Id: abc.access, CF-Access-Client-Secret: def456" +func getHubHeaders() (http.Header, error) { + raw, _ := utils.GetEnv("HUB_HEADERS") + if strings.TrimSpace(raw) == "" { + return nil, nil + } + var headers http.Header + for i, entry := range strings.FieldsFunc(raw, func(r rune) bool { + return r == ',' || r == '\n' || r == '\r' + }) { + // values are often secrets, so errors must not include them + name, value, ok := strings.Cut(entry, ":") + name, value = strings.TrimSpace(name), strings.TrimSpace(value) + if !ok || !httpguts.ValidHeaderFieldName(name) { + return nil, fmt.Errorf("%w: entry %d is not a valid 'Name: value' pair", errInvalidHubHeaders, i+1) + } + if !httpguts.ValidHeaderFieldValue(value) { + return nil, fmt.Errorf("%w: invalid value for header %q", errInvalidHubHeaders, name) + } + if headers == nil { + headers = make(http.Header) + } + headers.Add(name, value) + } + return headers, nil +} + // getOptions returns the WebSocket client options, creating them if necessary. // It configures the connection URL, TLS settings, and authentication headers. func (client *WebSocketClient) getOptions() *gws.ClientOption { @@ -195,6 +236,16 @@ func (client *WebSocketClient) getOptions() *gws.ClientOption { return proxy.FromEnvironment(), nil }, } + + // hub rejects a handshake without a valid X-Token or X-Beszel, ignoring user supplied values + for name, values := range client.hubHeaders { + if name == "X-Token" || name == "X-Beszel" { + slog.Warn("Ignoring reserved header in HUB_HEADERS", "header", name) + continue + } + client.options.RequestHeader[name] = values + } + return client.options } diff --git a/agent/client_test.go b/agent/client_test.go index 06d6c1a76..37ee8a60e 100644 --- a/agent/client_test.go +++ b/agent/client_test.go @@ -746,3 +746,127 @@ func TestWebSocketDeadlineCoversSlowCollection(t *testing.T) { t.Fatalf("WebSocket deadline %s is shorter than the slow-collection window of %s", wsDeadline, minimumDeadline) } } + +// TestGetHubHeaders covers HUB_HEADERS parsing +func TestGetHubHeaders(t *testing.T) { + testCases := []struct { + name string + raw string + want http.Header + wantErr bool + }{ + { + name: "unset", + raw: "", + want: nil, + }, + { + name: "two pairs with whitespace and trailing comma", + raw: " CF-Access-Client-Id:abc.access , CF-Access-Client-Secret: def456 ,", + want: http.Header{ + "Cf-Access-Client-Id": []string{"abc.access"}, + "Cf-Access-Client-Secret": []string{"def456"}, + }, + }, + { + name: "written across lines", + raw: "X-One: 1,\nX-Two: 2\n", + want: http.Header{"X-One": []string{"1"}, "X-Two": []string{"2"}}, + }, + { + name: "value may contain colons and equals", + raw: "Authorization: Bearer a=b:c", + want: http.Header{"Authorization": []string{"Bearer a=b:c"}}, + }, + { + name: "empty value", + raw: "X-Empty:", + want: http.Header{"X-Empty": []string{""}}, + }, + { + name: "repeated name adds value", + raw: "X-Multi: one, X-Multi: two", + want: http.Header{"X-Multi": []string{"one", "two"}}, + }, + { + name: "second pair missing colon", + raw: "X-One: 1, oops", + wantErr: true, + }, + { + name: "invalid header name", + raw: "Bad Header: value", + wantErr: true, + }, + { + name: "invalid header value", + raw: "X-Bad: va\x7flue", + wantErr: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("BESZEL_AGENT_HUB_HEADERS", tc.raw) + headers, err := getHubHeaders() + if tc.wantErr { + require.Error(t, err) + assert.ErrorIs(t, err, errInvalidHubHeaders) + assert.Nil(t, headers) + return + } + require.NoError(t, err) + assert.Equal(t, tc.want, headers) + }) + } +} + +// TestGetHubHeadersErrorOmitsValues verifies that parse errors do not leak +// header values, which are often secrets +func TestGetHubHeadersErrorOmitsValues(t *testing.T) { + const secret = "super-secret-value" + for _, raw := range []string{ + "CF-Access-Client-Secret: " + secret + "\x7f", + "Bad Name: " + secret, + secret, + } { + t.Setenv("BESZEL_AGENT_HUB_HEADERS", raw) + _, err := getHubHeaders() + require.ErrorIs(t, err, errInvalidHubHeaders) + assert.NotContains(t, err.Error(), secret) + } +} + +// TestWebSocketClient_GetOptionsHubHeaders verifies that +// HUB_HEADERS values reach the handshake, and user supplied X-BESZEL values are ignored +func TestWebSocketClient_GetOptionsHubHeaders(t *testing.T) { + agent := createTestAgent(t) + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") + t.Setenv("BESZEL_AGENT_HUB_HEADERS", + "CF-Access-Client-Id: abc.access, User-Agent: curl/8.0, x-token: hijack, X-BESZEL: 0.0.1") + + client, err := newWebSocketClient(agent) + require.NoError(t, err) + + header := client.getOptions().RequestHeader + assert.Equal(t, "abc.access", header.Get("Cf-Access-Client-Id")) + assert.Equal(t, "curl/8.0", header.Get("User-Agent")) + assert.Equal(t, "test-token", header.Get("X-Token")) + assert.Equal(t, beszel.Version, header.Get("X-Beszel")) +} + +// TestNewWebSocketClientInvalidHubHeaders verifies that a malformed +// HUB_HEADERS value fails client creation +func TestNewWebSocketClientInvalidHubHeaders(t *testing.T) { + agent := createTestAgent(t) + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") + t.Setenv("BESZEL_AGENT_HUB_HEADERS", "malfored-header") + + client, err := newWebSocketClient(agent) + + require.Error(t, err) + assert.Nil(t, client) + assert.ErrorIs(t, err, errInvalidHubHeaders) +} diff --git a/agent/connection_manager.go b/agent/connection_manager.go index c52a2d288..88f3a6511 100644 --- a/agent/connection_manager.go +++ b/agent/connection_manager.go @@ -168,7 +168,7 @@ func (c *ConnectionManager) Start(serverOptions ServerOptions) error { wsClient, err := newWebSocketClient(c.agent) if err != nil { var caCertErr *caCertFileError - if errors.As(err, &caCertErr) { + if errors.As(err, &caCertErr) || errors.Is(err, errInvalidHubHeaders) { return err } disableSSH, _ := utils.GetEnv("DISABLE_SSH")