fix: ensure role setting still happens when using async middleware

This fixes a bug where RLS policies would not be used correctly when
Yates is combined with async middleware.
See prisma/prisma#18276

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

src/index.ts

@@ -31,6 +31,12 @@ export type GetContextFn<ContextKeys extends string = string> = () => {
 	};
 } | null;
 
+declare module "@prisma/client" {
+	interface PrismaClient {
+		_executeRequest: (params: any) => Promise<any>;
+	}
+}
+
 /**
  * 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.
@@ -87,22 +93,42 @@ export const createClient = (prisma: PrismaClient, getContext: GetContextFn) =>
 					}
 
 					try {
-						const txResults: any[] = await prisma.$transaction([
+						// Because batch transactions inside a prisma client query extension can run out of order if used with async middleware,
+						// we need to run the logic inside an interactive transaction, however this brings a different set of problems in that the
+						// main query will no longer automatically run inside the transaction. We resolve this issue by manually executing the prisma request.
+						// See https://github.com/prisma/prisma/issues/18276
+						const queryResults = await prisma.$transaction(async (tx) => {
 							// Switch to the user role, We can't use a prepared statement here, due to limitations in PG not allowing prepared statements to be used in SET ROLE
-							prisma.$queryRawUnsafe(`SET ROLE ${pgRole}`),
+							await tx.$queryRawUnsafe(`SET ROLE ${pgRole}`);
 							// Now set all the context variables using `set_config` so that they can be used in RLS
-							...toPairs(context).map(([key, value]) => {
-								return prisma.$queryRaw`SELECT set_config(${key}, ${value.toString()},  true);`;
-							}),
-							...[
-								// Now call original function
-								// Conveniently, the `query` function will happily run inside the transaction.
-								query(args),
-								// Switch role back to admin user
-								prisma.$queryRawUnsafe("SET ROLE none"),
-							],
-						]);
-						const queryResults = txResults[txResults.length - 2];
+							for (const [key, value] of toPairs(context)) {
+								await tx.$queryRaw`SELECT set_config(${key}, ${value.toString()},  true);`;
+							}
+
+							// Inconveniently, the `query` function will not run inside an interactive transaction.
+							// We need to manually reconstruct the query, and attached the "secret" transaction ID.
+							// This ensures that the query will run inside the transaction AND that middlewares will not be re-applied
+
+							// https://github.com/prisma/prisma/blob/4.11.0/packages/client/src/runtime/getPrismaClient.ts#L1013
+							const txId = (tx as any)[Symbol.for("prisma.client.transaction.id")];
+
+							// See https://github.com/prisma/prisma/blob/4.11.0/packages/client/src/runtime/getPrismaClient.ts#L860
+							const result = await prisma._executeRequest({
+								args,
+								clientMethod: `${model.toLowerCase()}.${operation}`,
+								jsModelName: model.toLowerCase(),
+								action: operation,
+								model,
+								transaction: {
+									kind: "itx",
+									id: txId,
+								},
+							});
+							// Switch role back to admin user
+							await tx.$queryRawUnsafe("SET ROLE none");
+
+							return result;
+						});
 
 						return queryResults;
 					} catch (e) {

test/integration/middleware.spec.ts

@@ -103,4 +103,39 @@ describe("middlewares", () => {
 			}),
 		).rejects.toThrow();
 	});
+
+	it("should not be able to bypass RBAC when async middleware is used", async () => {
+		const prisma = new PrismaClient();
+
+		const middleware: Prisma.Middleware = async (params, next) => {
+			await "test";
+			return next(params);
+		};
+
+		prisma.$use(middleware);
+
+		const roleName = `USER_${uuid()}`;
+
+		const client = await setup({
+			prisma,
+			getRoles(abilities) {
+				return {
+					[roleName]: [abilities.Post.read],
+				};
+			},
+			getContext: () => {
+				return {
+					role: roleName,
+				};
+			},
+		});
+
+		await expect(
+			client.post.create({
+				data: {
+					title: `Test post from ${roleName}`,
+				},
+			}),
+		).rejects.toThrow();
+	});
 });