From 48e8e472a342c3e28f552ef5aae06962ee2da3a9 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Mon, 17 Jan 2022 15:20:54 +0200
Subject: [PATCH] Update state store from IntentAPI requests

---
 appservice/intent.go | 93 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 87 insertions(+), 6 deletions(-)

diff --git a/appservice/intent.go b/appservice/intent.go
index 65271d76..84a47793 100644
--- a/appservice/intent.go
+++ b/appservice/intent.go
@@ -7,6 +7,7 @@
 package appservice
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"strings"
@@ -90,7 +91,7 @@ func (intent *IntentAPI) EnsureJoined(roomID id.RoomID) error {
 			return fmt.Errorf("failed to ensure joined: %w", err)
 		}
 	}
-	intent.as.StateStore.SetMembership(resp.RoomID, intent.UserID, "join")
+	intent.as.StateStore.SetMembership(resp.RoomID, intent.UserID, event.MembershipJoin)
 	return nil
 }
 
@@ -108,32 +109,112 @@ func (intent *IntentAPI) SendMassagedMessageEvent(roomID id.RoomID, eventType ev
 	return intent.Client.SendMessageEvent(roomID, eventType, contentJSON, mautrix.ReqSendEvent{Timestamp: ts})
 }
 
+func (intent *IntentAPI) updateStoreWithOutgoingEvent(roomID id.RoomID, eventType event.Type, stateKey string, contentJSON interface{}, eventID id.EventID) {
+	fakeEvt := &event.Event{
+		StateKey: &stateKey,
+		Sender:   intent.UserID,
+		Type:     eventType,
+		ID:       eventID,
+		RoomID:   roomID,
+		Content:  event.Content{},
+	}
+	var err error
+	fakeEvt.Content.VeryRaw, err = json.Marshal(contentJSON)
+	if err != nil {
+		intent.Logger.Debugfln("Failed to marshal state event content to update state store: %v", err)
+		return
+	}
+	err = json.Unmarshal(fakeEvt.Content.VeryRaw, &fakeEvt.Content.Raw)
+	if err != nil {
+		intent.Logger.Debugfln("Failed to unmarshal state event content to update state store: %v", err)
+		return
+	}
+	err = fakeEvt.Content.ParseRaw(fakeEvt.Type)
+	if err != nil {
+		intent.Logger.Debugfln("Failed to parse state event content to update state store: %v", err)
+		return
+	}
+	intent.as.UpdateState(fakeEvt)
+}
+
 func (intent *IntentAPI) SendStateEvent(roomID id.RoomID, eventType event.Type, stateKey string, contentJSON interface{}) (*mautrix.RespSendEvent, error) {
 	if err := intent.EnsureJoined(roomID); err != nil {
 		return nil, err
 	}
-	return intent.Client.SendStateEvent(roomID, eventType, stateKey, contentJSON)
+	resp, err := intent.Client.SendStateEvent(roomID, eventType, stateKey, contentJSON)
+	if err == nil && resp != nil {
+		intent.updateStoreWithOutgoingEvent(roomID, eventType, stateKey, contentJSON, resp.EventID)
+	}
+	return resp, err
 }
 
 func (intent *IntentAPI) SendMassagedStateEvent(roomID id.RoomID, eventType event.Type, stateKey string, contentJSON interface{}, ts int64) (*mautrix.RespSendEvent, error) {
 	if err := intent.EnsureJoined(roomID); err != nil {
 		return nil, err
 	}
-	return intent.Client.SendMassagedStateEvent(roomID, eventType, stateKey, contentJSON, ts)
+	resp, err := intent.Client.SendMassagedStateEvent(roomID, eventType, stateKey, contentJSON, ts)
+	if err == nil && resp != nil {
+		intent.updateStoreWithOutgoingEvent(roomID, eventType, stateKey, contentJSON, resp.EventID)
+	}
+	return resp, err
 }
 
 func (intent *IntentAPI) StateEvent(roomID id.RoomID, eventType event.Type, stateKey string, outContent interface{}) error {
 	if err := intent.EnsureJoined(roomID); err != nil {
 		return err
 	}
-	return intent.Client.StateEvent(roomID, eventType, stateKey, outContent)
+	err := intent.Client.StateEvent(roomID, eventType, stateKey, outContent)
+	if err == nil {
+		intent.updateStoreWithOutgoingEvent(roomID, eventType, stateKey, outContent, "")
+	}
+	return err
 }
 
 func (intent *IntentAPI) State(roomID id.RoomID) (mautrix.RoomStateMap, error) {
 	if err := intent.EnsureJoined(roomID); err != nil {
 		return nil, err
 	}
-	return intent.Client.State(roomID)
+	state, err := intent.Client.State(roomID)
+	if err == nil {
+		for _, events := range state {
+			for _, evt := range events {
+				intent.as.UpdateState(evt)
+			}
+		}
+	}
+	return state, err
+}
+
+func (intent *IntentAPI) InviteUser(roomID id.RoomID, req *mautrix.ReqInviteUser) (resp *mautrix.RespInviteUser, err error) {
+	resp, err = intent.Client.InviteUser(roomID, req)
+	if err == nil {
+		intent.as.StateStore.SetMembership(roomID, req.UserID, event.MembershipInvite)
+	}
+	return
+}
+
+func (intent *IntentAPI) KickUser(roomID id.RoomID, req *mautrix.ReqKickUser) (resp *mautrix.RespKickUser, err error) {
+	resp, err = intent.Client.KickUser(roomID, req)
+	if err == nil {
+		intent.as.StateStore.SetMembership(roomID, req.UserID, event.MembershipLeave)
+	}
+	return
+}
+
+func (intent *IntentAPI) BanUser(roomID id.RoomID, req *mautrix.ReqBanUser) (resp *mautrix.RespBanUser, err error) {
+	resp, err = intent.Client.BanUser(roomID, req)
+	if err == nil {
+		intent.as.StateStore.SetMembership(roomID, req.UserID, event.MembershipBan)
+	}
+	return
+}
+
+func (intent *IntentAPI) UnbanUser(roomID id.RoomID, req *mautrix.ReqUnbanUser) (resp *mautrix.RespUnbanUser, err error) {
+	resp, err = intent.Client.UnbanUser(roomID, req)
+	if err == nil {
+		intent.as.StateStore.SetMembership(roomID, req.UserID, event.MembershipLeave)
+	}
+	return
 }
 
 func (intent *IntentAPI) Member(roomID id.RoomID, userID id.UserID) *event.MemberEventContent {
@@ -303,7 +384,7 @@ func (intent *IntentAPI) Members(roomID id.RoomID, req ...mautrix.ReqMembers) (r
 
 func (intent *IntentAPI) EnsureInvited(roomID id.RoomID, userID id.UserID) error {
 	if !intent.as.StateStore.IsInvited(roomID, userID) {
-		_, err := intent.Client.InviteUser(roomID, &mautrix.ReqInviteUser{
+		_, err := intent.InviteUser(roomID, &mautrix.ReqInviteUser{
 			UserID: userID,
 		})
 		if httpErr, ok := err.(mautrix.HTTPError); ok && httpErr.RespError != nil && strings.Contains(httpErr.RespError.Err, "is already in the room") {