3 files changed, 16 insertions, 4 deletions
diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 32d63bb5..aa7dc473 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -439,6 +439,7 @@ type QueryMembershipAtEventRequest struct {
 
 // QueryMembershipAtEventResponse is the response to QueryMembershipAtEventRequest.
 type QueryMembershipAtEventResponse struct {
-	// Memberships is a map from eventID to a list of events (if any).
+	// Memberships is a map from eventID to a list of events (if any). Events that
+	// do not have known state will return an empty array here.
 	Memberships map[string][]*gomatrixserverlib.HeaderedEvent `json:"memberships"`
 }
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index d08c5c49..b41a92e9 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -208,6 +208,9 @@ func (r *Queryer) QueryMembershipForUser(
 	return err
 }
 
+// QueryMembershipAtEvent returns the known memberships at a given event.
+// If the state before an event is not known, an empty list will be returned
+// for that event instead.
 func (r *Queryer) QueryMembershipAtEvent(
 	ctx context.Context,
 	request *api.QueryMembershipAtEventRequest,
@@ -237,7 +240,11 @@ func (r *Queryer) QueryMembershipAtEvent(
 	}
 
 	for _, eventID := range request.EventIDs {
-		stateEntry := stateEntries[eventID]
+		stateEntry, ok := stateEntries[eventID]
+		if !ok {
+			response.Memberships[eventID] = []*gomatrixserverlib.HeaderedEvent{}
+			continue
+		}
 		memberships, err := helpers.GetMembershipsAtState(ctx, r.DB, stateEntry, false)
 		if err != nil {
 			return fmt.Errorf("unable to get memberships at state: %w", err)
diff --git a/roomserver/state/state.go b/roomserver/state/state.go
index a40a2e9b..cb96d83e 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -18,6 +18,7 @@ package state
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"sort"
 	"sync"
@@ -134,11 +135,14 @@ func (v *StateResolution) LoadMembershipAtEvent(
 	for i := range eventIDs {
 		eventID := eventIDs[i]
 		snapshotNID, err := v.db.SnapshotNIDFromEventID(ctx, eventID)
-		if err != nil {
+		if err != nil && err != sql.ErrNoRows {
 			return nil, fmt.Errorf("LoadStateAtEvent.SnapshotNIDFromEventID failed for event %s : %w", eventID, err)
 		}
 		if snapshotNID == 0 {
-			return nil, fmt.Errorf("LoadStateAtEvent.SnapshotNIDFromEventID(%s) returned 0 NID, was this event stored?", eventID)
+			// If we don't know a state snapshot for this event then we can't calculate
+			// memberships at the time of the event, so skip over it. This means that
+			// it isn't guaranteed that the response map will contain every single event.
+			continue
 		}
 		snapshotNIDMap[snapshotNID] = append(snapshotNIDMap[snapshotNID], eventID)
 	}