From 03029711fe53d29c7c95f2f50e2c1face97256f5 Mon Sep 17 00:00:00 2001
From: chenyuqide <chenyuqide@outlook.com>
Date: Sun, 10 Apr 2022 18:57:15 +0800
Subject: [PATCH] Add client space api '/rooms/{roomId}/hierarchy'

---
 src/api/client_server/mod.rs   |   2 +
 src/api/client_server/space.rs | 242 +++++++++++++++++++++++++++++++++
 src/main.rs                    |   1 +
 3 files changed, 245 insertions(+)
 create mode 100644 src/api/client_server/space.rs

diff --git a/src/api/client_server/mod.rs b/src/api/client_server/mod.rs
index 6ed17e76..4cc000ed 100644
--- a/src/api/client_server/mod.rs
+++ b/src/api/client_server/mod.rs
@@ -20,6 +20,7 @@ mod report;
 mod room;
 mod search;
 mod session;
+mod space;
 mod state;
 mod sync;
 mod tag;
@@ -52,6 +53,7 @@ pub use report::*;
 pub use room::*;
 pub use search::*;
 pub use session::*;
+pub use space::*;
 pub use state::*;
 pub use sync::*;
 pub use tag::*;
diff --git a/src/api/client_server/space.rs b/src/api/client_server/space.rs
new file mode 100644
index 00000000..8cc9221f
--- /dev/null
+++ b/src/api/client_server/space.rs
@@ -0,0 +1,242 @@
+use std::{collections::HashSet, sync::Arc};
+
+use crate::{services, Error, PduEvent, Result, Ruma};
+use ruma::{
+    api::client::{
+        error::ErrorKind,
+        space::{get_hierarchy, SpaceHierarchyRoomsChunk, SpaceRoomJoinRule},
+    },
+    events::{
+        room::{
+            avatar::RoomAvatarEventContent,
+            canonical_alias::RoomCanonicalAliasEventContent,
+            create::RoomCreateEventContent,
+            guest_access::{GuestAccess, RoomGuestAccessEventContent},
+            history_visibility::{HistoryVisibility, RoomHistoryVisibilityEventContent},
+            join_rules::{JoinRule, RoomJoinRulesEventContent},
+            name::RoomNameEventContent,
+            topic::RoomTopicEventContent,
+        },
+        space::child::{HierarchySpaceChildEvent, SpaceChildEventContent},
+        StateEventType,
+    },
+    serde::Raw,
+    MilliSecondsSinceUnixEpoch, OwnedRoomId, RoomId,
+};
+use serde_json;
+use tracing::warn;
+
+/// # `GET /_matrix/client/v1/rooms/{room_id}/hierarchy``
+///
+/// Paginates over the space tree in a depth-first manner to locate child rooms of a given space.
+///
+/// - TODO: Use federation for unknown room.
+///
+pub async fn get_hierarchy_route(
+    body: Ruma<get_hierarchy::v1::IncomingRequest>,
+) -> Result<get_hierarchy::v1::Response> {
+    // from format is '{suggested_only}|{max_depth}|{skip}'
+    let (suggested_only, max_depth, start) = body
+        .from
+        .as_ref()
+        .map_or(
+            Some((
+                body.suggested_only,
+                body.max_depth
+                    .map_or(services().globals.hierarchy_max_depth(), |v| v.into())
+                    .min(services().globals.hierarchy_max_depth()),
+                0,
+            )),
+            |from| {
+                let mut p = from.split('|');
+                Some((
+                    p.next()?.trim().parse().ok()?,
+                    p.next()?
+                        .trim()
+                        .parse::<u64>()
+                        .ok()?
+                        .min(services().globals.hierarchy_max_depth()),
+                    p.next()?.trim().parse().ok()?,
+                ))
+            },
+        )
+        .ok_or(Error::BadRequest(ErrorKind::InvalidParam, "Invalid from"))?;
+
+    let limit = body.limit.map_or(20u64, |v| v.into()) as usize;
+    let mut skip = start;
+
+    // Set for avoid search in loop.
+    let mut room_set = HashSet::new();
+    let mut rooms_chunk: Vec<SpaceHierarchyRoomsChunk> = vec![];
+    let mut stack = vec![(0, body.room_id.clone())];
+
+    while let (Some((depth, room_id)), true) = (stack.pop(), rooms_chunk.len() < limit) {
+        let (childern, pdus): (Vec<_>, Vec<_>) = services()
+            .rooms
+            .state_accessor
+            .room_state_full(&room_id)
+            .await?
+            .into_iter()
+            .filter_map(|((e_type, key), pdu)| {
+                (e_type == StateEventType::SpaceChild && !room_set.contains(&room_id))
+                    .then_some((key, pdu))
+            })
+            .unzip();
+
+        if skip == 0 {
+            if rooms_chunk.len() < limit {
+                room_set.insert(room_id.clone());
+                rooms_chunk.push(get_room_chunk(room_id, suggested_only, pdus).await?);
+            }
+        } else {
+            skip -= 1;
+        }
+
+        if depth < max_depth {
+            childern.into_iter().rev().for_each(|key| {
+                stack.push((depth + 1, RoomId::parse(key).unwrap()));
+            });
+        }
+    }
+
+    Ok(get_hierarchy::v1::Response {
+        next_batch: (!stack.is_empty()).then_some(format!(
+            "{}|{}|{}",
+            suggested_only,
+            max_depth,
+            start + limit
+        )),
+        rooms: rooms_chunk,
+    })
+}
+
+async fn get_room_chunk(
+    room_id: OwnedRoomId,
+    suggested_only: bool,
+    phus: Vec<Arc<PduEvent>>,
+) -> Result<SpaceHierarchyRoomsChunk> {
+    Ok(SpaceHierarchyRoomsChunk {
+        canonical_alias: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomCanonicalAlias, "")
+            .ok()
+            .and_then(|s| {
+                serde_json::from_str(s?.content.get())
+                    .map(|c: RoomCanonicalAliasEventContent| c.alias)
+                    .ok()?
+            }),
+        name: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomName, "")
+            .ok()
+            .flatten()
+            .and_then(|s| {
+                serde_json::from_str(s.content.get())
+                    .map(|c: RoomNameEventContent| c.name)
+                    .ok()?
+            }),
+        num_joined_members: services()
+            .rooms
+            .state_cache
+            .room_joined_count(&room_id)?
+            .unwrap_or_else(|| {
+                warn!("Room {} has no member count", &room_id);
+                0
+            })
+            .try_into()
+            .expect("user count should not be that big"),
+        topic: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomTopic, "")
+            .ok()
+            .and_then(|s| {
+                serde_json::from_str(s?.content.get())
+                    .ok()
+                    .map(|c: RoomTopicEventContent| c.topic)
+            }),
+        world_readable: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomHistoryVisibility, "")?
+            .map_or(Ok(false), |s| {
+                serde_json::from_str(s.content.get())
+                    .map(|c: RoomHistoryVisibilityEventContent| {
+                        c.history_visibility == HistoryVisibility::WorldReadable
+                    })
+                    .map_err(|_| {
+                        Error::bad_database("Invalid room history visibility event in database.")
+                    })
+            })?,
+        guest_can_join: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomGuestAccess, "")?
+            .map_or(Ok(false), |s| {
+                serde_json::from_str(s.content.get())
+                    .map(|c: RoomGuestAccessEventContent| c.guest_access == GuestAccess::CanJoin)
+                    .map_err(|_| {
+                        Error::bad_database("Invalid room guest access event in database.")
+                    })
+            })?,
+        avatar_url: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomAvatar, "")
+            .ok()
+            .and_then(|s| {
+                serde_json::from_str(s?.content.get())
+                    .map(|c: RoomAvatarEventContent| c.url)
+                    .ok()?
+            }),
+        join_rule: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomJoinRules, "")?
+            .map(|s| {
+                serde_json::from_str(s.content.get())
+                    .map(|c: RoomJoinRulesEventContent| match c.join_rule {
+                        JoinRule::Invite => SpaceRoomJoinRule::Invite,
+                        JoinRule::Knock => SpaceRoomJoinRule::Knock,
+                        JoinRule::Private => SpaceRoomJoinRule::Private,
+                        JoinRule::Public => SpaceRoomJoinRule::Public,
+                        JoinRule::Restricted(_) => SpaceRoomJoinRule::Restricted,
+                        // Can't convert two type.
+                        JoinRule::_Custom(_) => SpaceRoomJoinRule::Private,
+                    })
+                    .map_err(|_| Error::bad_database("Invalid room join rules event in database."))
+            })
+            .ok_or_else(|| Error::bad_database("Invalid room join rules event in database."))??,
+        room_type: services()
+            .rooms
+            .state_accessor
+            .room_state_get(&room_id, &StateEventType::RoomCreate, "")
+            .map(|s| {
+                serde_json::from_str(s?.content.get())
+                    .map(|c: RoomCreateEventContent| c.room_type)
+                    .ok()?
+            })
+            .ok()
+            .flatten(),
+        children_state: phus
+            .into_iter()
+            .flat_map(|pdu| {
+                Some(HierarchySpaceChildEvent {
+                    // Ignore unsuggested rooms if suggested_only is set
+                    content: serde_json::from_str(pdu.content.get()).ok().filter(
+                        |pdu: &SpaceChildEventContent| {
+                            !suggested_only || pdu.suggested.unwrap_or(false)
+                        },
+                    )?,
+                    sender: pdu.sender.clone(),
+                    state_key: pdu.state_key.clone()?,
+                    origin_server_ts: MilliSecondsSinceUnixEpoch(pdu.origin_server_ts),
+                })
+            })
+            .filter_map(|hsce| Raw::<HierarchySpaceChildEvent>::new(&hsce).ok())
+            .collect::<Vec<_>>(),
+        room_id,
+    })
+}
diff --git a/src/main.rs b/src/main.rs
index e754b84f..75550231 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -318,6 +318,7 @@ fn routes() -> Router {
         .ruma_route(client_server::send_state_event_for_key_route)
         .ruma_route(client_server::get_state_events_route)
         .ruma_route(client_server::get_state_events_for_key_route)
+        .ruma_route(client_server::get_hierarchy_route)
         // Ruma doesn't have support for multiple paths for a single endpoint yet, and these routes
         // share one Ruma request / response type pair with {get,send}_state_event_for_key_route
         .route(