Validate duckgres.query_source against its allowed values (#948)

The duckgres.query_source session GUC is a billing dimension: it flows
verbatim from the session into the compute-usage bucket key
(duckgres_org_compute_usage.query_source) and onward into the billing
pull API. It previously accepted ANY string on every set path (simple
SET, split multi-statement batches, extended-protocol Parse, and the
-c startup option), handing clients unbounded cardinality and arbitrary
junk (10KB strings, control characters) in the billing table and its
exports.

Now the value is a closed enum, validated once at the transpiler
interception (transform.NormalizeQuerySource) so every SQL path rejects
an invalid value with SQLSTATE 22023 before anything is stored on the
session: 'invalid value for "duckgres.query_source": must be
"standard" or "endpoints"'. Matching is case-insensitive and
normalized to lowercase; empty resets to the default (unchanged
semantics); a SET whose value is not a single simple constant/identifier
is rejected instead of silently resetting. An invalid startup option
rejects the connection with FATAL 22023, matching the duckgres.worker_*
startup options. The rejection message deliberately does not echo the
offending value (it is arbitrary client input and error text lands in
logs and the admin recent-errors ring).

Defense in depth: server.ConnectionBilling clamps a non-canonical value
to "standard" at the metering boundary (warn log, value length only),
so a future validation bypass still cannot write junk into the billing
bucket key.

Author: benben

CLAUDE.md

@@ -470,7 +470,15 @@ admin APIs, never clearable via the admin API, so every org always has one;
 resolved from the config snapshot at connection end, where 0 can still appear
 only for an unknown org / stale snapshot); `query_source` is the
 `duckgres.query_source` session GUC (`standard` unless set; a mid-connection
-change bills the whole connection under the final value). Invariants for anyone
+change bills the whole connection under the final value). The GUC is a **closed
+enum validated at SET time** (`transform.NormalizeQuerySource`): only
+`standard` | `endpoints` (case-insensitive, normalized to lowercase; empty =
+reset to default) — anything else is rejected with `22023` on every set path
+(simple/batched SET, extended Parse, and the `-c` startup option, which rejects
+the connection like invalid `duckgres.worker_*` options), and
+`server.ConnectionBilling` clamps a non-canonical value to `standard` as
+defense in depth so client junk can never become a billing bucket key.
+Invariants for anyone
 touching this path:
 
 - **Metering is strictly best-effort and off the hot path.** A metering error

docs/design/billing-pull-api.md

@@ -43,7 +43,17 @@ Internally, one row per unique key, values accumulated:
   org's default team. `query_source` (`standard` | `endpoints`) is set by a session
   GUC (`duckgres.query_source`), defaulting to `standard` when unset; the meter
   reads it per connection. (If it's changed mid-connection the per-connection meter
-  uses a single value — same future refinement as team.)
+  uses a single value — same future refinement as team.) The GUC is **validated at
+  SET time** against the closed `{standard, endpoints}` set — case-insensitive,
+  normalized to lowercase; empty resets to the default; any other value is
+  rejected with SQLSTATE `22023` (`invalid value for "duckgres.query_source":
+  must be "standard" or "endpoints"`) on every set path (simple SET, batched,
+  extended-protocol Parse). An invalid `-c duckgres.query_source=…` startup
+  option rejects the connection (FATAL `22023`), matching the
+  `duckgres.worker_*` startup options. As defense in depth,
+  `server.ConnectionBilling` clamps a non-canonical value to `standard` at the
+  metering boundary, so client input can never put unbounded-cardinality junk
+  in the bucket key or the billing export.
 - Worker size (`cpu`, `mem_gib`) is part of the key, so different worker sizes
   accumulate — and bill — separately. Units match the values (vCPU / GiB).
 
@@ -167,7 +177,8 @@ sizes; stored as `NUMERIC` so grouping is exact.)
   at org creation on both the provisioning and admin APIs, and not clearable
   via the admin API; used as the bucket `team_id` — fixed per org, no per-team
   logic yet); a `duckgres.query_source` session GUC
-  (`standard` | `endpoints`, default `standard`) read by the meter; a bearer secret
+  (`standard` | `endpoints`, default `standard`, validated at SET time — anything
+  else is a `22023` error) read by the meter; a bearer secret
   for the API.
 
 ## Storage metric

server/conn.go

