fix: Increase test coverage and improve setup safety

This change introduces the use of advisory locks on startup to ensure
safety if multiple clients are bootstrapping simultaneously. Without the
lock you run the risk of concurrent updates to the same tables (mostly
when using `GRANT`).

Signed-off-by: Lucian Buzzo <lucian.buzzo@gmail.com>

prisma/migrations/20230125135410_add_tags_model/migration.sql

@@ -0,0 +1,25 @@
+-- CreateTable
+CREATE TABLE "Tag" (
+    "id" SERIAL NOT NULL,
+    "label" TEXT NOT NULL,
+
+    CONSTRAINT "Tag_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateTable
+CREATE TABLE "_PostToTag" (
+    "A" INTEGER NOT NULL,
+    "B" INTEGER NOT NULL
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "_PostToTag_AB_unique" ON "_PostToTag"("A", "B");
+
+-- CreateIndex
+CREATE INDEX "_PostToTag_B_index" ON "_PostToTag"("B");
+
+-- AddForeignKey
+ALTER TABLE "_PostToTag" ADD CONSTRAINT "_PostToTag_A_fkey" FOREIGN KEY ("A") REFERENCES "Post"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "_PostToTag" ADD CONSTRAINT "_PostToTag_B_fkey" FOREIGN KEY ("B") REFERENCES "Tag"("id") ON DELETE CASCADE ON UPDATE CASCADE;

prisma/schema.prisma

@@ -26,9 +26,16 @@ model Post {
   title     String   @db.VarChar(255)
   author    User?    @relation(fields: [authorId], references: [id])
   authorId  Int?
+  tags      Tag[]
+}
+
+model Tag {
+  id    Int    @id @default(autoincrement())
+  label String
+  posts Post[]
 }
 
 enum Role {
   USER
   ADMIN
-}
+}

src/index.ts

@@ -15,16 +15,25 @@ export type ModelAbilities = { [Model in Models]: { [op: string]: Ability } };
 
 export type GetContextFn = () => {
 	role: string;
-	context: {
+	context?: {
 		[key: string]: string;
 	};
 } | null;
 
+/**
+ * This function is used to take a lock that is automatically released at the end of the current transaction.
+ * This is very convenient for ensuring we don't hit concurrency issues when running setup code.
+ */
+const takeLock = (prisma: PrismaClient) =>
+	prisma.$executeRawUnsafe("SELECT pg_advisory_xact_lock(2142616474639426746);");
+
 export const createAbilityName = (model: string, ability: string) => {
 	return `${model}_${ability}_role`.toLowerCase();
 };
 
 export const createRoleName = (name: string) => {
+	// Esnure the role name only has lowercase alpha characters and underscores
+	// This also doubles as a check against SQL injection
 	const normalized = name.toLowerCase().replace("-", "_").replace(/[^a-z_]/g, "");
 	return `yates_role_${normalized}`;
 };
@@ -52,12 +61,6 @@ export const setupMiddleware = (prisma: PrismaClient, getContext: GetContextFn)
 
 		const pgRole = createRoleName(role);
 
-		// Check the role name only has lowercase alpha characters and underscores
-		// This also doubles as a check against SQL injection
-		if (pgRole.match(/[^a-z_]/)) {
-			throw new Error("Invalid role name.");
-		}
-
 		// Generate model class name from model params (PascalCase to camelCase)
 		const modelName = params.model.charAt(0).toLowerCase() + params.model.slice(1);
 
@@ -208,25 +211,33 @@ export const createRoles = async ({
 	for (const model in abilities) {
 		const table = model;
 
-		await prisma.$queryRawUnsafe(`ALTER table "${table}" enable row level security`);
+		await prisma.$transaction([
+			takeLock(prisma),
+			prisma.$queryRawUnsafe(`ALTER table "${table}" enable row level security;`),
+		]);
 
 		for (const slug in abilities[model as keyof typeof abilities]) {
 			const ability = abilities[model as keyof typeof abilities]![slug];
 			const roleName = createAbilityName(model, slug);
 
 			// Check if role already exists
-			await prisma.$queryRawUnsafe(`
-        do
-        $$
-        begin
-        if not exists (select * from pg_catalog.pg_roles where rolname = '${roleName}') then 
-          create role ${roleName};
+			await prisma.$transaction([
+				takeLock(prisma),
+				prisma.$queryRawUnsafe(`
+          do
+          $$
+          begin
+          if not exists (select * from pg_catalog.pg_roles where rolname = '${roleName}') then 
+            create role ${roleName};
+          end if;
+          end
+          $$
+          ;
+        `),
+				prisma.$queryRawUnsafe(`
           GRANT ${ability.operation} ON "${table}" TO ${roleName};
-        end if;
-        end
-        $$
-        ;
-      `);
+        `),
+			]);
 
 			if (ability.expression) {
 				await setRLS(prisma, table, roleName, ability.operation, ability.expression);
@@ -239,7 +250,7 @@ export const createRoles = async ({
 	// It's not possible to dynamically GRANT these to a shared user role, as the GRANT is not isolated per transaction and leads to broken permissions.
 	for (const key in roles) {
 		const role = createRoleName(key);
-		await prisma.$queryRawUnsafe(`
+		await prisma.$executeRawUnsafe(`
 		do
 		$$
 		begin
@@ -251,18 +262,6 @@ export const createRoles = async ({
 		;
 	`);
 
-		// Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
-		// This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
-		await prisma.$queryRawUnsafe(`
-		GRANT ALL ON ALL TABLES IN SCHEMA public TO ${role};
-	`);
-		await prisma.$queryRawUnsafe(`
-		GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO ${role};
-	`);
-		await prisma.$queryRawUnsafe(`
-		GRANT ALL ON SCHEMA public TO ${role};
-	`);
-
 		const wildCardAbilities = flatMap(abilities, (model, modelName) => {
 			return map(model, (params, slug) => {
 				return createAbilityName(modelName, slug);
@@ -273,7 +272,20 @@ export const createRoles = async ({
 			roleAbilities === "*"
 				? wildCardAbilities
 				: roleAbilities.map((ability) => createAbilityName(ability.model!, ability.slug!));
-		await prisma.$queryRawUnsafe(`GRANT ${rlsRoles.join(", ")} TO ${role}`);
+
+		// Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
+		// This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
+		await prisma.$transaction([
+			takeLock(prisma),
+			prisma.$executeRawUnsafe(`GRANT ALL ON ALL TABLES IN SCHEMA public TO ${role};`),
+			prisma.$executeRawUnsafe(`
+        GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO ${role};
+      `),
+			prisma.$executeRawUnsafe(`
+        GRANT ALL ON SCHEMA public TO ${role};
+      `),
+			prisma.$queryRawUnsafe(`GRANT ${rlsRoles.join(", ")} TO ${role}`),
+		]);
 	}
 };
 

test/integration/index.spec.ts

@@ -22,13 +22,40 @@ describe("setup", () => {
 			expect(getRoles.mock.calls).toHaveLength(1);
 			const abilities = getRoles.mock.calls[0][0];
 
-			expect(Object.keys(abilities)).toStrictEqual(["User", "Post"]);
+			expect(Object.keys(abilities)).toStrictEqual(["User", "Post", "Tag"]);
 			expect(Object.keys(abilities.User)).toStrictEqual(["create", "read", "update", "delete"]);
 			expect(Object.keys(abilities.Post)).toStrictEqual(["create", "read", "update", "delete"]);
+			expect(Object.keys(abilities.Tag)).toStrictEqual(["create", "read", "update", "delete"]);
 		});
 	});
 
 	describe("params.getContext()", () => {
+		it("should skip RBAC if .getContext() returns null", async () => {
+			const prisma = new PrismaClient();
+
+			const role = `USER_${uuid()}`;
+
+			await setup({
+				prisma,
+				getRoles(abilities) {
+					return {
+						[role]: [abilities.Post.read],
+					};
+				},
+				getContext: () => {
+					return null;
+				},
+			});
+
+			const post = await prisma.post.create({
+				data: {
+					title: "Test post",
+				},
+			});
+
+			expect(post.id).toBeDefined();
+		});
+
 		it("should allow a custom context to be set", async () => {
 			const prisma = new PrismaClient();