Add tests, more caching

roomserver/internal/query/query.go

@@ -521,14 +521,10 @@ func (r *Queryer) QueryMissingEvents(
 	response.Events = make([]*gomatrixserverlib.HeaderedEvent, 0, len(loadedEvents)-len(eventsToFilter))
 	for _, event := range loadedEvents {
 		if !eventsToFilter[event.EventID()] {
-			roomVersion, verr := r.roomVersion(event.RoomID())
-			if verr != nil {
-				return verr
-			}
 			if _, ok := redactEventIDs[event.EventID()]; ok {
 				event.Redact()
 			}
-			response.Events = append(response.Events, event.Headered(roomVersion))
+			response.Events = append(response.Events, event.Headered(info.RoomVersion))
 		}
 	}
 
@@ -718,14 +714,6 @@ func (r *Queryer) QueryRoomVersionForRoom(
 	return nil
 }
 
-func (r *Queryer) roomVersion(roomID string) (gomatrixserverlib.RoomVersion, error) {
-	var res api.QueryRoomVersionForRoomResponse
-	err := r.QueryRoomVersionForRoom(context.Background(), &api.QueryRoomVersionForRoomRequest{
-		RoomID: roomID,
-	}, &res)
-	return res.RoomVersion, err
-}
-
 func (r *Queryer) QueryPublishedRooms(
 	ctx context.Context,
 	req *api.QueryPublishedRoomsRequest,
@@ -910,8 +898,8 @@ func (r *Queryer) QueryRestrictedJoinAllowed(ctx context.Context, req *api.Query
 	if err = json.Unmarshal(joinRulesEvent.Content(), &joinRules); err != nil {
 		return fmt.Errorf("json.Unmarshal: %w", err)
 	}
-	// If the join rule isn't "restricted" then there's nothing more to do.
+	// If the join rule isn't "restricted" or "knock_restricted" then there's nothing more to do.
-	res.Restricted = joinRules.JoinRule == spec.Restricted
+	res.Restricted = joinRules.JoinRule == spec.Restricted || joinRules.JoinRule == spec.KnockRestricted
 	if !res.Restricted {
 		return nil
 	}
@@ -932,9 +920,9 @@ func (r *Queryer) QueryRestrictedJoinAllowed(ctx context.Context, req *api.Query
 	if err != nil {
 		return fmt.Errorf("r.DB.GetStateEvent: %w", err)
 	}
-	var powerLevels gomatrixserverlib.PowerLevelContent
+	powerLevels, err := powerLevelsEvent.PowerLevels()
-	if err = json.Unmarshal(powerLevelsEvent.Content(), &powerLevels); err != nil {
+	if err != nil {
-		return fmt.Errorf("json.Unmarshal: %w", err)
+		return fmt.Errorf("unable to get powerlevels: %w", err)
 	}
 	// Step through the join rules and see if the user matches any of them.
 	for _, rule := range joinRules.Allow {

roomserver/roomserver_test.go

@@ -582,3 +582,175 @@ func TestRedaction(t *testing.T) {
 		}
 	})
 }
+
+func TestQueryRestrictedJoinAllowed(t *testing.T) {
+	alice := test.NewUser(t)
+	bob := test.NewUser(t)
+
+	// a room we don't create in the database
+	allowedByRoomNotExists := test.NewRoom(t, alice)
+
+	// a room we create in the database, used for authorisation
+	allowedByRoomExists := test.NewRoom(t, alice)
+	allowedByRoomExists.CreateAndInsert(t, bob, spec.MRoomMember, map[string]interface{}{
+		"membership": spec.Join,
+	}, test.WithStateKey(bob.ID))
+
+	testCases := []struct {
+		name            string
+		prepareRoomFunc func(t *testing.T) *test.Room
+		wantResponse    api.QueryRestrictedJoinAllowedResponse
+	}{
+		{
+			name: "public room unrestricted",
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				return test.NewRoom(t, alice)
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Resident: true,
+			},
+		},
+		{
+			name: "room version without restrictions",
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				return test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV7))
+			},
+		},
+		{
+			name: "restricted only", // bob is not allowed to join
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				r := test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV8))
+				r.CreateAndInsert(t, alice, spec.MRoomJoinRules, map[string]interface{}{
+					"join_rule": spec.Restricted,
+				}, test.WithStateKey(""))
+				return r
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Resident:   true,
+				Restricted: true,
+			},
+		},
+		{
+			name: "knock_restricted",
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				r := test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV8))
+				r.CreateAndInsert(t, alice, spec.MRoomJoinRules, map[string]interface{}{
+					"join_rule": spec.KnockRestricted,
+				}, test.WithStateKey(""))
+				return r
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Resident:   true,
+				Restricted: true,
+			},
+		},
+		{
+			name: "restricted with pending invite", // bob should be allowed to join
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				r := test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV8))
+				r.CreateAndInsert(t, alice, spec.MRoomJoinRules, map[string]interface{}{
+					"join_rule": spec.Restricted,
+				}, test.WithStateKey(""))
+				r.CreateAndInsert(t, alice, spec.MRoomMember, map[string]interface{}{
+					"membership": spec.Invite,
+				}, test.WithStateKey(bob.ID))
+				return r
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Resident:   true,
+				Restricted: true,
+				Allowed:    true,
+			},
+		},
+		{
+			name: "restricted with allowed room_id, but missing room", // bob should not be allowed to join, as we don't know about the room
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				r := test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV10))
+				r.CreateAndInsert(t, alice, spec.MRoomJoinRules, map[string]interface{}{
+					"join_rule": spec.KnockRestricted,
+					"allow": []map[string]interface{}{
+						{
+							"room_id": allowedByRoomNotExists.ID,
+							"type":    spec.MRoomMembership,
+						},
+					},
+				}, test.WithStateKey(""))
+				r.CreateAndInsert(t, bob, spec.MRoomMember, map[string]interface{}{
+					"membership":                       spec.Join,
+					"join_authorised_via_users_server": alice.ID,
+				}, test.WithStateKey(bob.ID))
+				return r
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Restricted: true,
+			},
+		},
+		{
+			name: "restricted with allowed room_id", // bob should be allowed to join, as we know about the room
+			prepareRoomFunc: func(t *testing.T) *test.Room {
+				r := test.NewRoom(t, alice, test.RoomVersion(gomatrixserverlib.RoomVersionV10))
+				r.CreateAndInsert(t, alice, spec.MRoomJoinRules, map[string]interface{}{
+					"join_rule": spec.KnockRestricted,
+					"allow": []map[string]interface{}{
+						{
+							"room_id": allowedByRoomExists.ID,
+							"type":    spec.MRoomMembership,
+						},
+					},
+				}, test.WithStateKey(""))
+				r.CreateAndInsert(t, bob, spec.MRoomMember, map[string]interface{}{
+					"membership":                       spec.Join,
+					"join_authorised_via_users_server": alice.ID,
+				}, test.WithStateKey(bob.ID))
+				return r
+			},
+			wantResponse: api.QueryRestrictedJoinAllowedResponse{
+				Resident:      true,
+				Restricted:    true,
+				Allowed:       true,
+				AuthorisedVia: alice.ID,
+			},
+		},
+	}
+
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		cfg, processCtx, close := testrig.CreateConfig(t, dbType)
+		natsInstance := jetstream.NATSInstance{}
+		defer close()
+
+		cm := sqlutil.NewConnectionManager(processCtx, cfg.Global.DatabaseOptions)
+		caches := caching.NewRistrettoCache(128*1024*1024, time.Hour, caching.DisableMetrics)
+
+		rsAPI := roomserver.NewInternalAPI(processCtx, cfg, cm, &natsInstance, caches, caching.DisableMetrics)
+		rsAPI.SetFederationAPI(nil, nil)
+
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				if tc.prepareRoomFunc == nil {
+					t.Fatal("missing prepareRoomFunc")
+				}
+				testRoom := tc.prepareRoomFunc(t)
+				// Create the room
+				if err := api.SendEvents(processCtx.Context(), rsAPI, api.KindNew, testRoom.Events(), "test", "test", "test", nil, false); err != nil {
+					t.Errorf("failed to send events: %v", err)
+				}
+
+				if err := api.SendEvents(processCtx.Context(), rsAPI, api.KindNew, allowedByRoomExists.Events(), "test", "test", "test", nil, false); err != nil {
+					t.Errorf("failed to send events: %v", err)
+				}
+
+				req := api.QueryRestrictedJoinAllowedRequest{
+					UserID: bob.ID,
+					RoomID: testRoom.ID,
+				}
+				res := api.QueryRestrictedJoinAllowedResponse{}
+				if err := rsAPI.QueryRestrictedJoinAllowed(processCtx.Context(), &req, &res); err != nil {
+					t.Fatal(err)
+				}
+				if !reflect.DeepEqual(tc.wantResponse, res) {
+					t.Fatalf("unexpected response, want %#v - got %#v", tc.wantResponse, res)
+				}
+			})
+		}
+	})
+}

