forked from wangziqi/gongxue-base
316 lines
10 KiB
JavaScript
316 lines
10 KiB
JavaScript
import pg from 'pg';
|
||
|
||
const { Pool } = pg;
|
||
|
||
const databaseUrl = process.env.DATABASE_URL || 'postgresql://postgres:postgres@127.0.0.1:54322/postgres';
|
||
|
||
const ids = {
|
||
mainTenant: '00000000-0000-0000-0000-000000000001',
|
||
partnerTenant: '00000000-0000-0000-0000-000000000901',
|
||
mainUser: '00000000-0000-0000-0000-000000000101',
|
||
partnerAdminUser: '00000000-0000-0000-0000-000000000907',
|
||
};
|
||
|
||
const readChecks = [
|
||
{
|
||
name: 'tenant_branding',
|
||
sql: 'select tenant_id::text as tenant_id, brand_name as label from public.tenant_branding order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'tenant_settings',
|
||
sql: 'select tenant_id::text as tenant_id, public_config::text as label from public.tenant_settings order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'tenant_domains',
|
||
sql: 'select tenant_id::text as tenant_id, host as label from public.tenant_domains order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'tenant_memberships',
|
||
sql: 'select tenant_id::text as tenant_id, role as label from public.tenant_memberships order by tenant_id, role',
|
||
expectMain: true,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'regions',
|
||
sql: 'select tenant_id::text as tenant_id, name as label from public.regions order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'questions',
|
||
sql: 'select tenant_id::text as tenant_id, legacy_id as label from public.questions order by tenant_id, id',
|
||
expectMain: true,
|
||
expectPartner: false,
|
||
},
|
||
{
|
||
name: 'student_profiles',
|
||
sql: 'select tenant_id::text as tenant_id, user_id::text as label from public.student_profiles order by tenant_id, user_id',
|
||
expectMain: true,
|
||
expectPartner: false,
|
||
},
|
||
{
|
||
name: 'orders',
|
||
sql: 'select tenant_id::text as tenant_id, order_no as label from public.orders order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: false,
|
||
},
|
||
{
|
||
name: 'content_assets',
|
||
sql: 'select tenant_id::text as tenant_id, asset_key as label from public.content_assets order by tenant_id',
|
||
expectMain: true,
|
||
expectPartner: false,
|
||
},
|
||
{
|
||
name: 'tenant_subscriptions',
|
||
sql: 'select tenant_id::text as tenant_id, plan_code as label from public.tenant_subscriptions order by tenant_id',
|
||
expectMain: false,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'tenant_invoices',
|
||
sql: 'select tenant_id::text as tenant_id, invoice_no as label from public.tenant_invoices order by tenant_id',
|
||
expectMain: false,
|
||
expectPartner: true,
|
||
},
|
||
{
|
||
name: 'tenant_usage_records',
|
||
sql: 'select tenant_id::text as tenant_id, metric_key as label from public.tenant_usage_records order by tenant_id, metric_key',
|
||
expectMain: false,
|
||
expectPartner: true,
|
||
},
|
||
];
|
||
|
||
const writeChecks = [
|
||
{
|
||
name: 'tenant_branding cross-tenant update',
|
||
setupRole: 'authenticated',
|
||
tenantId: ids.mainTenant,
|
||
roleClaim: 'authenticated',
|
||
sql: `
|
||
update public.tenant_branding
|
||
set slogan = slogan
|
||
where tenant_id = $1::uuid
|
||
returning tenant_id::text
|
||
`,
|
||
params: [ids.partnerTenant],
|
||
},
|
||
{
|
||
name: 'regions cross-tenant insert',
|
||
setupRole: 'authenticated',
|
||
tenantId: ids.mainTenant,
|
||
roleClaim: 'authenticated',
|
||
sql: `
|
||
insert into public.regions (tenant_id, legacy_id, name, code, sort_order, is_active)
|
||
values ($1::uuid, 'rls-cross-tenant-probe', 'RLS 越权探针', 'RLS-X', 9999, false)
|
||
returning tenant_id::text
|
||
`,
|
||
params: [ids.partnerTenant],
|
||
expectErrorCode: '42501',
|
||
},
|
||
];
|
||
|
||
const probeGrantStatements = [
|
||
'grant usage on schema app to authenticated, anon',
|
||
'grant usage on schema public to authenticated, anon',
|
||
`grant select on
|
||
public.tenant_branding,
|
||
public.tenant_settings,
|
||
public.tenant_domains,
|
||
public.tenant_memberships,
|
||
public.regions,
|
||
public.questions,
|
||
public.student_profiles,
|
||
public.orders,
|
||
public.content_assets,
|
||
public.tenant_subscriptions,
|
||
public.tenant_invoices,
|
||
public.tenant_usage_records
|
||
to authenticated, anon`,
|
||
'grant update on public.tenant_branding to authenticated',
|
||
'grant insert on public.regions to authenticated',
|
||
];
|
||
|
||
const pool = new Pool({ connectionString: databaseUrl });
|
||
const results = [];
|
||
|
||
function pass(name, detail = {}) {
|
||
results.push({ status: 'pass', name, detail });
|
||
console.log(`PASS ${name}`);
|
||
}
|
||
|
||
function fail(name, message, detail = {}) {
|
||
results.push({ status: 'fail', name, message, detail });
|
||
console.error(`FAIL ${name}: ${message}`);
|
||
if (Object.keys(detail).length > 0) console.error(JSON.stringify(detail, null, 2));
|
||
}
|
||
|
||
function assert(condition, name, message, detail = {}) {
|
||
if (condition) pass(name, detail);
|
||
else fail(name, message, detail);
|
||
}
|
||
|
||
function onlyTenantRows(rows, tenantId) {
|
||
return rows.every(row => row.tenant_id === tenantId);
|
||
}
|
||
|
||
function hasTenantRows(rows, tenantId) {
|
||
return rows.some(row => row.tenant_id === tenantId);
|
||
}
|
||
|
||
async function withRlsContext(client, { dbRole = 'authenticated', tenantId = '', roleClaim = 'authenticated', sub = '' }, action) {
|
||
await client.query('begin');
|
||
try {
|
||
for (const statement of probeGrantStatements) await client.query(statement);
|
||
await client.query(`set local role ${dbRole}`);
|
||
if (tenantId) {
|
||
await client.query("select set_config('request.jwt.claim.tenant_id', $1, true)", [tenantId]);
|
||
} else {
|
||
await client.query("select set_config('request.jwt.claim.tenant_id', '', true)");
|
||
}
|
||
await client.query("select set_config('request.jwt.claim.role', $1, true)", [roleClaim]);
|
||
await client.query("select set_config('request.jwt.claim.app_role', $1, true)", [roleClaim]);
|
||
if (sub) await client.query("select set_config('request.jwt.claim.sub', $1, true)", [sub]);
|
||
|
||
const value = await action();
|
||
await client.query('rollback');
|
||
return value;
|
||
} catch (error) {
|
||
await client.query('rollback').catch(() => {});
|
||
throw error;
|
||
}
|
||
}
|
||
|
||
async function queryAs(client, context, sql, params = []) {
|
||
return withRlsContext(client, context, async () => {
|
||
const result = await client.query(sql, params);
|
||
return result.rows;
|
||
});
|
||
}
|
||
|
||
async function runReadIsolationChecks(client) {
|
||
for (const check of readChecks) {
|
||
const mainRows = await queryAs(client, { tenantId: ids.mainTenant }, check.sql);
|
||
assert(
|
||
onlyTenantRows(mainRows, ids.mainTenant),
|
||
`rls.read.${check.name}.main_no_leak`,
|
||
'主租户上下文不应看到其它租户数据',
|
||
{ rows: mainRows },
|
||
);
|
||
assert(
|
||
hasTenantRows(mainRows, ids.mainTenant) === check.expectMain,
|
||
`rls.read.${check.name}.main_expected_seed`,
|
||
'主租户 seed 数据存在性不符合预期',
|
||
{ expected: check.expectMain, rows: mainRows },
|
||
);
|
||
|
||
const partnerRows = await queryAs(client, { tenantId: ids.partnerTenant }, check.sql);
|
||
assert(
|
||
onlyTenantRows(partnerRows, ids.partnerTenant),
|
||
`rls.read.${check.name}.partner_no_leak`,
|
||
'伙伴租户上下文不应看到其它租户数据',
|
||
{ rows: partnerRows },
|
||
);
|
||
assert(
|
||
hasTenantRows(partnerRows, ids.partnerTenant) === check.expectPartner,
|
||
`rls.read.${check.name}.partner_expected_seed`,
|
||
'伙伴租户 seed 数据存在性不符合预期',
|
||
{ expected: check.expectPartner, rows: partnerRows },
|
||
);
|
||
|
||
const anonymousRows = await queryAs(client, { dbRole: 'anon', tenantId: '', roleClaim: 'anon' }, check.sql);
|
||
assert(
|
||
anonymousRows.length === 0,
|
||
`rls.read.${check.name}.no_tenant_claim_empty`,
|
||
'没有 tenant_id claim 的上下文不应看到租户数据',
|
||
{ rows: anonymousRows },
|
||
);
|
||
}
|
||
}
|
||
|
||
async function runPlatformAdminChecks(client) {
|
||
for (const check of readChecks) {
|
||
const rows = await queryAs(client, { tenantId: '', roleClaim: 'platform_admin', sub: ids.mainUser }, check.sql);
|
||
const hasMain = hasTenantRows(rows, ids.mainTenant);
|
||
const hasPartner = hasTenantRows(rows, ids.partnerTenant);
|
||
assert(
|
||
hasMain === check.expectMain && hasPartner === check.expectPartner,
|
||
`rls.platform_admin.${check.name}.cross_tenant_visibility`,
|
||
'平台管理员 RLS 旁路应只暴露当前表已有的多租户 seed 数据',
|
||
{ expectedMain: check.expectMain, expectedPartner: check.expectPartner, rows },
|
||
);
|
||
}
|
||
}
|
||
|
||
async function runWriteIsolationChecks(client) {
|
||
for (const check of writeChecks) {
|
||
try {
|
||
const rows = await queryAs(
|
||
client,
|
||
{ dbRole: check.setupRole, tenantId: check.tenantId, roleClaim: check.roleClaim },
|
||
check.sql,
|
||
check.params,
|
||
);
|
||
|
||
assert(
|
||
rows.length === 0,
|
||
`rls.write.${check.name}`,
|
||
'跨租户写入不应成功返回任何行',
|
||
{ rows },
|
||
);
|
||
} catch (error) {
|
||
if (check.expectErrorCode && error.code === check.expectErrorCode) {
|
||
pass(`rls.write.${check.name}`, { expectedErrorCode: error.code });
|
||
} else {
|
||
fail(`rls.write.${check.name}`, error.message, { code: error.code });
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
async function verifySeed(client) {
|
||
const result = await client.query(
|
||
`
|
||
select tenant_id::text, count(*)::int as count
|
||
from public.tenant_branding
|
||
where tenant_id in ($1::uuid, $2::uuid)
|
||
group by tenant_id
|
||
`,
|
||
[ids.mainTenant, ids.partnerTenant],
|
||
);
|
||
const counts = new Map(result.rows.map(row => [row.tenant_id, Number(row.count)]));
|
||
assert(
|
||
counts.get(ids.mainTenant) === 1 && counts.get(ids.partnerTenant) === 1,
|
||
'rls.seed.main_and_partner_tenants',
|
||
'需要先运行 npm run db:smoke-seed,确保主租户和伙伴租户 seed 都存在',
|
||
{ rows: result.rows },
|
||
);
|
||
}
|
||
|
||
async function main() {
|
||
const client = await pool.connect();
|
||
try {
|
||
await verifySeed(client);
|
||
await runReadIsolationChecks(client);
|
||
await runPlatformAdminChecks(client);
|
||
await runWriteIsolationChecks(client);
|
||
} finally {
|
||
client.release();
|
||
await pool.end();
|
||
}
|
||
|
||
const failed = results.filter(item => item.status === 'fail');
|
||
console.log(`\nRLS tenant isolation checks: ${results.length - failed.length} passed, ${failed.length} failed.`);
|
||
if (failed.length > 0) process.exitCode = 1;
|
||
}
|
||
|
||
main().catch(error => {
|
||
console.error(error);
|
||
process.exitCode = 1;
|
||
});
|