feat(agent): support custom WebSocket headers with HUB_HEADERS (#2585)

Co-authored-by: henrygd <hank@henrygd.me>

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
 }
 

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)
+}

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")