Merge branch 'main' into site-v2

server/src/auth.rs

@@ -18,7 +18,7 @@ use scrypt::{
 use serde::{Deserialize, Serialize};
 use std::{borrow::Cow, convert::TryFrom, sync::Arc};
 use surf::{StatusCode, Url};
-use tide::Server;
+use tide::{log, Server};
 use zrpc::auth as zed_auth;
 
 static CURRENT_GITHUB_USER: &'static str = "current_github_user";
@@ -121,6 +121,7 @@ pub fn add_routes(app: &mut Server<Arc<AppState>>) {
 struct NativeAppSignInParams {
     native_app_port: String,
     native_app_public_key: String,
+    impersonate: Option<String>,
 }
 
 async fn get_sign_in(mut request: Request) -> tide::Result {
@@ -142,11 +143,15 @@ async fn get_sign_in(mut request: Request) -> tide::Result {
 
     let app_sign_in_params: Option<NativeAppSignInParams> = request.query().ok();
     if let Some(query) = app_sign_in_params {
-        redirect_url
-            .query_pairs_mut()
+        let mut redirect_query = redirect_url.query_pairs_mut();
+        redirect_query
             .clear()
             .append_pair("native_app_port", &query.native_app_port)
             .append_pair("native_app_public_key", &query.native_app_public_key);
+
+        if let Some(impersonate) = &query.impersonate {
+            redirect_query.append_pair("impersonate", impersonate);
+        }
     }
 
     let (auth_url, csrf_token) = request
@@ -222,7 +227,20 @@ async fn get_auth_callback(mut request: Request) -> tide::Result {
     // When signing in from the native app, generate a new access token for the current user. Return
     // a redirect so that the user's browser sends this access token to the locally-running app.
     if let Some((user, app_sign_in_params)) = user.zip(query.native_app_sign_in_params) {
-        let access_token = create_access_token(request.db(), user.id).await?;
+        let mut user_id = user.id;
+        if let Some(impersonated_login) = app_sign_in_params.impersonate {
+            log::info!("attempting to impersonate user @{}", impersonated_login);
+            if let Some(user) = request.db().get_users_by_ids([user_id]).await?.first() {
+                if user.admin {
+                    user_id = request.db().create_user(&impersonated_login, false).await?;
+                    log::info!("impersonating user {}", user_id.0);
+                } else {
+                    log::info!("refusing to impersonate user");
+                }
+            }
+        }
+
+        let access_token = create_access_token(request.db(), user_id).await?;
         let native_app_public_key =
             zed_auth::PublicKey::try_from(app_sign_in_params.native_app_public_key.clone())
                 .context("failed to parse app public key")?;
@@ -232,7 +250,7 @@ async fn get_auth_callback(mut request: Request) -> tide::Result {
 
         return Ok(tide::Redirect::new(&format!(
             "http://127.0.0.1:{}?user_id={}&access_token={}",
-            app_sign_in_params.native_app_port, user.id.0, encrypted_access_token,
+            app_sign_in_params.native_app_port, user_id.0, encrypted_access_token,
         ))
         .into());
     }

server/src/bin/seed.rs

@@ -27,8 +27,12 @@ async fn main() {
     let zed_users = ["nathansobo", "maxbrunsfeld", "as-cii", "iamnbutler"];
     let mut zed_user_ids = Vec::<UserId>::new();
     for zed_user in zed_users {
-        if let Some(user_id) = db.get_user(zed_user).await.expect("failed to fetch user") {
-            zed_user_ids.push(user_id);
+        if let Some(user) = db
+            .get_user_by_github_login(zed_user)
+            .await
+            .expect("failed to fetch user")
+        {
+            zed_user_ids.push(user.id);
         } else {
             zed_user_ids.push(
                 db.create_user(zed_user, true)

server/src/db.rs

@@ -97,27 +97,12 @@ impl Db {
 
     // users
 
-    #[allow(unused)] // Help rust-analyzer
-    #[cfg(any(test, feature = "seed-support"))]
-    pub async fn get_user(&self, github_login: &str) -> Result<Option<UserId>> {
-        test_support!(self, {
-            let query = "
-                SELECT id
-                FROM users
-                WHERE github_login = $1
-            ";
-            sqlx::query_scalar(query)
-                .bind(github_login)
-                .fetch_optional(&self.pool)
-                .await
-        })
-    }
-
     pub async fn create_user(&self, github_login: &str, admin: bool) -> Result<UserId> {
         test_support!(self, {
             let query = "
                 INSERT INTO users (github_login, admin)
                 VALUES ($1, $2)
+                ON CONFLICT (github_login) DO UPDATE SET github_login = excluded.github_login
                 RETURNING id
             ";
             sqlx::query_scalar(query)
@@ -138,51 +123,17 @@ impl Db {
 
     pub async fn get_users_by_ids(
         &self,
-        requester_id: UserId,
-        ids: impl Iterator<Item = UserId>,
+        ids: impl IntoIterator<Item = UserId>,
     ) -> Result<Vec<User>> {
-        let mut include_requester = false;
-        let ids = ids
-            .map(|id| {
-                if id == requester_id {
-                    include_requester = true;
-                }
-                id.0
-            })
-            .collect::<Vec<_>>();
-
+        let ids = ids.into_iter().map(|id| id.0).collect::<Vec<_>>();
         test_support!(self, {
-            // Only return users that are in a common channel with the requesting user.
-            // Also allow the requesting user to return their own data, even if they aren't
-            // in any channels.
             let query = "
-                SELECT
-                    users.*
-                FROM
-                    users, channel_memberships
-                WHERE
-                    users.id = ANY ($1) AND
-                    channel_memberships.user_id = users.id AND
-                    channel_memberships.channel_id IN (
-                        SELECT channel_id
-                        FROM channel_memberships
-                        WHERE channel_memberships.user_id = $2
-                    )
-                UNION
-                SELECT
-                    users.*
-                FROM
-                    users
-                WHERE
-                    $3 AND users.id = $2
+                SELECT users.*
+                FROM users
+                WHERE users.id = ANY ($1)
             ";
 
-            sqlx::query_as(query)
-                .bind(&ids)
-                .bind(requester_id)
-                .bind(include_requester)
-                .fetch_all(&self.pool)
-                .await
+            sqlx::query_as(query).bind(&ids).fetch_all(&self.pool).await
         })
     }
 
@@ -613,45 +564,11 @@ pub mod tests {
         let friend1 = db.create_user("friend-1", false).await.unwrap();
         let friend2 = db.create_user("friend-2", false).await.unwrap();
         let friend3 = db.create_user("friend-3", false).await.unwrap();
-        let stranger = db.create_user("stranger", false).await.unwrap();
-
-        // A user can read their own info, even if they aren't in any channels.
-        assert_eq!(
-            db.get_users_by_ids(
-                user,
-                [user, friend1, friend2, friend3, stranger].iter().copied()
-            )
-            .await
-            .unwrap(),
-            vec![User {
-                id: user,
-                github_login: "user".to_string(),
-                admin: false,
-            },],
-        );
-
-        // A user can read the info of any other user who is in a shared channel
-        // with them.
-        let org = db.create_org("test org", "test-org").await.unwrap();
-        let chan1 = db.create_org_channel(org, "channel-1").await.unwrap();
-        let chan2 = db.create_org_channel(org, "channel-2").await.unwrap();
-        let chan3 = db.create_org_channel(org, "channel-3").await.unwrap();
-
-        db.add_channel_member(chan1, user, false).await.unwrap();
-        db.add_channel_member(chan2, user, false).await.unwrap();
-        db.add_channel_member(chan1, friend1, false).await.unwrap();
-        db.add_channel_member(chan1, friend2, false).await.unwrap();
-        db.add_channel_member(chan2, friend2, false).await.unwrap();
-        db.add_channel_member(chan2, friend3, false).await.unwrap();
-        db.add_channel_member(chan3, stranger, false).await.unwrap();
 
         assert_eq!(
-            db.get_users_by_ids(
-                user,
-                [user, friend1, friend2, friend3, stranger].iter().copied()
-            )
-            .await
-            .unwrap(),
+            db.get_users_by_ids([user, friend1, friend2, friend3])
+                .await
+                .unwrap(),
             vec![
                 User {
                     id: user,
@@ -675,18 +592,6 @@ pub mod tests {
                 }
             ]
         );
-
-        // The user's own info is only returned if they request it.
-        assert_eq!(
-            db.get_users_by_ids(user, [friend1].iter().copied())
-                .await
-                .unwrap(),
-            vec![User {
-                id: friend1,
-                github_login: "friend-1".to_string(),
-                admin: false,
-            },]
-        )
     }
 
     #[gpui::test]

server/src/rpc/store.rs

@@ -0,0 +1,615 @@
+use crate::db::{ChannelId, UserId};
+use anyhow::anyhow;
+use std::collections::{hash_map, HashMap, HashSet};
+use zrpc::{proto, ConnectionId};
+
+#[derive(Default)]
+pub struct Store {
+    connections: HashMap<ConnectionId, ConnectionState>,
+    connections_by_user_id: HashMap<UserId, HashSet<ConnectionId>>,
+    worktrees: HashMap<u64, Worktree>,
+    visible_worktrees_by_user_id: HashMap<UserId, HashSet<u64>>,
+    channels: HashMap<ChannelId, Channel>,
+    next_worktree_id: u64,
+}
+
+struct ConnectionState {
+    user_id: UserId,
+    worktrees: HashSet<u64>,
+    channels: HashSet<ChannelId>,
+}
+
+pub struct Worktree {
+    pub host_connection_id: ConnectionId,
+    pub collaborator_user_ids: Vec<UserId>,
+    pub root_name: String,
+    pub share: Option<WorktreeShare>,
+}
+
+pub struct WorktreeShare {
+    pub guest_connection_ids: HashMap<ConnectionId, ReplicaId>,
+    pub active_replica_ids: HashSet<ReplicaId>,
+    pub entries: HashMap<u64, proto::Entry>,
+}
+
+#[derive(Default)]
+pub struct Channel {
+    pub connection_ids: HashSet<ConnectionId>,
+}
+
+pub type ReplicaId = u16;
+
+#[derive(Default)]
+pub struct RemovedConnectionState {
+    pub hosted_worktrees: HashMap<u64, Worktree>,
+    pub guest_worktree_ids: HashMap<u64, Vec<ConnectionId>>,
+    pub collaborator_ids: HashSet<UserId>,
+}
+
+pub struct JoinedWorktree<'a> {
+    pub replica_id: ReplicaId,
+    pub worktree: &'a Worktree,
+}
+
+pub struct UnsharedWorktree {
+    pub connection_ids: Vec<ConnectionId>,
+    pub collaborator_ids: Vec<UserId>,
+}
+
+pub struct LeftWorktree {
+    pub connection_ids: Vec<ConnectionId>,
+    pub collaborator_ids: Vec<UserId>,
+}
+
+impl Store {
+    pub fn add_connection(&mut self, connection_id: ConnectionId, user_id: UserId) {
+        self.connections.insert(
+            connection_id,
+            ConnectionState {
+                user_id,
+                worktrees: Default::default(),
+                channels: Default::default(),
+            },
+        );
+        self.connections_by_user_id
+            .entry(user_id)
+            .or_default()
+            .insert(connection_id);
+    }
+
+    pub fn remove_connection(
+        &mut self,
+        connection_id: ConnectionId,
+    ) -> tide::Result<RemovedConnectionState> {
+        let connection = if let Some(connection) = self.connections.remove(&connection_id) {
+            connection
+        } else {
+            return Err(anyhow!("no such connection"))?;
+        };
+
+        for channel_id in &connection.channels {
+            if let Some(channel) = self.channels.get_mut(&channel_id) {
+                channel.connection_ids.remove(&connection_id);
+            }
+        }
+
+        let user_connections = self
+            .connections_by_user_id
+            .get_mut(&connection.user_id)
+            .unwrap();
+        user_connections.remove(&connection_id);
+        if user_connections.is_empty() {
+            self.connections_by_user_id.remove(&connection.user_id);
+        }
+
+        let mut result = RemovedConnectionState::default();
+        for worktree_id in connection.worktrees.clone() {
+            if let Ok(worktree) = self.remove_worktree(worktree_id, connection_id) {
+                result
+                    .collaborator_ids
+                    .extend(worktree.collaborator_user_ids.iter().copied());
+                result.hosted_worktrees.insert(worktree_id, worktree);
+            } else if let Some(worktree) = self.leave_worktree(connection_id, worktree_id) {
+                result
+                    .guest_worktree_ids
+                    .insert(worktree_id, worktree.connection_ids);
+                result.collaborator_ids.extend(worktree.collaborator_ids);
+            }
+        }
+
+        #[cfg(test)]
+        self.check_invariants();
+
+        Ok(result)
+    }
+
+    #[cfg(test)]
+    pub fn channel(&self, id: ChannelId) -> Option<&Channel> {
+        self.channels.get(&id)
+    }
+
+    pub fn join_channel(&mut self, connection_id: ConnectionId, channel_id: ChannelId) {
+        if let Some(connection) = self.connections.get_mut(&connection_id) {
+            connection.channels.insert(channel_id);
+            self.channels
+                .entry(channel_id)
+                .or_default()
+                .connection_ids
+                .insert(connection_id);
+        }
+    }
+
+    pub fn leave_channel(&mut self, connection_id: ConnectionId, channel_id: ChannelId) {
+        if let Some(connection) = self.connections.get_mut(&connection_id) {
+            connection.channels.remove(&channel_id);
+            if let hash_map::Entry::Occupied(mut entry) = self.channels.entry(channel_id) {
+                entry.get_mut().connection_ids.remove(&connection_id);
+                if entry.get_mut().connection_ids.is_empty() {
+                    entry.remove();
+                }
+            }
+        }
+    }
+
+    pub fn user_id_for_connection(&self, connection_id: ConnectionId) -> tide::Result<UserId> {
+        Ok(self
+            .connections
+            .get(&connection_id)
+            .ok_or_else(|| anyhow!("unknown connection"))?
+            .user_id)
+    }
+
+    pub fn connection_ids_for_user<'a>(
+        &'a self,
+        user_id: UserId,
+    ) -> impl 'a + Iterator<Item = ConnectionId> {
+        self.connections_by_user_id
+            .get(&user_id)
+            .into_iter()
+            .flatten()
+            .copied()
+    }
+
+    pub fn collaborators_for_user(&self, user_id: UserId) -> Vec<proto::Collaborator> {
+        let mut collaborators = HashMap::new();
+        for worktree_id in self
+            .visible_worktrees_by_user_id
+            .get(&user_id)
+            .unwrap_or(&HashSet::new())
+        {
+            let worktree = &self.worktrees[worktree_id];
+
+            let mut guests = HashSet::new();
+            if let Ok(share) = worktree.share() {
+                for guest_connection_id in share.guest_connection_ids.keys() {
+                    if let Ok(user_id) = self.user_id_for_connection(*guest_connection_id) {
+                        guests.insert(user_id.to_proto());
+                    }
+                }
+            }
+
+            if let Ok(host_user_id) = self.user_id_for_connection(worktree.host_connection_id) {
+                collaborators
+                    .entry(host_user_id)
+                    .or_insert_with(|| proto::Collaborator {
+                        user_id: host_user_id.to_proto(),
+                        worktrees: Vec::new(),
+                    })
+                    .worktrees
+                    .push(proto::WorktreeMetadata {
+                        id: *worktree_id,
+                        root_name: worktree.root_name.clone(),
+                        is_shared: worktree.share.is_some(),
+                        guests: guests.into_iter().collect(),
+                    });
+            }
+        }
+
+        collaborators.into_values().collect()
+    }
+
+    pub fn add_worktree(&mut self, worktree: Worktree) -> u64 {
+        let worktree_id = self.next_worktree_id;
+        for collaborator_user_id in &worktree.collaborator_user_ids {
+            self.visible_worktrees_by_user_id
+                .entry(*collaborator_user_id)
+                .or_default()
+                .insert(worktree_id);
+        }
+        self.next_worktree_id += 1;
+        if let Some(connection) = self.connections.get_mut(&worktree.host_connection_id) {
+            connection.worktrees.insert(worktree_id);
+        }
+        self.worktrees.insert(worktree_id, worktree);
+
+        #[cfg(test)]
+        self.check_invariants();
+
+        worktree_id
+    }
+
+    pub fn remove_worktree(
+        &mut self,
+        worktree_id: u64,
+        acting_connection_id: ConnectionId,
+    ) -> tide::Result<Worktree> {
+        let worktree = if let hash_map::Entry::Occupied(e) = self.worktrees.entry(worktree_id) {
+            if e.get().host_connection_id != acting_connection_id {
+                Err(anyhow!("not your worktree"))?;
+            }
+            e.remove()
+        } else {
+            return Err(anyhow!("no such worktree"))?;
+        };
+
+        if let Some(connection) = self.connections.get_mut(&worktree.host_connection_id) {
+            connection.worktrees.remove(&worktree_id);
+        }
+
+        if let Some(share) = &worktree.share {
+            for connection_id in share.guest_connection_ids.keys() {
+                if let Some(connection) = self.connections.get_mut(connection_id) {
+                    connection.worktrees.remove(&worktree_id);
+                }
+            }
+        }
+
+        for collaborator_user_id in &worktree.collaborator_user_ids {
+            if let Some(visible_worktrees) = self
+                .visible_worktrees_by_user_id
+                .get_mut(&collaborator_user_id)
+            {
+                visible_worktrees.remove(&worktree_id);
+            }
+        }
+
+        #[cfg(test)]
+        self.check_invariants();
+
+        Ok(worktree)
+    }
+
+    pub fn share_worktree(
+        &mut self,
+        worktree_id: u64,
+        connection_id: ConnectionId,
+        entries: HashMap<u64, proto::Entry>,
+    ) -> Option<Vec<UserId>> {
+        if let Some(worktree) = self.worktrees.get_mut(&worktree_id) {
+            if worktree.host_connection_id == connection_id {
+                worktree.share = Some(WorktreeShare {
+                    guest_connection_ids: Default::default(),
+                    active_replica_ids: Default::default(),
+                    entries,
+                });
+                return Some(worktree.collaborator_user_ids.clone());
+            }
+        }
+        None
+    }
+
+    pub fn unshare_worktree(
+        &mut self,
+        worktree_id: u64,
+        acting_connection_id: ConnectionId,
+    ) -> tide::Result<UnsharedWorktree> {
+        let worktree = if let Some(worktree) = self.worktrees.get_mut(&worktree_id) {
+            worktree
+        } else {
+            return Err(anyhow!("no such worktree"))?;
+        };
+
+        if worktree.host_connection_id != acting_connection_id {
+            return Err(anyhow!("not your worktree"))?;
+        }
+
+        let connection_ids = worktree.connection_ids();
+        let collaborator_ids = worktree.collaborator_user_ids.clone();
+        if let Some(share) = worktree.share.take() {
+            for connection_id in share.guest_connection_ids.into_keys() {
+                if let Some(connection) = self.connections.get_mut(&connection_id) {
+                    connection.worktrees.remove(&worktree_id);
+                }
+            }
+
+            #[cfg(test)]
+            self.check_invariants();
+
+            Ok(UnsharedWorktree {
+                connection_ids,
+                collaborator_ids,
+            })
+        } else {
+            Err(anyhow!("worktree is not shared"))?
+        }
+    }
+
+    pub fn join_worktree(
+        &mut self,
+        connection_id: ConnectionId,
+        user_id: UserId,
+        worktree_id: u64,
+    ) -> tide::Result<JoinedWorktree> {
+        let connection = self
+            .connections
+            .get_mut(&connection_id)
+            .ok_or_else(|| anyhow!("no such connection"))?;
+        let worktree = self
+            .worktrees
+            .get_mut(&worktree_id)
+            .and_then(|worktree| {
+                if worktree.collaborator_user_ids.contains(&user_id) {
+                    Some(worktree)
+                } else {
+                    None
+                }
+            })
+            .ok_or_else(|| anyhow!("no such worktree"))?;
+
+        let share = worktree.share_mut()?;
+        connection.worktrees.insert(worktree_id);
+
+        let mut replica_id = 1;
+        while share.active_replica_ids.contains(&replica_id) {
+            replica_id += 1;
+        }
+        share.active_replica_ids.insert(replica_id);
+        share.guest_connection_ids.insert(connection_id, replica_id);
+
+        #[cfg(test)]
+        self.check_invariants();
+
+        Ok(JoinedWorktree {
+            replica_id,
+            worktree: &self.worktrees[&worktree_id],
+        })
+    }
+
+    pub fn leave_worktree(
+        &mut self,
+        connection_id: ConnectionId,
+        worktree_id: u64,
+    ) -> Option<LeftWorktree> {
+        let worktree = self.worktrees.get_mut(&worktree_id)?;
+        let share = worktree.share.as_mut()?;
+        let replica_id = share.guest_connection_ids.remove(&connection_id)?;
+        share.active_replica_ids.remove(&replica_id);
+
+        if let Some(connection) = self.connections.get_mut(&connection_id) {
+            connection.worktrees.remove(&worktree_id);
+        }
+
+        let connection_ids = worktree.connection_ids();
+        let collaborator_ids = worktree.collaborator_user_ids.clone();
+
+        #[cfg(test)]
+        self.check_invariants();
+
+        Some(LeftWorktree {
+            connection_ids,
+            collaborator_ids,
+        })
+    }
+
+    pub fn update_worktree(
+        &mut self,
+        connection_id: ConnectionId,
+        worktree_id: u64,
+        removed_entries: &[u64],
+        updated_entries: &[proto::Entry],
+    ) -> tide::Result<Vec<ConnectionId>> {
+        let worktree = self.write_worktree(worktree_id, connection_id)?;
+        let share = worktree.share_mut()?;
+        for entry_id in removed_entries {
+            share.entries.remove(&entry_id);
+        }
+        for entry in updated_entries {
+            share.entries.insert(entry.id, entry.clone());
+        }
+        Ok(worktree.connection_ids())
+    }
+
+    pub fn worktree_host_connection_id(
+        &self,
+        connection_id: ConnectionId,
+        worktree_id: u64,
+    ) -> tide::Result<ConnectionId> {
+        Ok(self
+            .read_worktree(worktree_id, connection_id)?
+            .host_connection_id)
+    }
+
+    pub fn worktree_guest_connection_ids(
+        &self,
+        connection_id: ConnectionId,
+        worktree_id: u64,
+    ) -> tide::Result<Vec<ConnectionId>> {
+        Ok(self
+            .read_worktree(worktree_id, connection_id)?
+            .share()?
+            .guest_connection_ids
+            .keys()
+            .copied()
+            .collect())
+    }
+
+    pub fn worktree_connection_ids(
+        &self,
+        connection_id: ConnectionId,
+        worktree_id: u64,
+    ) -> tide::Result<Vec<ConnectionId>> {
+        Ok(self
+            .read_worktree(worktree_id, connection_id)?
+            .connection_ids())
+    }
+
+    pub fn channel_connection_ids(&self, channel_id: ChannelId) -> Option<Vec<ConnectionId>> {
+        Some(self.channels.get(&channel_id)?.connection_ids())
+    }
+
+    fn read_worktree(
+        &self,
+        worktree_id: u64,
+        connection_id: ConnectionId,
+    ) -> tide::Result<&Worktree> {
+        let worktree = self
+            .worktrees
+            .get(&worktree_id)
+            .ok_or_else(|| anyhow!("worktree not found"))?;
+
+        if worktree.host_connection_id == connection_id
+            || worktree
+                .share()?
+                .guest_connection_ids
+                .contains_key(&connection_id)
+        {
+            Ok(worktree)
+        } else {
+            Err(anyhow!(
+                "{} is not a member of worktree {}",
+                connection_id,
+                worktree_id
+            ))?
+        }
+    }
+
+    fn write_worktree(
+        &mut self,
+        worktree_id: u64,
+        connection_id: ConnectionId,
+    ) -> tide::Result<&mut Worktree> {
+        let worktree = self
+            .worktrees
+            .get_mut(&worktree_id)
+            .ok_or_else(|| anyhow!("worktree not found"))?;
+
+        if worktree.host_connection_id == connection_id
+            || worktree.share.as_ref().map_or(false, |share| {
+                share.guest_connection_ids.contains_key(&connection_id)
+            })
+        {
+            Ok(worktree)
+        } else {
+            Err(anyhow!(
+                "{} is not a member of worktree {}",
+                connection_id,
+                worktree_id
+            ))?
+        }
+    }
+
+    #[cfg(test)]
+    fn check_invariants(&self) {
+        for (connection_id, connection) in &self.connections {
+            for worktree_id in &connection.worktrees {
+                let worktree = &self.worktrees.get(&worktree_id).unwrap();
+                if worktree.host_connection_id != *connection_id {
+                    assert!(worktree
+                        .share()
+                        .unwrap()
+                        .guest_connection_ids
+                        .contains_key(connection_id));
+                }
+            }
+            for channel_id in &connection.channels {
+                let channel = self.channels.get(channel_id).unwrap();
+                assert!(channel.connection_ids.contains(connection_id));
+            }
+            assert!(self
+                .connections_by_user_id
+                .get(&connection.user_id)
+                .unwrap()
+                .contains(connection_id));
+        }
+
+        for (user_id, connection_ids) in &self.connections_by_user_id {
+            for connection_id in connection_ids {
+                assert_eq!(
+                    self.connections.get(connection_id).unwrap().user_id,
+                    *user_id
+                );
+            }
+        }
+
+        for (worktree_id, worktree) in &self.worktrees {
+            let host_connection = self.connections.get(&worktree.host_connection_id).unwrap();
+            assert!(host_connection.worktrees.contains(worktree_id));
+
+            for collaborator_id in &worktree.collaborator_user_ids {
+                let visible_worktree_ids = self
+                    .visible_worktrees_by_user_id
+                    .get(collaborator_id)
+                    .unwrap();
+                assert!(visible_worktree_ids.contains(worktree_id));
+            }
+
+            if let Some(share) = &worktree.share {
+                for guest_connection_id in share.guest_connection_ids.keys() {
+                    let guest_connection = self.connections.get(guest_connection_id).unwrap();
+                    assert!(guest_connection.worktrees.contains(worktree_id));
+                }
+                assert_eq!(
+                    share.active_replica_ids.len(),
+                    share.guest_connection_ids.len(),
+                );
+                assert_eq!(
+                    share.active_replica_ids,
+                    share
+                        .guest_connection_ids
+                        .values()
+                        .copied()
+                        .collect::<HashSet<_>>(),
+                );
+            }
+        }
+
+        for (user_id, visible_worktree_ids) in &self.visible_worktrees_by_user_id {
+            for worktree_id in visible_worktree_ids {
+                let worktree = self.worktrees.get(worktree_id).unwrap();
+                assert!(worktree.collaborator_user_ids.contains(user_id));
+            }
+        }
+
+        for (channel_id, channel) in &self.channels {
+            for connection_id in &channel.connection_ids {
+                let connection = self.connections.get(connection_id).unwrap();
+                assert!(connection.channels.contains(channel_id));
+            }
+        }
+    }
+}
+
+impl Worktree {
+    pub fn connection_ids(&self) -> Vec<ConnectionId> {
+        if let Some(share) = &self.share {
+            share
+                .guest_connection_ids
+                .keys()
+                .copied()
+                .chain(Some(self.host_connection_id))
+                .collect()
+        } else {
+            vec![self.host_connection_id]
+        }
+    }
+
+    pub fn share(&self) -> tide::Result<&WorktreeShare> {
+        Ok(self
+            .share
+            .as_ref()
+            .ok_or_else(|| anyhow!("worktree is not shared"))?)
+    }
+
+    fn share_mut(&mut self) -> tide::Result<&mut WorktreeShare> {
+        Ok(self
+            .share
+            .as_mut()
+            .ok_or_else(|| anyhow!("worktree is not shared"))?)
+    }
+}
+
+impl Channel {
+    fn connection_ids(&self) -> Vec<ConnectionId> {
+        self.connection_ids.iter().copied().collect()
+    }
+}