6 files changed, 105 insertions, 0 deletions
diff --git a/roomserver/api/api.go b/roomserver/api/api.go
index 2495157a..159c1829 100644
--- a/roomserver/api/api.go
+++ b/roomserver/api/api.go
@@ -89,6 +89,13 @@ type RoomserverInternalAPI interface {
 		response *QueryMembershipsForRoomResponse,
 	) error
 
+	// Query if we think we're still in a room.
+	QueryServerJoinedToRoom(
+		ctx context.Context,
+		request *QueryServerJoinedToRoomRequest,
+		response *QueryServerJoinedToRoomResponse,
+	) error
+
 	// Query whether a server is allowed to see an event
 	QueryServerAllowedToSeeEvent(
 		ctx context.Context,
diff --git a/roomserver/api/api_trace.go b/roomserver/api/api_trace.go
index b7accb9a..5fabbc21 100644
--- a/roomserver/api/api_trace.go
+++ b/roomserver/api/api_trace.go
@@ -134,6 +134,16 @@ func (t *RoomserverInternalAPITrace) QueryMembershipsForRoom(
 	return err
 }
 
+func (t *RoomserverInternalAPITrace) QueryServerJoinedToRoom(
+	ctx context.Context,
+	req *QueryServerJoinedToRoomRequest,
+	res *QueryServerJoinedToRoomResponse,
+) error {
+	err := t.Impl.QueryServerJoinedToRoom(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryServerJoinedToRoom req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
 func (t *RoomserverInternalAPITrace) QueryServerAllowedToSeeEvent(
 	ctx context.Context,
 	req *QueryServerAllowedToSeeEventRequest,
diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 67a217c8..5d61e862 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -140,6 +140,22 @@ type QueryMembershipsForRoomResponse struct {
 	HasBeenInRoom bool `json:"has_been_in_room"`
 }
 
+// QueryServerJoinedToRoomRequest is a request to QueryServerJoinedToRoom
+type QueryServerJoinedToRoomRequest struct {
+	// Server name of the server to find
+	ServerName gomatrixserverlib.ServerName `json:"server_name"`
+	// ID of the room to see if we are still joined to
+	RoomID string `json:"room_id"`
+}
+
+// QueryMembershipsForRoomResponse is a response to QueryServerJoinedToRoom
+type QueryServerJoinedToRoomResponse struct {
+	// True if the room exists on the server
+	RoomExists bool `json:"room_exists"`
+	// True if we still believe that we are participating in the room
+	IsInRoom bool `json:"is_in_room"`
+}
+
 // QueryServerAllowedToSeeEventRequest is a request to QueryServerAllowedToSeeEvent
 type QueryServerAllowedToSeeEventRequest struct {
 	// The event ID to look up invites in.
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index fb981447..58cb4493 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -227,6 +227,50 @@ func (r *Queryer) QueryMembershipsForRoom(
 	return nil
 }
 
+// QueryServerJoinedToRoom implements api.RoomserverInternalAPI
+func (r *Queryer) QueryServerJoinedToRoom(
+	ctx context.Context,
+	request *api.QueryServerJoinedToRoomRequest,
+	response *api.QueryServerJoinedToRoomResponse,
+) error {
+	info, err := r.DB.RoomInfo(ctx, request.RoomID)
+	if err != nil {
+		return fmt.Errorf("r.DB.RoomInfo: %w", err)
+	}
+	if info == nil || info.IsStub {
+		return nil
+	}
+	response.RoomExists = true
+
+	eventNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false)
+	if err != nil {
+		return fmt.Errorf("r.DB.GetMembershipEventNIDsForRoom: %w", err)
+	}
+	if len(eventNIDs) == 0 {
+		return nil
+	}
+
+	events, err := r.DB.Events(ctx, eventNIDs)
+	if err != nil {
+		return fmt.Errorf("r.DB.Events: %w", err)
+	}
+
+	for _, e := range events {
+		if e.Type() == gomatrixserverlib.MRoomMember && e.StateKey() != nil {
+			_, serverName, err := gomatrixserverlib.SplitID('@', *e.StateKey())
+			if err != nil {
+				continue
+			}
+			if serverName == request.ServerName {
+				response.IsInRoom = true
+				break
+			}
+		}
+	}
+
+	return nil
+}
+
 // QueryServerAllowedToSeeEvent implements api.RoomserverInternalAPI
 func (r *Queryer) QueryServerAllowedToSeeEvent(
 	ctx context.Context,
diff --git a/roomserver/inthttp/client.go b/roomserver/inthttp/client.go
index f2510c75..3dd3edaf 100644
--- a/roomserver/inthttp/client.go
+++ b/roomserver/inthttp/client.go
@@ -38,6 +38,7 @@ const (
 	RoomserverQueryEventsByIDPath              = "/roomserver/queryEventsByID"
 	RoomserverQueryMembershipForUserPath       = "/roomserver/queryMembershipForUser"
 	RoomserverQueryMembershipsForRoomPath      = "/roomserver/queryMembershipsForRoom"
+	RoomserverQueryServerJoinedToRoomPath      = "/roomserver/queryServerJoinedToRoomPath"
 	RoomserverQueryServerAllowedToSeeEventPath = "/roomserver/queryServerAllowedToSeeEvent"
 	RoomserverQueryMissingEventsPath           = "/roomserver/queryMissingEvents"
 	RoomserverQueryStateAndAuthChainPath       = "/roomserver/queryStateAndAuthChain"
@@ -312,6 +313,19 @@ func (h *httpRoomserverInternalAPI) QueryMembershipsForRoom(
 	return httputil.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
 }
 
+// QueryMembershipsForRoom implements RoomserverQueryAPI
+func (h *httpRoomserverInternalAPI) QueryServerJoinedToRoom(
+	ctx context.Context,
+	request *api.QueryServerJoinedToRoomRequest,
+	response *api.QueryServerJoinedToRoomResponse,
+) error {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryServerJoinedToRoom")
+	defer span.Finish()
+
+	apiURL := h.roomserverURL + RoomserverQueryServerJoinedToRoomPath
+	return httputil.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
+}
+
 // QueryServerAllowedToSeeEvent implements RoomserverQueryAPI
 func (h *httpRoomserverInternalAPI) QueryServerAllowedToSeeEvent(
 	ctx context.Context,
diff --git a/roomserver/inthttp/server.go b/roomserver/inthttp/server.go
index 8ffa9cf9..c7e541dd 100644
--- a/roomserver/inthttp/server.go
+++ b/roomserver/inthttp/server.go
@@ -168,6 +168,20 @@ func AddRoutes(r api.RoomserverInternalAPI, internalAPIMux *mux.Router) {
 		}),
 	)
 	internalAPIMux.Handle(
+		RoomserverQueryServerJoinedToRoomPath,
+		httputil.MakeInternalAPI("queryServerJoinedToRoom", func(req *http.Request) util.JSONResponse {
+			var request api.QueryServerJoinedToRoomRequest
+			var response api.QueryServerJoinedToRoomResponse
+			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
+				return util.ErrorResponse(err)
+			}
+			if err := r.QueryServerJoinedToRoom(req.Context(), &request, &response); err != nil {
+				return util.ErrorResponse(err)
+			}
+			return util.JSONResponse{Code: http.StatusOK, JSON: &response}
+		}),
+	)
+	internalAPIMux.Handle(
 		RoomserverQueryServerAllowedToSeeEventPath,
 		httputil.MakeInternalAPI("queryServerAllowedToSeeEvent", func(req *http.Request) util.JSONResponse {
 			var request api.QueryServerAllowedToSeeEventRequest