@@ -27,6 +27,7 @@ import (
 	"github.com/posthog/duckgres/server/usersecrets"
 	"github.com/posthog/duckgres/server/wire"
 	"github.com/posthog/duckgres/transpiler"
+	"github.com/posthog/duckgres/transpiler/transform"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -196,7 +197,9 @@ type clientConn struct {
 	// This is a prerequisite for pull-based compute billing (usage buckets are
 	// keyed by query_source). Set via `SET duckgres.query_source = '...'` (or a
 	// `-c duckgres.query_source=...` startup option) and read back via
-	// `SHOW duckgres.query_source` / current_setting.
+	// `SHOW duckgres.query_source` / current_setting. Every set path validates
+	// against the closed {standard, endpoints} set (22023 on anything else), so
+	// this only ever holds "", "standard", or "endpoints".
 	querySource string
 
 	// Provisioned worker pod size for compute-usage billing (remote/k8s backend
@@ -265,13 +268,31 @@ func (c *clientConn) QuerySource() string {
 }
 
 // setQuerySource records a client-supplied `duckgres.query_source` value on the
-// session. Pass-through: any string is accepted (only "standard"/"endpoints" are
-// meaningful downstream, but other values are stored, not rejected). An empty
-// value resets to the default.
+// session. Callers must pass an already-validated canonical value ("standard",
+// "endpoints", or "" = reset to default): the transpiler validates SET
+// statements (transform.NormalizeQuerySource, rejecting anything else with
+// 22023 before this is reached) and applyStartupQuerySource validates the
+// startup option. The setter itself stays dumb; ConnectionBilling additionally
+// clamps at the metering boundary as defense in depth.
 func (c *clientConn) setQuerySource(value string) {
 	c.querySource = value
 }
 
+// applyStartupQuerySource validates and applies a `-c duckgres.query_source=…`
+// startup-option value. The GUC is a billing dimension, so the value must be
+// in the closed {standard, endpoints} set (case-insensitive, normalized to
+// lowercase; empty = default). An invalid value returns the 22023 error for
+// the caller to reject the connection with, mirroring how the control plane
+// rejects invalid duckgres.worker_* startup options.
+func (c *clientConn) applyStartupQuerySource(raw string) error {
+	norm, err := transform.NormalizeQuerySource(raw)
+	if err != nil {
+		return err
+	}
+	c.setQuerySource(norm)
+	return nil
+}
+
 // generateSecretKey is a thin alias for wire.GenerateSecretKey kept for the
 // internal call sites in this file. New code should call wire.GenerateSecretKey
 // directly.
@@ -1017,9 +1038,16 @@ func (c *clientConn) handleStartup() error {
 
 		// Honor a `-c duckgres.query_source=...` startup option (libpq `options`
 		// keyword / PGOPTIONS). Other GUCs in `options` are not applied here.
+		// An invalid value rejects the connection with FATAL 22023 — the same
+		// treatment resolveWorkerProfile gives an invalid duckgres.worker_*
+		// startup option, and what PostgreSQL itself does with an invalid
+		// `options` GUC value.
 		if opts := ParseStartupOptions(params["options"]); len(opts) > 0 {
 			if v, ok := opts[querySourceGUCName]; ok {
-				c.setQuerySource(v)
+				if err := c.applyStartupQuerySource(v); err != nil {
+					c.sendError("FATAL", "22023", err.Error())
+					return fmt.Errorf("invalid %s startup option", querySourceGUCName)
+				}
 			}
 		}
 

server/exports.go

@@ -10,6 +10,7 @@ import (
 	"github.com/posthog/duckgres/server/observe"
 	"github.com/posthog/duckgres/server/sessionmeta"
 	"github.com/posthog/duckgres/server/wire"
+	"github.com/posthog/duckgres/transpiler/transform"
 )
 
 // Exported wrappers for protocol functions used by the control plane worker.
@@ -213,11 +214,25 @@ func SetConnectionWorkerSize(cc *clientConn, millicores, mib int64) {
 // size). Call at the same teardown point as CloseConnectionMetrics. A
 // mid-connection GUC change is not split: the whole connection is metered
 // under the final value (documented in docs/design/billing-pull-api.md).
+//
+// The query source is clamped to the closed {standard, endpoints} set as
+// defense in depth: every set path already validates (22023 at SET / startup
+// time), so a non-canonical value here means a validation bypass — degrade it
+// to the default rather than writing unbounded-cardinality client input into
+// the billing bucket key (and onward into billing exports).
 func ConnectionBilling(cc *clientConn) (orgID, querySource string, millicores, mib int64, dur time.Duration) {
 	if cc == nil {
 		return "", "", 0, 0, 0
 	}
-	return cc.orgID, cc.QuerySource(), cc.workerMillicores, cc.workerMiB, time.Since(cc.backendStart)
+	qs := cc.QuerySource()
+	if qs != transform.QuerySourceStandard && qs != transform.QuerySourceEndpoints {
+		// Should be unreachable. Log the length, not the value — it is
+		// arbitrary client input.
+		cc.logger().Warn("Non-canonical duckgres.query_source at billing teardown; degrading to default.",
+			"value_len", len(qs))
+		qs = defaultQuerySource
+	}
+	return cc.orgID, qs, cc.workerMillicores, cc.workerMiB, time.Since(cc.backendStart)
 }
 
 // CancelClientConn cancels the context of a clientConn.