From b8efcf99cc977cfc18e854637f494a4f314fc364 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Thu, 30 Sep 2021 16:05:56 +0300
Subject: [PATCH] Add method to get full room state

---
 appservice/intent.go |  9 +++++++-
 client.go            | 49 ++++++++++++++++++++++++++++++++++++++++++++
 room.go              |  6 ++++--
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/appservice/intent.go b/appservice/intent.go
index 725e4ab..459eb9b 100644
--- a/appservice/intent.go
+++ b/appservice/intent.go
@@ -121,13 +121,20 @@ func (intent *IntentAPI) SendMassagedStateEvent(roomID id.RoomID, eventType even
 	return intent.Client.SendMassagedStateEvent(roomID, eventType, stateKey, contentJSON, ts)
 }
 
-func (intent *IntentAPI) StateEvent(roomID id.RoomID, eventType event.Type, stateKey string, outContent interface{}) (err error) {
+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)
 }
 
+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)
+}
+
 func (intent *IntentAPI) Member(roomID id.RoomID, userID id.UserID) *event.MemberEventContent {
 	member, ok := intent.as.StateStore.TryGetMember(roomID, userID)
 	if !ok {
diff --git a/client.go b/client.go
index 45bdf5a..cac8d91 100644
--- a/client.go
+++ b/client.go
@@ -1078,6 +1078,55 @@ func (cli *Client) StateEvent(roomID id.RoomID, eventType event.Type, stateKey s
 	return
 }
 
+// parseRoomStateArray parses a JSON array as a stream and stores the events inside it in a room state map.
+func parseRoomStateArray(_ *http.Request, res *http.Response, responseJSON interface{}) ([]byte, error) {
+	response := make(RoomStateMap)
+	responsePtr := responseJSON.(*interface{})
+	*responsePtr = response
+	dec := json.NewDecoder(res.Body)
+
+	arrayStart, err := dec.Token()
+	if err != nil {
+		return nil, err
+	} else if arrayStart != json.Delim('[') {
+		return nil, fmt.Errorf("expected array start, got %+v", arrayStart)
+	}
+
+	for i := 1; dec.More(); i++ {
+		var evt *event.Event
+		err = dec.Decode(&evt)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse state array item #%d: %v", i, err)
+		}
+		subMap, ok := response[evt.Type]
+		if !ok {
+			subMap = make(map[string]*event.Event)
+			response[evt.Type] = subMap
+		}
+		subMap[*evt.StateKey] = evt
+	}
+
+	arrayEnd, err := dec.Token()
+	if err != nil {
+		return nil, err
+	} else if arrayEnd != json.Delim(']') {
+		return nil, fmt.Errorf("expected array end, got %+v", arrayStart)
+	}
+	return nil, nil
+}
+
+// State gets all state in a room.
+// See https://matrix.org/docs/spec/client_server/r0.6.1#get-matrix-client-r0-rooms-roomid-state
+func (cli *Client) State(roomID id.RoomID) (stateMap RoomStateMap, err error) {
+	_, err = cli.MakeFullRequest(FullRequest{
+		Method:       http.MethodGet,
+		URL:          cli.BuildURL("rooms", roomID, "state"),
+		ResponseJSON: &stateMap,
+		Handler:      parseRoomStateArray,
+	})
+	return
+}
+
 // UploadLink uploads an HTTP URL and then returns an MXC URI.
 func (cli *Client) UploadLink(link string) (*RespMediaUpload, error) {
 	res, err := cli.Client.Get(link)
diff --git a/room.go b/room.go
index 588097f..c3ddb7e 100644
--- a/room.go
+++ b/room.go
@@ -5,10 +5,12 @@ import (
 	"maunium.net/go/mautrix/id"
 )
 
+type RoomStateMap = map[event.Type]map[string]*event.Event
+
 // Room represents a single Matrix room.
 type Room struct {
 	ID    id.RoomID
-	State map[event.Type]map[string]*event.Event
+	State RoomStateMap
 }
 
 // UpdateState updates the room's current state with the given Event. This will clobber events based
@@ -47,6 +49,6 @@ func NewRoom(roomID id.RoomID) *Room {
 	// Init the State map and return a pointer to the Room
 	return &Room{
 		ID:    roomID,
-		State: make(map[event.Type]map[string]*event.Event),
+		State: make(RoomStateMap),
 	}
 }