SQL Injection

Proof-of-concept (conceptual, executable steps below): Background: Prior to this patch, the code interpolated the resource string directly into a partition name (embeddings_<resource>) and into a DDL statement such as CREATE TABLE IF NOT EXISTS embeddings_<resource> PARTITION OF <parent> FOR VALUES IN ('<resource>'). If an attacker supplied a crafted resource containing SQL-control characters, they could terminate the DDL early and inject additional SQL, potentially dropping tables or altering schema dependencies. Malicious input example (vulnerable path): - resource = "evil'); DROP TABLE vector_backfill_jobs;--" - This would yield a generated SQL similar to: CREATE TABLE IF NOT EXISTS embeddings_evil'); DROP TABLE vector_backfill_jobs;-- PARTITION OF embeddings_parent FOR VALUES IN ('evil') Expected impact: The injected DROP TABLE vector_backfill_jobs; would be executed as part of the same statement sequence, causing unintended destructive operation if the database processes multiple statements coming from a single ExecContext call. With the fix: The code now rejects unsafe resource inputs via sanitizeIdentifier(resource) and returns an error before any DDL is constructed. Safe inputs are strictly validated (e.g., non-empty, sanitized form), and only then are partitions created with validated leaf names. The patch also adds tests that assert unsafe inputs are rejected and that resource partitions are created safely. If you want to experiment locally, you can simulate the vulnerability by invoking EnsureResourcePartition with the crafted resource and observing that the call is rejected on patched versions, while it would proceed (and allow SQL injection) on pre-patch versions.

