From 0780228b2ab56c2812b43fda6902abac4159b42e Mon Sep 17 00:00:00 2001
From: Guilherme Gazzo <guilhermegazzo@gmail.com>
Date: Mon, 16 Dec 2024 16:24:36 -0300
Subject: [PATCH] chore: save event as staging after validation

---
 packages/homeserver/src/plugins/mongodb.ts    | 13 ++++++++++
 .../routes/federation/sendTransaction.spec.ts | 12 ++++++++-
 .../src/routes/federation/sendTransaction.ts  | 25 ++++++-------------
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/packages/homeserver/src/plugins/mongodb.ts b/packages/homeserver/src/plugins/mongodb.ts
index 274a1504..220d2098 100644
--- a/packages/homeserver/src/plugins/mongodb.ts
+++ b/packages/homeserver/src/plugins/mongodb.ts
@@ -3,6 +3,7 @@ import type { InferContext } from "elysia";
 import { type Db, MongoClient } from "mongodb";
 
 import type { EventBase } from "@hs/core/src/events/eventBase";
+import { generateId } from "../authentication";
 
 export interface Server {
 	_id: string;
@@ -153,6 +154,17 @@ export const routerWithMongodb = (db: Db) =>
 				);
 			};
 
+			const createStagingEvent = async (event: EventBase) => {
+				const id = generateId(event);
+				await eventsCollection.insertOne({
+					_id: id,
+					event,
+					staged: true,
+				});
+
+				return id;
+			};
+
 			return {
 				serversCollection,
 				getValidPublicKeyFromLocal,
@@ -163,6 +175,7 @@ export const routerWithMongodb = (db: Db) =>
 				getMissingEventsByDeep,
 				getLastEvent,
 				getAuthEvents,
+				createStagingEvent,
 			};
 		})(),
 	);
diff --git a/packages/homeserver/src/routes/federation/sendTransaction.spec.ts b/packages/homeserver/src/routes/federation/sendTransaction.spec.ts
index a77ef331..0cc68c09 100644
--- a/packages/homeserver/src/routes/federation/sendTransaction.spec.ts
+++ b/packages/homeserver/src/routes/federation/sendTransaction.spec.ts
@@ -48,6 +48,9 @@ describe("/send/:txnId", () => {
 							return;
 						},
 					},
+					createStagingEvent: async () => {
+						return;
+					},
 					serversCollection: {
 						findOne: async () => {
 							return;
@@ -83,6 +86,7 @@ describe("/send/:txnId", () => {
 				}),
 			);
 
+			const data = await resp.json();
 			expect(resp.status).toBe(400);
 		});
 
@@ -180,9 +184,12 @@ describe("/send/:txnId", () => {
 			);
 
 			const data = await resp.json();
+			const id = generateId(signedPdu);
 			expect(resp.status).toBe(200);
 			expect(data).toHaveProperty("pdus");
-			expect(data.pdus).toBeEmptyObject();
+			expect(data.pdus).toStrictEqual({
+				[id]: {},
+			});
 		});
 	});
 });
@@ -228,6 +235,9 @@ describe("/send/:txnId using real case", () => {
 							return;
 						},
 					},
+					createStagingEvent: async () => {
+						return;
+					},
 					serversCollection: {
 						findOne: async () => {
 							return;
diff --git a/packages/homeserver/src/routes/federation/sendTransaction.ts b/packages/homeserver/src/routes/federation/sendTransaction.ts
index 9179474a..408775e6 100644
--- a/packages/homeserver/src/routes/federation/sendTransaction.ts
+++ b/packages/homeserver/src/routes/federation/sendTransaction.ts
@@ -30,7 +30,7 @@ export const sendTransactionRoute = new Elysia().put(
 
 		const {
 			config,
-			mongo: { eventsCollection },
+			mongo: { eventsCollection, createStagingEvent },
 		} = context;
 
 		const { pdus, edus = [] } = body as any;
@@ -167,28 +167,17 @@ export const sendTransactionRoute = new Elysia().put(
 			for (const [roomId, pdus] of pdusByRoomId) {
 				// const roomVersion = getRoomVersion
 				for (const pdu of pdus) {
-					if (
-						!(await validatePdu(pdu).catch((e) => {
-							console.error("error validating pdu", e);
-							return true;
-						}))
-					) {
+					try {
+						await validatePdu(pdu);
 						resultPDUs[`${generateId(pdu)}`] = {};
+						void createStagingEvent(pdu);
+					} catch (e) {
+						console.error("error validating pdu", e);
+						resultPDUs[`${generateId(pdu)}`] = e as any;
 					}
 				}
 			}
 
-			await eventsCollection
-				.insertMany(
-					pdus.map((event) => ({
-						_id: generateId(event),
-						event,
-					})),
-				)
-				.catch((e) => {
-					console.error("error saving event", e);
-				});
-
 			return {
 				pdus: resultPDUs,
 			};