mrvc/roominfotree/roominfotree.go

package roominfotree

import (
	"context"
	"log"
	"strings"
	"sync"

	"github.com/matrix-org/gomatrixserverlib/fclient"
	"github.com/matrix-org/gomatrixserverlib/spec"
	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/id"
)

type homeserverServerVersionInfo struct {
	homeserver        string
	serverVersionInfo fclient.Version
}

type RoomInfoTree map[id.RoomID](*RoomInfo)

type RoomInfo struct {
	MemberCount     uint                             `json:"memberCount"`
	GivenAliases    []string                         `json:"givenAliases"`
	Aliases         []string                         `json:"aliases"`
	MaxRoomVersions map[string](*MaxRoomVersionInfo) `json:"maxRoomVersions"`
}

type MaxRoomVersionInfo struct {
	MemberCount uint                     `json:"memberCount"`
	Servers     map[string](*ServerInfo) `json:"servers"`
}

type ServerInfo struct {
	MemberCount uint                      `json:"memberCount"`
	Versions    map[string](*VersionInfo) `json:"versions"`
}

type VersionInfo struct {
	MemberCount uint                         `json:"memberCount"`
	Homeservers map[string](*HomeserverInfo) `json:"homeservers,omitempty"`
}

type HomeserverInfo struct {
	MemberCount uint `json:"memberCount"`
}

var unknownServerVersionInfo = fclient.Version{
	Server: struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	}{
		Name:    "unknown",
		Version: "unknown",
	},
}

// Resolves the given list of rooms to a list of room IDs and a map of given aliases by room ID.
func resolveRooms(rooms []string, client *mautrix.Client) ([]id.RoomID, map[id.RoomID]([]string)) {
	roomIDSet := make(map[id.RoomID]bool)
	aliasSetByRoomID := make(map[id.RoomID](map[string]bool))
	for _, room := range rooms {
		// Check, if given room is an alias and try to resolve it into a room id.
		if strings.HasPrefix(room, "#") {
			resolvedAlias, err := client.ResolveAlias(context.Background(), id.RoomAlias(room))
			if err != nil {
				log.Printf("Error resolving given alias (%s) to room ID.\n", room)
				log.Fatal(err)
			}

			roomIDSet[resolvedAlias.RoomID] = true

			aliasSet, ok := aliasSetByRoomID[resolvedAlias.RoomID]
			if !ok {
				aliasSet = make(map[string]bool)
				aliasSetByRoomID[resolvedAlias.RoomID] = aliasSet
			}
			aliasSet[room] = true
		} else {
			roomIDSet[id.RoomID(room)] = true
		}
	}

	roomIDs := make([]id.RoomID, 0, len(roomIDSet))
	for roomID := range roomIDSet {
		roomIDs = append(roomIDs, roomID)
	}

	aliasesByRoomID := make(map[id.RoomID]([]string))
	for roomID, aliasSet := range aliasSetByRoomID {
		aliases := make([]string, 0, len(aliasSet))
		for alias := range aliasSet {
			aliases = append(aliases, alias)
		}

		aliasesByRoomID[roomID] = aliases
	}

	return roomIDs, aliasesByRoomID
}

