From f802598abbf787435278b51f8b50e316b5e6a1f1 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Fri, 11 Nov 2022 10:04:24 +0100
Subject: [PATCH] Assing stateKeyNID if we don't have the expected count

---
 roomserver/storage/shared/storage.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 4455ec3bf..a4dda62df 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -110,6 +110,19 @@ func (d *Database) eventStateKeyNIDs(
 	for eventStateKey, nid := range nids {
 		result[eventStateKey] = nid
 	}
+	// We received some nids, but are still missing some, work out which and create them
+	if len(eventStateKeys) < len(result) {
+		for _, eventStateKey := range eventStateKeys {
+			if _, ok := result[eventStateKey]; ok {
+				continue
+			}
+			nid, err := d.assignStateKeyNID(ctx, txn, eventStateKey)
+			if err != nil {
+				return result, err
+			}
+			result[eventStateKey] = nid
+		}
+	}
 	return result, nil
 }