introduce option for recursively checking all child rooms as well

README.md

@@ -50,6 +50,7 @@ Version Support:
 
 - Show member count broken down by maximum room version, Matrix server implementation and version and optionally homeserver name.
 - Check multiple rooms at once by setting the `--room` flag multiple times.
+- Recursively check all the child rooms for the given rooms (spaces) by setting the `--recursive` flag.
 
 ## Usage
 

config/config.go

@@ -23,6 +23,7 @@ type configOptions struct {
 	userID                       configOption[string]
 	token                        configOption[string]
 	rooms                        configOption[rooms]
+	recursive                    configOption[bool]
 	printHomeserverMemberCount   configOption[bool]
 	homeserverVersionInfoTimeout configOption[time.Duration]
 }
@@ -31,6 +32,7 @@ type Config struct {
 	UserID                       string
 	Token                        string
 	Rooms                        rooms
+	Recursive                    bool
 	PrintHomeserverMemberCount   bool
 	HomeserverVersionInfoTimeout time.Duration
 }
@@ -39,6 +41,7 @@ var configOpts = configOptions{
 	userID:                       configOption[string]{"user-id", "", "The Matrix user ID to use.", "MRVC_USER_ID"},
 	token:                        configOption[string]{"token", "", "The Matrix access token to use.", "MRVC_TOKEN"},
 	rooms:                        configOption[rooms]{"room", []string{}, "The Matrix room to check. The flag can be set multiple times to check multiple rooms.", "MRVC_ROOM"},
+	recursive:                    configOption[bool]{"recursive", false, "Recursively check the child rooms for the given rooms (spaces) as well.", "MRVC_RECURSIVE"},
 	printHomeserverMemberCount:   configOption[bool]{"print-homeserver-member-count", false, "Print the member count for each homeserver.", "MRVC_PRINT_HOMESERVER_MEMBER_COUNT"},
 	homeserverVersionInfoTimeout: configOption[time.Duration]{"homeserver-version-info-timeout", time.Second * 5, "Timeout for getting the homeservers version information per homeserver.", "MRVC_HOMESERVER_VERSION_INFO_TIMEOUT"},
 }
@@ -85,6 +88,7 @@ func (configOpt configOption[T]) getConfigValueWithError(configFlag *T, visitedF
 
 var userIdFlag = flag.String(configOpts.userID.getFlagArgs())
 var tokenFlag = flag.String(configOpts.token.getFlagArgs())
+var recursiveFlag = flag.Bool(configOpts.recursive.getFlagArgs())
 var printHomeserverMemberCountFlag = flag.Bool(configOpts.printHomeserverMemberCount.getFlagArgs())
 var homeserverVersionInfoTimeoutFlag = flag.Duration(configOpts.homeserverVersionInfoTimeout.getFlagArgs())
 
@@ -118,6 +122,14 @@ func Get() Config {
 	if err != nil {
 		log.Fatal("A Matrix room must be provided.")
 	}
+	config.Recursive = configOpts.recursive.getConfigValueWithDefault(recursiveFlag, visitedFlags, func(envVar string) bool {
+		parsedEnvVar, err := strconv.ParseBool(envVar)
+		if err != nil {
+			log.Printf("Error parsing %s:\n", configOpts.recursive.envVarName)
+			log.Fatal(err)
+		}
+		return parsedEnvVar
+	})
 	config.PrintHomeserverMemberCount = configOpts.printHomeserverMemberCount.getConfigValueWithDefault(printHomeserverMemberCountFlag, visitedFlags, func(envVar string) bool {
 		parsedEnvVar, err := strconv.ParseBool(envVar)
 		if err != nil {

main.go

@@ -73,7 +73,7 @@ func main() {
 		fclient.WithTimeout(config.HomeserverVersionInfoTimeout),
 	)
 
-	roomInfoTree := roominfotree.Get(config.Rooms, client, federationClient)
+	roomInfoTree := roominfotree.Get(config.Rooms, config.Recursive, client, federationClient)
 
 	for roomID, roomInfo := range roomInfoTree {
 		fmt.Println("Room:")

roominfotree/roominfotree.go

@@ -97,6 +97,60 @@ func resolveRooms(rooms []string, client *mautrix.Client) ([]id.RoomID, map[id.R
 	return roomIDs, aliasesByRoomID
 }
 
+func getChildRoomChunks(roomID id.RoomID, client *mautrix.Client) []*mautrix.ChildRoomsChunk {
+	var maxDepth *int
+	maxDepth = nil
+	suggestedOnly := false
+
+	var childRoomChunks []*mautrix.ChildRoomsChunk
+
+	hierarchyRequest := mautrix.ReqHierarchy{
+		MaxDepth:      maxDepth,
+		SuggestedOnly: suggestedOnly,
+	}
+	finishedRequesting := false
+	for !finishedRequesting {
+		hierarchyResponse, err := client.Hierarchy(context.Background(), roomID, &hierarchyRequest)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		childRoomChunks = append(childRoomChunks, hierarchyResponse.Rooms...)
+
+		if hierarchyResponse.NextBatch == "" {
+			finishedRequesting = true
+		} else {
+			hierarchyRequest = mautrix.ReqHierarchy{
+				From:          hierarchyResponse.NextBatch,
+				Limit:         hierarchyRequest.Limit,
+				MaxDepth:      hierarchyRequest.MaxDepth,
+				SuggestedOnly: hierarchyRequest.SuggestedOnly,
+			}
+		}
+	}
+
+	return childRoomChunks
+}
+
+func addChildRooms(givenRoomIDs []id.RoomID, client *mautrix.Client) []id.RoomID {
+	roomIDSet := make(map[id.RoomID]bool)
+
+	for _, givenRoomID := range givenRoomIDs {
+		roomIDSet[givenRoomID] = true
+
+		childRoomChunks := getChildRoomChunks(givenRoomID, client)
+		for _, childRoomChunk := range childRoomChunks {
+			roomIDSet[childRoomChunk.PublicRoomInfo.RoomID] = true
+		}
+	}
+
+	roomIDs := make([]id.RoomID, 0, len(roomIDSet))
+	for roomID := range roomIDSet {
+		roomIDs = append(roomIDs, roomID)
+	}
+	return roomIDs
+}
+
 func getMembersByHomeserverByRoomID(roomIDs []id.RoomID, client *mautrix.Client) map[id.RoomID](map[string]uint) {
 	joinedMembersByRoomID := make(map[id.RoomID]*mautrix.RespJoinedMembers)
 	for _, roomID := range roomIDs {
@@ -174,8 +228,11 @@ func getServerVersionInfoByHomeserver(homeservers []string, federationClient *fc
 	return serverVersionInfoByHomeserver
 }
 
-func Get(rooms []string, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
+func Get(rooms []string, recursive bool, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
 	roomIDs, aliasesByRoomID := resolveRooms(rooms, client)
+	if recursive {
+		roomIDs = addChildRooms(roomIDs, client)
+	}
 	membersByHomeserverByRoomID := getMembersByHomeserverByRoomID(roomIDs, client)
 	homeservers := getHomeservers(membersByHomeserverByRoomID)
 	serverVersionInfoByHomeserver := getServerVersionInfoByHomeserver(homeservers, federationClient)