roomserver/storage/shared/storage.go

@@ -1164,7 +1164,7 @@ func (d *Database) GetStateEvent(ctx context.Context, roomID, evType, stateKey s
 	if roomInfo.IsStub() {
 		return nil, nil
 	}
-	eventTypeNID, err := d.EventTypesTable.SelectEventTypeNID(ctx, nil, evType)
+	eventTypeNID, err := d.GetOrCreateEventTypeNID(ctx, evType)
 	if err == sql.ErrNoRows {
 		// No rooms have an event of this type, otherwise we'd have an event type NID
 		return nil, nil
@@ -1172,7 +1172,7 @@ func (d *Database) GetStateEvent(ctx context.Context, roomID, evType, stateKey s
 	if err != nil {
 		return nil, err
 	}
-	stateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, stateKey)
+	stateKeyNID, err := d.GetOrCreateEventStateKeyNID(ctx, &stateKey)
 	if err == sql.ErrNoRows {
 		// No rooms have a state event with this state key, otherwise we'd have an state key NID
 		return nil, nil
@@ -1201,6 +1201,10 @@ func (d *Database) GetStateEvent(ctx context.Context, roomID, evType, stateKey s
 	// return the event requested
 	for _, e := range entries {
 		if e.EventTypeNID == eventTypeNID && e.EventStateKeyNID == stateKeyNID {
+			cachedEvent, ok := d.Cache.GetRoomServerEvent(e.EventNID)
+			if ok {
+				return cachedEvent.Headered(roomInfo.RoomVersion), nil
+			}
 			data, err := d.EventJSONTable.BulkSelectEventJSON(ctx, nil, []types.EventNID{e.EventNID})
 			if err != nil {
 				return nil, err
@@ -1324,7 +1328,7 @@ func (d *Database) GetBulkStateContent(ctx context.Context, roomIDs []string, tu
 	}
 	// we don't bother failing the request if we get asked for event types we don't know about, as all that would result in is no matches which
 	// isn't a failure.
-	eventTypeNIDMap, err := d.EventTypesTable.BulkSelectEventTypeNID(ctx, nil, eventTypes)
+	eventTypeNIDMap, err := d.eventTypeNIDs(ctx, nil, eventTypes)
 	if err != nil {
 		return nil, fmt.Errorf("GetBulkStateContent: failed to map event type nids: %w", err)
 	}