diff --git a/pkg/storage/unified/resource/search_vector_test.go b/pkg/storage/unified/resource/search_vector_test.go index 1ee98b006c7b4..103aad327327a 100644 --- a/pkg/storage/unified/resource/search_vector_test.go +++ b/pkg/storage/unified/resource/search_vector_test.go @@ -103,6 +103,10 @@ func (f *fakeVectorBackend) TryAcquireReconcilerLock(context.Context) (func(), b func (f *fakeVectorBackend) ListIncompleteBackfillJobs(context.Context, string) ([]vector.BackfillJob, error) { return nil, nil } +func (f *fakeVectorBackend) EnsureResourcePartition(context.Context, string) error { return nil } +func (f *fakeVectorBackend) CreateBackfillJob(_ context.Context, _, _ string, _ int64) error { + return nil +} func (f *fakeVectorBackend) UpdateBackfillJobCheckpoint(context.Context, int64, string, string) error { return nil } diff --git a/pkg/storage/unified/search/embed/backfill/fakes_test.go b/pkg/storage/unified/search/embed/backfill/fakes_test.go index 648d2af6ba8ba..beb77314c6aba 100644 --- a/pkg/storage/unified/search/embed/backfill/fakes_test.go +++ b/pkg/storage/unified/search/embed/backfill/fakes_test.go @@ -270,6 +270,10 @@ func (f *fakeVector) SetLatestRV(context.Context, int64) error { return nil } func (f *fakeVector) TryAcquireReconcilerLock(context.Context) (func(), bool, error) { return func() {}, true, nil } +func (f *fakeVector) EnsureResourcePartition(context.Context, string) error { return nil } +func (f *fakeVector) CreateBackfillJob(_ context.Context, _, _ string, _ int64) error { + return nil +} func (f *fakeVector) ListIncompleteBackfillJobs(_ context.Context, model string) ([]vector.BackfillJob, error) { f.mu.Lock() defer f.mu.Unlock() diff --git a/pkg/storage/unified/search/embed/reconciler/fakes_test.go b/pkg/storage/unified/search/embed/reconciler/fakes_test.go index 3fb3f5441f495..08b4f867ec2b4 100644 --- a/pkg/storage/unified/search/embed/reconciler/fakes_test.go +++ b/pkg/storage/unified/search/embed/reconciler/fakes_test.go @@ -307,6 +307,10 @@ func (f *fakeVector) SetLatestRV(_ context.Context, rv int64) error { func (f *fakeVector) ListIncompleteBackfillJobs(context.Context, string) ([]vector.BackfillJob, error) { return nil, nil } +func (f *fakeVector) EnsureResourcePartition(context.Context, string) error { return nil } +func (f *fakeVector) CreateBackfillJob(context.Context, string, string, int64) error { + return nil +} func (f *fakeVector) UpdateBackfillJobCheckpoint(context.Context, int64, string, string) error { return nil } diff --git a/pkg/storage/unified/search/vector/data/vector_backfill_jobs_create.sql b/pkg/storage/unified/search/vector/data/vector_backfill_jobs_create.sql new file mode 100644 index 0000000000000..115050428e110 --- /dev/null +++ b/pkg/storage/unified/search/vector/data/vector_backfill_jobs_create.sql @@ -0,0 +1,5 @@ +INSERT INTO vector_backfill_jobs + ({{ .Ident "model" }}, {{ .Ident "resource" }}, {{ .Ident "stopping_rv" }}, {{ .Ident "is_complete" }}) + VALUES ({{ .Arg .Model }}, {{ .Arg .Resource }}, {{ .Arg .StoppingRV }}, FALSE) + ON CONFLICT ({{ .Ident "model" }}, {{ .Ident "resource" }}) DO NOTHING +; diff --git a/pkg/storage/unified/search/vector/integration_test.go b/pkg/storage/unified/search/vector/integration_test.go index 09c5ed1542a5c..43b7f1479f982 100644 --- a/pkg/storage/unified/search/vector/integration_test.go +++ b/pkg/storage/unified/search/vector/integration_test.go @@ -89,6 +89,8 @@ func cleanIntegrationState(t *testing.T, engine *xorm.Engine) { `DELETE FROM vector_promoted WHERE namespace LIKE 'integration-test%'`) _, _ = engine.DB().ExecContext(ctx, `UPDATE vector_latest_rv SET latest_rv = 0 WHERE id = 1`) + _, _ = engine.DB().ExecContext(ctx, + `DELETE FROM vector_backfill_jobs WHERE model = $1`, testModel) } func TestIntegrationVectorUpsertAndSearch(t *testing.T) { @@ -387,6 +389,21 @@ func TestIntegrationVectorGetLatestRV(t *testing.T) { assert.Equal(t, int64(100), rv) } +func TestIntegrationVectorCreateBackfillJob(t *testing.T) { + backend, _, ctx := setupIntegrationTest(t) + + require.NoError(t, backend.CreateBackfillJob(ctx, testModel, testResource, 100)) + + // Second insert for the same (model, resource) is a no-op (ON CONFLICT + // DO NOTHING): the original row is preserved, not overwritten with 200. + require.NoError(t, backend.CreateBackfillJob(ctx, testModel, testResource, 200)) + + jobs, err := backend.ListIncompleteBackfillJobs(ctx, testModel) + require.NoError(t, err) + require.Len(t, jobs, 1, "exactly one job exists after the conflicting insert") + assert.Equal(t, int64(100), jobs[0].StoppingRV, "original stopping_rv preserved") +} + func TestIntegrationVectorReconcilerLock(t *testing.T) { backend, _, ctx := setupIntegrationTest(t) @@ -481,3 +498,45 @@ func makeEmbedding(a, b float32) []float32 { e[1] = b return e } + +func TestEnsureResourcePartition_RejectsUnsafeResource(t *testing.T) { + b := &pgvectorBackend{} + for _, res := range []string{"", "Dashboards", "dash-boards", "a.b", "drop;table", "with space"} { + require.Error(t, b.EnsureResourcePartition(context.Background(), res), + "resource %q must be rejected", res) + } +} + +func TestIntegrationVectorEnsureResourcePartition(t *testing.T) { + backend, engine, ctx := setupIntegrationTest(t) + pg := backend.(*pgvectorBackend) + + const res = "testpartition" + leaf := subtreeName(res) + idx := leaf + "_metadata_idx" + drop := func() { _, _ = engine.DB().ExecContext(ctx, fmt.Sprintf(`DROP TABLE IF EXISTS %s`, leaf)) } + drop() + t.Cleanup(drop) + + // Absent before creation. + ready, err := pg.resourcePartitionReady(ctx, leaf, idx) + require.NoError(t, err) + require.False(t, ready) + + // Create it: partition + metadata index both present. + require.NoError(t, backend.EnsureResourcePartition(ctx, res)) + ready, err = pg.resourcePartitionReady(ctx, leaf, idx) + require.NoError(t, err) + assert.True(t, ready, "leaf attached as partition and metadata index present") + + // Heals a missing index: drop it, retry must recreate it. + _, err = engine.DB().ExecContext(ctx, fmt.Sprintf(`DROP INDEX IF EXISTS %s`, idx)) + require.NoError(t, err) + require.NoError(t, backend.EnsureResourcePartition(ctx, res)) + ready, err = pg.resourcePartitionReady(ctx, leaf, idx) + require.NoError(t, err) + assert.True(t, ready, "missing index recreated on retry") + + // Idempotent: a second call (fast path) is a no-op, no error. + require.NoError(t, backend.EnsureResourcePartition(ctx, res)) +} diff --git a/pkg/storage/unified/search/vector/pgvector.go b/pkg/storage/unified/search/vector/pgvector.go index 994e838ee6b9d..69f93a18ab75d 100644 --- a/pkg/storage/unified/search/vector/pgvector.go +++ b/pkg/storage/unified/search/vector/pgvector.go @@ -416,6 +416,95 @@ func (b *pgvectorBackend) ListIncompleteBackfillJobs(ctx context.Context, model return out, nil } +// EnsureResourcePartition creates the embeddings_<resource> partition leaf and +// its metadata index. The per-resource advisory lock serializes the attach +// (ACCESS EXCLUSIVE on the parent) so concurrent replicas don't race it. +func (b *pgvectorBackend) EnsureResourcePartition(ctx context.Context, resource string) error { + // resource is interpolated unquoted into the DDL below; reject anything + // sanitizeIdentifier would alter to keep it injection-safe. + if resource == "" || sanitizeIdentifier(resource) != resource { + return fmt.Errorf("ensure partition: unsafe resource %q", resource) + } + leaf := subtreeName(resource) // embeddings_<resource> + idx := leaf + "_metadata_idx" + + // Fast path: skip the lock + DDL only when BOTH leaf and index exist. + // Checking the index too lets a retry finish a prior attempt that created + // the leaf but failed before the index. + ready, err := b.resourcePartitionReady(ctx, leaf, idx) + if err != nil { + return fmt.Errorf("check partition %s: %w", leaf, err) + } + if ready { + return nil + } + + conn, err := b.db.SqlDB().Conn(ctx) + if err != nil { + return fmt.Errorf("ensure partition conn: %w", err) + } + defer func() { _ = conn.Close() }() + + lockName := "vectorpartition_" + resource + if _, err := conn.ExecContext(ctx, + "SELECT pg_advisory_lock(hashtext($1)::bigint)", lockName); err != nil { + return fmt.Errorf("ensure partition lock: %w", err) + } + defer func() { + _, _ = conn.ExecContext(context.Background(), + "SELECT pg_advisory_unlock(hashtext($1)::bigint)", lockName) + }() + + if _, err := conn.ExecContext(ctx, fmt.Sprintf( + `CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES IN ('%s')`, + leaf, unifiedParent, resource, + )); err != nil { + return fmt.Errorf("create partition %s: %w", leaf, err) + } + if _, err := conn.ExecContext(ctx, fmt.Sprintf( + `CREATE INDEX IF NOT EXISTS %s ON %s USING GIN (metadata)`, + idx, leaf, + )); err != nil { + return fmt.Errorf("create metadata index on %s: %w", leaf, err) + } + return nil +} + +// resourcePartitionReady reports whether leaf is attached as a partition of +// the parent (pg_inherits, not to_regclass, so a same-named unrelated table +// can't match) AND its metadata index exists. +func (b *pgvectorBackend) resourcePartitionReady(ctx context.Context, leaf, idx string) (bool, error) { + var ready bool + err := b.db.QueryRowContext(ctx, ` + SELECT + EXISTS ( + SELECT 1 FROM pg_inherits i + JOIN pg_class c ON c.oid = i.inhrelid + JOIN pg_class p ON p.oid = i.inhparent + WHERE p.relname = $1 AND c.relname = $2 + ) + AND EXISTS ( + SELECT 1 FROM pg_class WHERE relname = $3 AND relkind = 'i' + )`, unifiedParent, leaf, idx).Scan(&ready) + if err != nil { + return false, err + } + return ready, nil +} + +func (b *pgvectorBackend) CreateBackfillJob(ctx context.Context, model, resource string, stoppingRV int64) error { + req := &sqlVectorBackfillJobsCreateRequest{ + SQLTemplate: sqltemplate.New(b.dialect), + Model: model, + Resource: resource, + StoppingRV: stoppingRV, + } + if _, err := dbutil.Exec(ctx, b.db, sqlVectorBackfillJobsCreate, req); err != nil { + return fmt.Errorf("create backfill job (%s,%s): %w", model, resource, err) + } + return nil +} + func (b *pgvectorBackend) UpdateBackfillJobCheckpoint(ctx context.Context, id int64, lastSeenKey string, lastErr string) error { req := &sqlVectorBackfillJobsUpdateRequest{ SQLTemplate: sqltemplate.New(b.dialect), diff --git a/pkg/storage/unified/search/vector/queries.go b/pkg/storage/unified/search/vector/queries.go index 6c1b4419ba436..cfd059b752a85 100644 --- a/pkg/storage/unified/search/vector/queries.go +++ b/pkg/storage/unified/search/vector/queries.go @@ -36,6 +36,7 @@ var ( sqlVectorCollectionExists = mustTemplate("vector_collection_exists.sql") sqlVectorCollectionSearch = mustTemplate("vector_collection_search.sql") sqlVectorBackfillJobsList = mustTemplate("vector_backfill_jobs_list.sql") + sqlVectorBackfillJobsCreate = mustTemplate("vector_backfill_jobs_create.sql") sqlVectorBackfillJobsUpdate = mustTemplate("vector_backfill_jobs_update.sql") sqlVectorBackfillJobsSetError = mustTemplate("vector_backfill_jobs_set_error.sql") sqlVectorBackfillJobsComplete = mustTemplate("vector_backfill_jobs_complete.sql") @@ -158,6 +159,26 @@ func (r *sqlVectorBackfillJobsListRequest) Results() (*sqlVectorBackfillJobsList return &cp, nil } +type sqlVectorBackfillJobsCreateRequest struct { + sqltemplate.SQLTemplate + Model string + Resource string + StoppingRV int64 +} + +func (r *sqlVectorBackfillJobsCreateRequest) Validate() error { + if r.Model == "" { + return fmt.Errorf("missing model") + } + if r.Resource == "" { + return fmt.Errorf("missing resource") + } + if r.StoppingRV <= 0 { + return fmt.Errorf("stopping_rv must be positive") + } + return nil +} + type sqlVectorBackfillJobsUpdateRequest struct { sqltemplate.SQLTemplate ID int64 diff --git a/pkg/storage/unified/search/vector/queries_test.go b/pkg/storage/unified/search/vector/queries_test.go index d27f1d0123717..c063732717e9f 100644 --- a/pkg/storage/unified/search/vector/queries_test.go +++ b/pkg/storage/unified/search/vector/queries_test.go @@ -103,6 +103,17 @@ func TestVectorQueries(t *testing.T) { }, }, }, + sqlVectorBackfillJobsCreate: { + { + Name: "simple", + Data: &sqlVectorBackfillJobsCreateRequest{ + SQLTemplate: mocks.NewTestingSQLTemplate(), + Model: "text-embedding-005", + Resource: "dashboards", + StoppingRV: 12345, + }, + }, + }, sqlVectorBackfillJobsUpdate: { { Name: "simple", diff --git a/pkg/storage/unified/search/vector/store.go b/pkg/storage/unified/search/vector/store.go index 95bcdb659d064..02ebcbdbc3ee6 100644 --- a/pkg/storage/unified/search/vector/store.go +++ b/pkg/storage/unified/search/vector/store.go @@ -75,6 +75,13 @@ type VectorBackend interface { // own. Operators add rows via SQL migrations; the resource embedder drains them. ListIncompleteBackfillJobs(ctx context.Context, model string) ([]BackfillJob, error) + // EnsureResourcePartition creates the embeddings_<resource> partition leaf (idempotent). + EnsureResourcePartition(ctx context.Context, resource string) error + + // CreateBackfillJob creates a backfill job for (model, resource, stoppingRV). + // No-op if a job already exists for (model, resource). + CreateBackfillJob(ctx context.Context, model, resource string, stoppingRV int64) error + // UpdateBackfillJobCheckpoint writes the cursor + optional error after // each processed resource. Best-effort — race with another writer is // acceptable since the resource embedder is single-goroutine. diff --git a/pkg/storage/unified/search/vector/testdata/postgres--vector_backfill_jobs_create-simple.sql b/pkg/storage/unified/search/vector/testdata/postgres--vector_backfill_jobs_create-simple.sql new file mode 100644 index 0000000000000..a85e0cd1da4df --- /dev/null +++ b/pkg/storage/unified/search/vector/testdata/postgres--vector_backfill_jobs_create-simple.sql @@ -0,0 +1,5 @@ +INSERT INTO vector_backfill_jobs + ("model", "resource", "stopping_rv", "is_complete") + VALUES ('text-embedding-005', 'dashboards', 12345, FALSE) + ON CONFLICT ("model", "resource") DO NOTHING +;