func getChildRoomChunks(roomID id.RoomID, recursionMaxDepth *int, recursionSuggestedOnly bool, client *mautrix.Client) []*mautrix.ChildRoomsChunk {
	var childRoomChunks []*mautrix.ChildRoomsChunk

	hierarchyRequest := mautrix.ReqHierarchy{
		MaxDepth:      recursionMaxDepth,
		SuggestedOnly: recursionSuggestedOnly,
	}
	finishedRequesting := false
	for !finishedRequesting {
		hierarchyResponse, err := client.Hierarchy(context.Background(), roomID, &hierarchyRequest)
		if err != nil {
			log.Printf("Error getting child rooms for given room (%s).\n", roomID)
			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, recursionMaxDepth *int, recursionSuggestedOnly bool, client *mautrix.Client) []id.RoomID {
	roomIDSet := make(map[id.RoomID]bool)

	for _, givenRoomID := range givenRoomIDs {
		roomIDSet[givenRoomID] = true

		childRoomChunks := getChildRoomChunks(givenRoomID, recursionMaxDepth, recursionSuggestedOnly, 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 getAliasesByRoomID(roomIDs []id.RoomID, client *mautrix.Client) map[id.RoomID]([]id.RoomAlias) {
	aliasesByRoomID := make(map[id.RoomID]([]id.RoomAlias))

	for _, roomID := range roomIDs {
		aliasListResp, err := client.GetAliases(context.Background(), roomID)
		if err != nil {
			log.Printf("Error getting aliases for room (%s).\n", roomID)
			log.Fatal(err)
		}

		aliasesByRoomID[roomID] = aliasListResp.Aliases
	}

	return aliasesByRoomID
}

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 {
		joinedMembers, err := client.JoinedMembers(context.Background(), roomID)
		if err != nil {
			log.Printf("Error getting members for room (%s).\n", roomID)
			log.Fatal(err)
		}
		joinedMembersByRoomID[roomID] = joinedMembers
	}

	membersByHomeserverByRoomID := make(map[id.RoomID](map[string]uint))
	for roomID, joinedMembers := range joinedMembersByRoomID {
		membersByHomeserver := make(map[string]uint)
		for userID := range joinedMembers.Joined {
			membersByHomeserver[userID.Homeserver()]++
		}
		membersByHomeserverByRoomID[roomID] = membersByHomeserver
	}

	return membersByHomeserverByRoomID
}

// Gets the homeservers appearing in the given membersByHomeserverByRoomID map tree.
func getHomeservers(membersByHomeserverByRoomID map[id.RoomID]map[string]uint) []string {
	homeserverSet := make(map[string]bool)
	for _, membersByHomeserver := range membersByHomeserverByRoomID {
		for hs := range membersByHomeserver {
			homeserverSet[hs] = true
		}
	}

	homeservers := make([]string, 0, len(homeserverSet))
	for hs := range homeserverSet {
		homeservers = append(homeservers, hs)
	}

	return homeservers
}

func getServerVersionInfoByHomeserver(homeservers []string, federationClient *fclient.Client) map[string](fclient.Version) {
	homeserverChannel := make(chan string)
	go func() {
		for _, hs := range homeservers {
			homeserverChannel <- hs
		}
		close(homeserverChannel)
	}()

	homeserverSVInfoChannel := make(chan homeserverServerVersionInfo)
	var wg sync.WaitGroup
	const numHomeserverSVInfoReceivers = 100
	wg.Add(numHomeserverSVInfoReceivers)
	for i := 0; i < numHomeserverSVInfoReceivers; i++ {
		go func() {
			for hs := range homeserverChannel {
				serverVersionInfo, err := federationClient.GetVersion(context.Background(), spec.ServerName(hs))
				if err != nil {
					serverVersionInfo = unknownServerVersionInfo
				}
				homeserverSVInfoChannel <- homeserverServerVersionInfo{hs, serverVersionInfo}
			}
			wg.Done()
		}()
	}

	go func() {
		wg.Wait()
		close(homeserverSVInfoChannel)
	}()

	serverVersionInfoByHomeserver := make(map[string](fclient.Version))
	for homeserverSVInfo := range homeserverSVInfoChannel {
		serverVersionInfoByHomeserver[homeserverSVInfo.homeserver] = homeserverSVInfo.serverVersionInfo
	}
	return serverVersionInfoByHomeserver
}

func Get(rooms []string, recursive bool, recursionMaxDepth *int, recursionSuggestedOnly bool, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
	roomIDs, givenAliasesByRoomID := resolveRooms(rooms, client)
	if recursive {
		roomIDs = addChildRooms(roomIDs, recursionMaxDepth, recursionSuggestedOnly, client)
	}
	aliasesByRoomID := getAliasesByRoomID(roomIDs, client)
	membersByHomeserverByRoomID := getMembersByHomeserverByRoomID(roomIDs, client)
	homeservers := getHomeservers(membersByHomeserverByRoomID)
	serverVersionInfoByHomeserver := getServerVersionInfoByHomeserver(homeservers, federationClient)

	roomInfoTree := make(RoomInfoTree)

	for roomID, membersByHomeserver := range membersByHomeserverByRoomID {
		roomInfo, ok := roomInfoTree[roomID]
		if !ok {
			roomInfo = &RoomInfo{
				GivenAliases:    givenAliasesByRoomID[roomID],
				MaxRoomVersions: make(map[string]*MaxRoomVersionInfo),
			}

			aliases := aliasesByRoomID[roomID]
			aliasStrings := make([]string, 0, len(aliases))
			for _, alias := range aliases {
				aliasStrings = append(aliasStrings, alias.String())
			}
			roomInfo.Aliases = aliasStrings

			roomInfoTree[roomID] = roomInfo
		}

		for hs, members := range membersByHomeserver {
			serverVersionInfo := serverVersionInfoByHomeserver[hs]

			maxRoomVersion := getMaxRoomVersion(serverVersionInfo)

			// Sort into roomInfoTree and add to counters.
			maxRoomVersionInfo, ok := roomInfo.MaxRoomVersions[maxRoomVersion]
			if !ok {
				maxRoomVersionInfo = &MaxRoomVersionInfo{Servers: make(map[string]*ServerInfo)}
				roomInfo.MaxRoomVersions[maxRoomVersion] = maxRoomVersionInfo
			}
			serverInfo, ok := maxRoomVersionInfo.Servers[serverVersionInfo.Server.Name]
			if !ok {
				serverInfo = &ServerInfo{Versions: make(map[string]*VersionInfo)}
				maxRoomVersionInfo.Servers[serverVersionInfo.Server.Name] = serverInfo
			}
			versionInfo, ok := serverInfo.Versions[serverVersionInfo.Server.Version]
			if !ok {
				versionInfo = &VersionInfo{Homeservers: make(map[string]*HomeserverInfo)}
				serverInfo.Versions[serverVersionInfo.Server.Version] = versionInfo
			}
			homeserverInfo, ok := versionInfo.Homeservers[hs]
			if !ok {
				homeserverInfo = &HomeserverInfo{}
				versionInfo.Homeservers[hs] = homeserverInfo
			}

			homeserverInfo.MemberCount = members
			versionInfo.MemberCount += members
			serverInfo.MemberCount += members
			maxRoomVersionInfo.MemberCount += members
			roomInfo.MemberCount += members
		}
	}

	return roomInfoTree
}