pgsql-broker/tests/integration/rls_test.go

package integration

import (
	"context"
	"database/sql"
	"fmt"
	"testing"

	_ "github.com/lib/pq"
	"github.com/stretchr/testify/require"
)

// TestRLSTenantIsolation verifies that a non-superuser role without
// BYPASSRLS, connected via broker_set_tenant, only ever sees jobs and
// dependency rows for its own tenant -- the core guarantee behind the
// broker_jobs/broker_job_dependency FORCE ROW LEVEL SECURITY policies.
func TestRLSTenantIsolation(t *testing.T) {
	ctx := context.Background()
	adminDB := setupStage5Schema(t)

	// A restricted role, no BYPASSRLS, mirroring sql/roles/0001_roles.sql's
	// broker_runtime (superuser test DB roles already bypass RLS entirely,
	// so this test would be meaningless against the "user" role).
	_, err := adminDB.Exec("DROP ROLE IF EXISTS test_broker_runtime")
	require.NoError(t, err)
	_, err = adminDB.Exec("CREATE ROLE test_broker_runtime LOGIN PASSWORD 'test-pass' NOSUPERUSER NOBYPASSRLS")
	require.NoError(t, err)
	t.Cleanup(func() {
		if _, err := adminDB.Exec("REASSIGN OWNED BY test_broker_runtime TO CURRENT_USER"); err != nil {
			t.Logf("warning: failed to reassign objects owned by test_broker_runtime: %v", err)
		}
		if _, err := adminDB.Exec("DROP OWNED BY test_broker_runtime"); err != nil {
			t.Logf("warning: failed to drop grants owned by test_broker_runtime: %v", err)
		}
		if _, err := adminDB.Exec("DROP ROLE IF EXISTS test_broker_runtime"); err != nil {
			t.Logf("warning: failed to drop role test_broker_runtime: %v", err)
		}
	})

	for _, stmt := range []string{
		"GRANT USAGE ON SCHEMA broker TO test_broker_runtime",
		"GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA broker TO test_broker_runtime",
		"GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA broker TO test_broker_runtime",
		"GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA broker TO test_broker_runtime",
	} {
		_, err = adminDB.Exec(stmt)
		require.NoError(t, err)
	}

	runtimeDB, err := sql.Open("postgres",
		fmt.Sprintf("user=test_broker_runtime password=test-pass dbname=broker_test host=%s port=%d sslmode=disable options='-c search_path=broker,public'", testDBHost(), testDBPort()))
	require.NoError(t, err)
	defer runtimeDB.Close()
	require.NoError(t, runtimeDB.Ping())

	// broker_set_tenant uses SET LOCAL semantics (set_config(..., true)), so
	// it only takes effect for the remainder of the transaction it runs in.
	// Callers must set the tenant and perform the tenant-scoped operation in
	// the same explicit transaction (as worker.go's processJobs does) -- two
	// separate autocommitted statements would each run in their own
	// transaction and the tenant setting would not carry over.
	addJobAsTenant := func(tenant, name string) int64 {
		conn, err := runtimeDB.Conn(ctx)
		require.NoError(t, err)
		defer conn.Close()

		tx, err := conn.BeginTx(ctx, nil)
		require.NoError(t, err)

		_, err = tx.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", tenant)
		require.NoError(t, err)

		var retval int
		var errmsg string
		var jobID int64
		err = tx.QueryRowContext(ctx, `
			SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_add_job(
				$1, 'SELECT 1', 1, 0, 'sql', NULL, NULL, NULL, NULL, 1
			)`, name,
		).Scan(&retval, &errmsg, &jobID)
		require.NoError(t, err)
		require.Equal(t, 0, retval, errmsg)

		require.NoError(t, tx.Commit())
		return jobID
	}

	tenantAJob := addJobAsTenant("tenant-a", "tenant-a-job")
	tenantBJob := addJobAsTenant("tenant-b", "tenant-b-job")
	require.NotEqual(t, tenantAJob, tenantBJob)

	// As tenant-a, only tenant-a's job must be visible. All of tenant-a's
	// checks (including the broker_get claim below) share one explicit
	// transaction so the SET LOCAL tenant context stays in effect throughout.
	connA, err := runtimeDB.Conn(ctx)
	require.NoError(t, err)
	defer connA.Close()
	txA, err := connA.BeginTx(ctx, nil)
	require.NoError(t, err)
	defer txA.Rollback()

	_, err = txA.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", "tenant-a")
	require.NoError(t, err)

	var visibleCount int
	err = txA.QueryRowContext(ctx, "SELECT COUNT(*) FROM broker.broker_jobs WHERE id_broker_jobs IN ($1, $2)",
		tenantAJob, tenantBJob).Scan(&visibleCount)
	require.NoError(t, err)
	require.Equal(t, 1, visibleCount, "tenant-a must see only its own job, not tenant-b's")

	var visibleName string
	err = txA.QueryRowContext(ctx, "SELECT job_name FROM broker.broker_jobs WHERE id_broker_jobs = $1", tenantAJob).Scan(&visibleName)
	require.NoError(t, err)
	require.Equal(t, "tenant-a-job", visibleName)

	// As tenant-b, only tenant-b's job must be visible.
	connB, err := runtimeDB.Conn(ctx)
	require.NoError(t, err)
	defer connB.Close()
	txB, err := connB.BeginTx(ctx, nil)
	require.NoError(t, err)
	defer txB.Rollback()

	_, err = txB.ExecContext(ctx, "SELECT broker.broker_set_tenant($1)", "tenant-b")
	require.NoError(t, err)

	err = txB.QueryRowContext(ctx, "SELECT COUNT(*) FROM broker.broker_jobs WHERE id_broker_jobs IN ($1, $2)",
		tenantAJob, tenantBJob).Scan(&visibleCount)
	require.NoError(t, err)
	require.Equal(t, 1, visibleCount, "tenant-b must see only its own job, not tenant-a's")
	require.NoError(t, txB.Commit())

	// broker_get run under tenant-a's context must never be able to claim
	// tenant-b's job.
	var claimedID sql.NullInt64
	var getRetval int
	var getErrmsg string
	err = txA.QueryRowContext(ctx,
		"SELECT p_retval, p_errmsg, p_job_id FROM broker.broker_get($1, NULL, $2)", 1, 60,
	).Scan(&getRetval, &getErrmsg, &claimedID)
	require.NoError(t, err)
	require.Equal(t, 0, getRetval, getErrmsg)
	require.True(t, claimedID.Valid)
	require.Equal(t, tenantAJob, claimedID.Int64, "tenant-a's broker_get must only ever claim tenant-a's own job")
	require.NoError(t, txA.Commit())
}