From f0e4678d0a9286b94c636d58a2f13a79f2e5ed62 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Fri, 11 Nov 2022 07:23:32 +0100
Subject: [PATCH] Kick guest users on revocation of guest access

---
 roomserver/internal/api.go                |   1 +
 roomserver/internal/input/input.go        |   3 +
 roomserver/internal/input/input_events.go | 105 ++++++++++++++++++++++
 3 files changed, 109 insertions(+)

diff --git a/roomserver/internal/api.go b/roomserver/internal/api.go
index 1a11586a5..18fa236df 100644
--- a/roomserver/internal/api.go
+++ b/roomserver/internal/api.go
@@ -194,6 +194,7 @@ func (r *RoomserverInternalAPI) SetFederationAPI(fsAPI fsAPI.RoomserverFederatio
 
 func (r *RoomserverInternalAPI) SetUserAPI(userAPI userapi.RoomserverUserAPI) {
 	r.Leaver.UserAPI = userAPI
+	r.Inputer.UserAPI = userAPI
 }
 
 func (r *RoomserverInternalAPI) SetAppserviceAPI(asAPI asAPI.AppServiceInternalAPI) {
diff --git a/roomserver/internal/input/input.go b/roomserver/internal/input/input.go
index f5099ca11..f6c6ddd2e 100644
--- a/roomserver/internal/input/input.go
+++ b/roomserver/internal/input/input.go
@@ -23,6 +23,8 @@ import (
 	"sync"
 	"time"
 
+	userapi "github.com/matrix-org/dendrite/userapi/api"
+
 	"github.com/Arceliar/phony"
 	"github.com/getsentry/sentry-go"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -87,6 +89,7 @@ type Inputer struct {
 	workers             sync.Map // room ID -> *worker
 
 	Queryer *query.Queryer
+	UserAPI userapi.RoomserverUserAPI
 }
 
 // If a room consumer is inactive for a while then we will allow NATS
diff --git a/roomserver/internal/input/input_events.go b/roomserver/internal/input/input_events.go
index 60160e8e5..f2bd4ba49 100644
--- a/roomserver/internal/input/input_events.go
+++ b/roomserver/internal/input/input_events.go
@@ -19,8 +19,11 @@ package input
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"errors"
 	"fmt"
+	userAPI "github.com/matrix-org/dendrite/userapi/api"
+	"github.com/tidwall/gjson"
 	"time"
 
 	"github.com/matrix-org/gomatrixserverlib"
@@ -428,6 +431,10 @@ func (r *Inputer) processRoomEvent(
 		}
 	}
 
+	if err = r.kickGuests(ctx, event, roomInfo); err != nil {
+		logrus.WithError(err).Error("failed to kick guest users on m.room.guest_access revocation")
+	}
+
 	// Everything was OK — the latest events updater didn't error and
 	// we've sent output events. Finally, generate a hook call.
 	hooks.Run(hooks.KindNewEventPersisted, headered)
@@ -709,3 +716,101 @@ func (r *Inputer) calculateAndSetState(
 	succeeded = true
 	return nil
 }
+
+// kickGuests kicks guests users from m.room.guest_access rooms, if guest access is now prohibited.
+func (r *Inputer) kickGuests(ctx context.Context, event *gomatrixserverlib.Event, roomInfo *types.RoomInfo) error {
+	if event.Type() != gomatrixserverlib.MRoomGuestAccess {
+		return nil
+	}
+	if gjson.GetBytes(event.Content(), "guest_access").String() == "can_join" {
+		return nil
+	}
+	membershipNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, roomInfo.RoomNID, true, true)
+	if err != nil {
+		return err
+	}
+
+	memberEvents, err := r.DB.Events(ctx, membershipNIDs)
+	if err != nil {
+		return err
+	}
+
+	inputEvents := make([]api.InputRoomEvent, 0, len(memberEvents))
+	latestReq := &api.QueryLatestEventsAndStateRequest{
+		RoomID: event.RoomID(),
+	}
+	latestRes := &api.QueryLatestEventsAndStateResponse{}
+	if err = r.Queryer.QueryLatestEventsAndState(ctx, latestReq, latestRes); err != nil {
+		return err
+	}
+
+	prevEvents := latestRes.LatestEvents
+	for _, memberEvent := range memberEvents {
+		if memberEvent.StateKey() == nil {
+			continue
+		}
+
+		localpart, senderDomain, err := gomatrixserverlib.SplitID('@', *memberEvent.StateKey())
+		if err != nil {
+			continue
+		}
+
+		accountRes := &userAPI.QueryAccountByLocalpartResponse{}
+		if err = r.UserAPI.QueryAccountByLocalpart(ctx, &userAPI.QueryAccountByLocalpartRequest{Localpart: localpart}, accountRes); err != nil {
+			return err
+		}
+		if accountRes.Account == nil {
+			continue
+		}
+
+		if accountRes.Account.AccountType != userAPI.AccountTypeGuest {
+			continue
+		}
+
+		var memberContent gomatrixserverlib.MemberContent
+		if err = json.Unmarshal(memberEvent.Content(), &memberContent); err != nil {
+			return err
+		}
+		memberContent.Membership = gomatrixserverlib.Leave
+
+		stateKey := *memberEvent.StateKey()
+		fledglingEvent := &gomatrixserverlib.EventBuilder{
+			RoomID:     event.RoomID(),
+			Type:       gomatrixserverlib.MRoomMember,
+			StateKey:   &stateKey,
+			Sender:     stateKey,
+			PrevEvents: prevEvents,
+		}
+
+		if fledglingEvent.Content, err = json.Marshal(memberContent); err != nil {
+			return err
+		}
+
+		eventsNeeded, err := gomatrixserverlib.StateNeededForEventBuilder(fledglingEvent)
+		if err != nil {
+			return err
+		}
+
+		event, err := eventutil.BuildEvent(ctx, fledglingEvent, r.Cfg.Matrix, time.Now(), &eventsNeeded, latestRes)
+		if err != nil {
+			return err
+		}
+
+		inputEvents = append(inputEvents, api.InputRoomEvent{
+			Kind:         api.KindNew,
+			Event:        event,
+			Origin:       senderDomain,
+			SendAsServer: string(senderDomain),
+		})
+		prevEvents = []gomatrixserverlib.EventReference{
+			event.EventReference(),
+		}
+	}
+
+	inputReq := &api.InputRoomEventsRequest{
+		InputRoomEvents: inputEvents,
+		Asynchronous:    true,
+	}
+	inputRes := &api.InputRoomEventsResponse{}
+	return r.InputRoomEvents(ctx, inputReq, inputRes)
+}