1mod store;
   2
   3use crate::{
   4    auth,
   5    db::{self, ChannelId, MessageId, ProjectId, User, UserId},
   6    AppState, Result,
   7};
   8use anyhow::anyhow;
   9use async_tungstenite::tungstenite::{
  10    protocol::CloseFrame as TungsteniteCloseFrame, Message as TungsteniteMessage,
  11};
  12use axum::{
  13    body::Body,
  14    extract::{
  15        ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage},
  16        ConnectInfo, WebSocketUpgrade,
  17    },
  18    headers::{Header, HeaderName},
  19    http::StatusCode,
  20    middleware,
  21    response::IntoResponse,
  22    routing::get,
  23    Extension, Router, TypedHeader,
  24};
  25use collections::{HashMap, HashSet};
  26use futures::{
  27    channel::{mpsc, oneshot},
  28    future::{self, BoxFuture},
  29    stream::FuturesUnordered,
  30    FutureExt, SinkExt, StreamExt, TryStreamExt,
  31};
  32use lazy_static::lazy_static;
  33use prometheus::{register_int_gauge, IntGauge};
  34use rpc::{
  35    proto::{self, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, RequestMessage},
  36    Connection, ConnectionId, Peer, Receipt, TypedEnvelope,
  37};
  38use serde::{Serialize, Serializer};
  39use std::{
  40    any::TypeId,
  41    future::Future,
  42    marker::PhantomData,
  43    net::SocketAddr,
  44    ops::{Deref, DerefMut},
  45    os::unix::prelude::OsStrExt,
  46    rc::Rc,
  47    sync::{
  48        atomic::{AtomicBool, Ordering::SeqCst},
  49        Arc,
  50    },
  51    time::Duration,
  52};
  53use time::OffsetDateTime;
  54use tokio::{
  55    sync::{Mutex, MutexGuard},
  56    time::Sleep,
  57};
  58use tower::ServiceBuilder;
  59use tracing::{info_span, instrument, Instrument};
  60
  61pub use store::{Store, Worktree};
  62
  63lazy_static! {
  64    static ref METRIC_CONNECTIONS: IntGauge =
  65        register_int_gauge!("connections", "number of connections").unwrap();
  66    static ref METRIC_REGISTERED_PROJECTS: IntGauge =
  67        register_int_gauge!("registered_projects", "number of registered projects").unwrap();
  68    static ref METRIC_ACTIVE_PROJECTS: IntGauge =
  69        register_int_gauge!("active_projects", "number of active projects").unwrap();
  70    static ref METRIC_SHARED_PROJECTS: IntGauge = register_int_gauge!(
  71        "shared_projects",
  72        "number of open projects with one or more guests"
  73    )
  74    .unwrap();
  75}
  76
  77type MessageHandler =
  78    Box<dyn Send + Sync + Fn(Arc<Server>, Box<dyn AnyTypedEnvelope>) -> BoxFuture<'static, ()>>;
  79
  80struct Response<R> {
  81    server: Arc<Server>,
  82    receipt: Receipt<R>,
  83    responded: Arc<AtomicBool>,
  84}
  85
  86impl<R: RequestMessage> Response<R> {
  87    fn send(self, payload: R::Response) -> Result<()> {
  88        self.responded.store(true, SeqCst);
  89        self.server.peer.respond(self.receipt, payload)?;
  90        Ok(())
  91    }
  92}
  93
  94pub struct Server {
  95    peer: Arc<Peer>,
  96    pub(crate) store: Mutex<Store>,
  97    app_state: Arc<AppState>,
  98    handlers: HashMap<TypeId, MessageHandler>,
  99    notifications: Option<mpsc::UnboundedSender<()>>,
 100}
 101
 102pub trait Executor: Send + Clone {
 103    type Sleep: Send + Future;
 104    fn spawn_detached<F: 'static + Send + Future<Output = ()>>(&self, future: F);
 105    fn sleep(&self, duration: Duration) -> Self::Sleep;
 106}
 107
 108#[derive(Clone)]
 109pub struct RealExecutor;
 110
 111const MESSAGE_COUNT_PER_PAGE: usize = 100;
 112const MAX_MESSAGE_LEN: usize = 1024;
 113
 114pub(crate) struct StoreGuard<'a> {
 115    guard: MutexGuard<'a, Store>,
 116    _not_send: PhantomData<Rc<()>>,
 117}
 118
 119#[derive(Serialize)]
 120pub struct ServerSnapshot<'a> {
 121    peer: &'a Peer,
 122    #[serde(serialize_with = "serialize_deref")]
 123    store: StoreGuard<'a>,
 124}
 125
 126pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
 127where
 128    S: Serializer,
 129    T: Deref<Target = U>,
 130    U: Serialize,
 131{
 132    Serialize::serialize(value.deref(), serializer)
 133}
 134
 135impl Server {
 136    pub fn new(
 137        app_state: Arc<AppState>,
 138        notifications: Option<mpsc::UnboundedSender<()>>,
 139    ) -> Arc<Self> {
 140        let mut server = Self {
 141            peer: Peer::new(),
 142            app_state,
 143            store: Default::default(),
 144            handlers: Default::default(),
 145            notifications,
 146        };
 147
 148        server
 149            .add_request_handler(Server::ping)
 150            .add_request_handler(Server::create_room)
 151            .add_request_handler(Server::join_room)
 152            .add_message_handler(Server::leave_room)
 153            .add_request_handler(Server::call)
 154            .add_request_handler(Server::cancel_call)
 155            .add_message_handler(Server::decline_call)
 156            .add_request_handler(Server::update_participant_location)
 157            .add_request_handler(Server::share_project)
 158            .add_message_handler(Server::unshare_project)
 159            .add_request_handler(Server::join_project)
 160            .add_message_handler(Server::leave_project)
 161            .add_message_handler(Server::update_project)
 162            .add_message_handler(Server::register_project_activity)
 163            .add_request_handler(Server::update_worktree)
 164            .add_message_handler(Server::update_worktree_extensions)
 165            .add_message_handler(Server::start_language_server)
 166            .add_message_handler(Server::update_language_server)
 167            .add_message_handler(Server::update_diagnostic_summary)
 168            .add_request_handler(Server::forward_project_request::<proto::GetHover>)
 169            .add_request_handler(Server::forward_project_request::<proto::GetDefinition>)
 170            .add_request_handler(Server::forward_project_request::<proto::GetTypeDefinition>)
 171            .add_request_handler(Server::forward_project_request::<proto::GetReferences>)
 172            .add_request_handler(Server::forward_project_request::<proto::SearchProject>)
 173            .add_request_handler(Server::forward_project_request::<proto::GetDocumentHighlights>)
 174            .add_request_handler(Server::forward_project_request::<proto::GetProjectSymbols>)
 175            .add_request_handler(Server::forward_project_request::<proto::OpenBufferForSymbol>)
 176            .add_request_handler(Server::forward_project_request::<proto::OpenBufferById>)
 177            .add_request_handler(Server::forward_project_request::<proto::OpenBufferByPath>)
 178            .add_request_handler(Server::forward_project_request::<proto::GetCompletions>)
 179            .add_request_handler(
 180                Server::forward_project_request::<proto::ApplyCompletionAdditionalEdits>,
 181            )
 182            .add_request_handler(Server::forward_project_request::<proto::GetCodeActions>)
 183            .add_request_handler(Server::forward_project_request::<proto::ApplyCodeAction>)
 184            .add_request_handler(Server::forward_project_request::<proto::PrepareRename>)
 185            .add_request_handler(Server::forward_project_request::<proto::PerformRename>)
 186            .add_request_handler(Server::forward_project_request::<proto::ReloadBuffers>)
 187            .add_request_handler(Server::forward_project_request::<proto::FormatBuffers>)
 188            .add_request_handler(Server::forward_project_request::<proto::CreateProjectEntry>)
 189            .add_request_handler(Server::forward_project_request::<proto::RenameProjectEntry>)
 190            .add_request_handler(Server::forward_project_request::<proto::CopyProjectEntry>)
 191            .add_request_handler(Server::forward_project_request::<proto::DeleteProjectEntry>)
 192            .add_message_handler(Server::create_buffer_for_peer)
 193            .add_request_handler(Server::update_buffer)
 194            .add_message_handler(Server::update_buffer_file)
 195            .add_message_handler(Server::buffer_reloaded)
 196            .add_message_handler(Server::buffer_saved)
 197            .add_request_handler(Server::save_buffer)
 198            .add_request_handler(Server::get_channels)
 199            .add_request_handler(Server::get_users)
 200            .add_request_handler(Server::fuzzy_search_users)
 201            .add_request_handler(Server::request_contact)
 202            .add_request_handler(Server::remove_contact)
 203            .add_request_handler(Server::respond_to_contact_request)
 204            .add_request_handler(Server::join_channel)
 205            .add_message_handler(Server::leave_channel)
 206            .add_request_handler(Server::send_channel_message)
 207            .add_request_handler(Server::follow)
 208            .add_message_handler(Server::unfollow)
 209            .add_message_handler(Server::update_followers)
 210            .add_request_handler(Server::get_channel_messages)
 211            .add_message_handler(Server::update_diff_base)
 212            .add_request_handler(Server::get_private_user_info);
 213
 214        Arc::new(server)
 215    }
 216
 217    fn add_message_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
 218    where
 219        F: 'static + Send + Sync + Fn(Arc<Self>, TypedEnvelope<M>) -> Fut,
 220        Fut: 'static + Send + Future<Output = Result<()>>,
 221        M: EnvelopedMessage,
 222    {
 223        let prev_handler = self.handlers.insert(
 224            TypeId::of::<M>(),
 225            Box::new(move |server, envelope| {
 226                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
 227                let span = info_span!(
 228                    "handle message",
 229                    payload_type = envelope.payload_type_name()
 230                );
 231                span.in_scope(|| {
 232                    tracing::info!(
 233                        payload_type = envelope.payload_type_name(),
 234                        "message received"
 235                    );
 236                });
 237                let future = (handler)(server, *envelope);
 238                async move {
 239                    if let Err(error) = future.await {
 240                        tracing::error!(%error, "error handling message");
 241                    }
 242                }
 243                .instrument(span)
 244                .boxed()
 245            }),
 246        );
 247        if prev_handler.is_some() {
 248            panic!("registered a handler for the same message twice");
 249        }
 250        self
 251    }
 252
 253    /// Handle a request while holding a lock to the store. This is useful when we're registering
 254    /// a connection but we want to respond on the connection before anybody else can send on it.
 255    fn add_request_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
 256    where
 257        F: 'static + Send + Sync + Fn(Arc<Self>, TypedEnvelope<M>, Response<M>) -> Fut,
 258        Fut: Send + Future<Output = Result<()>>,
 259        M: RequestMessage,
 260    {
 261        let handler = Arc::new(handler);
 262        self.add_message_handler(move |server, envelope| {
 263            let receipt = envelope.receipt();
 264            let handler = handler.clone();
 265            async move {
 266                let responded = Arc::new(AtomicBool::default());
 267                let response = Response {
 268                    server: server.clone(),
 269                    responded: responded.clone(),
 270                    receipt: envelope.receipt(),
 271                };
 272                match (handler)(server.clone(), envelope, response).await {
 273                    Ok(()) => {
 274                        if responded.load(std::sync::atomic::Ordering::SeqCst) {
 275                            Ok(())
 276                        } else {
 277                            Err(anyhow!("handler did not send a response"))?
 278                        }
 279                    }
 280                    Err(error) => {
 281                        server.peer.respond_with_error(
 282                            receipt,
 283                            proto::Error {
 284                                message: error.to_string(),
 285                            },
 286                        )?;
 287                        Err(error)
 288                    }
 289                }
 290            }
 291        })
 292    }
 293
 294    /// Start a long lived task that records which users are active in which projects.
 295    pub fn start_recording_project_activity<E: 'static + Executor>(
 296        self: &Arc<Self>,
 297        interval: Duration,
 298        executor: E,
 299    ) {
 300        executor.spawn_detached({
 301            let this = Arc::downgrade(self);
 302            let executor = executor.clone();
 303            async move {
 304                let mut period_start = OffsetDateTime::now_utc();
 305                let mut active_projects = Vec::<(UserId, ProjectId)>::new();
 306                loop {
 307                    let sleep = executor.sleep(interval);
 308                    sleep.await;
 309                    let this = if let Some(this) = this.upgrade() {
 310                        this
 311                    } else {
 312                        break;
 313                    };
 314
 315                    active_projects.clear();
 316                    active_projects.extend(this.store().await.projects().flat_map(
 317                        |(project_id, project)| {
 318                            project.guests.values().chain([&project.host]).filter_map(
 319                                |collaborator| {
 320                                    if !collaborator.admin
 321                                        && collaborator
 322                                            .last_activity
 323                                            .map_or(false, |activity| activity > period_start)
 324                                    {
 325                                        Some((collaborator.user_id, *project_id))
 326                                    } else {
 327                                        None
 328                                    }
 329                                },
 330                            )
 331                        },
 332                    ));
 333
 334                    let period_end = OffsetDateTime::now_utc();
 335                    this.app_state
 336                        .db
 337                        .record_user_activity(period_start..period_end, &active_projects)
 338                        .await
 339                        .trace_err();
 340                    period_start = period_end;
 341                }
 342            }
 343        });
 344    }
 345
 346    pub fn handle_connection<E: Executor>(
 347        self: &Arc<Self>,
 348        connection: Connection,
 349        address: String,
 350        user: User,
 351        mut send_connection_id: Option<oneshot::Sender<ConnectionId>>,
 352        executor: E,
 353    ) -> impl Future<Output = Result<()>> {
 354        let mut this = self.clone();
 355        let user_id = user.id;
 356        let login = user.github_login;
 357        let span = info_span!("handle connection", %user_id, %login, %address);
 358        async move {
 359            let (connection_id, handle_io, mut incoming_rx) = this
 360                .peer
 361                .add_connection(connection, {
 362                    let executor = executor.clone();
 363                    move |duration| {
 364                        let timer = executor.sleep(duration);
 365                        async move {
 366                            timer.await;
 367                        }
 368                    }
 369                });
 370
 371            tracing::info!(%user_id, %login, %connection_id, %address, "connection opened");
 372            this.peer.send(connection_id, proto::Hello { peer_id: connection_id.0 })?;
 373            tracing::info!(%user_id, %login, %connection_id, %address, "sent hello message");
 374
 375            if let Some(send_connection_id) = send_connection_id.take() {
 376                let _ = send_connection_id.send(connection_id);
 377            }
 378
 379            if !user.connected_once {
 380                this.peer.send(connection_id, proto::ShowContacts {})?;
 381                this.app_state.db.set_user_connected_once(user_id, true).await?;
 382            }
 383
 384            let (contacts, invite_code) = future::try_join(
 385                this.app_state.db.get_contacts(user_id),
 386                this.app_state.db.get_invite_code_for_user(user_id)
 387            ).await?;
 388
 389            {
 390                let mut store = this.store().await;
 391                let incoming_call = store.add_connection(connection_id, user_id, user.admin);
 392                if let Some(incoming_call) = incoming_call {
 393                    this.peer.send(connection_id, incoming_call)?;
 394                }
 395
 396                this.peer.send(connection_id, store.build_initial_contacts_update(contacts))?;
 397
 398                if let Some((code, count)) = invite_code {
 399                    this.peer.send(connection_id, proto::UpdateInviteInfo {
 400                        url: format!("{}{}", this.app_state.config.invite_link_prefix, code),
 401                        count,
 402                    })?;
 403                }
 404            }
 405            this.update_user_contacts(user_id).await?;
 406
 407            let handle_io = handle_io.fuse();
 408            futures::pin_mut!(handle_io);
 409
 410            // Handlers for foreground messages are pushed into the following `FuturesUnordered`.
 411            // This prevents deadlocks when e.g., client A performs a request to client B and
 412            // client B performs a request to client A. If both clients stop processing further
 413            // messages until their respective request completes, they won't have a chance to
 414            // respond to the other client's request and cause a deadlock.
 415            //
 416            // This arrangement ensures we will attempt to process earlier messages first, but fall
 417            // back to processing messages arrived later in the spirit of making progress.
 418            let mut foreground_message_handlers = FuturesUnordered::new();
 419            loop {
 420                let next_message = incoming_rx.next().fuse();
 421                futures::pin_mut!(next_message);
 422                futures::select_biased! {
 423                    result = handle_io => {
 424                        if let Err(error) = result {
 425                            tracing::error!(?error, %user_id, %login, %connection_id, %address, "error handling I/O");
 426                        }
 427                        break;
 428                    }
 429                    _ = foreground_message_handlers.next() => {}
 430                    message = next_message => {
 431                        if let Some(message) = message {
 432                            let type_name = message.payload_type_name();
 433                            let span = tracing::info_span!("receive message", %user_id, %login, %connection_id, %address, type_name);
 434                            let span_enter = span.enter();
 435                            if let Some(handler) = this.handlers.get(&message.payload_type_id()) {
 436                                let notifications = this.notifications.clone();
 437                                let is_background = message.is_background();
 438                                let handle_message = (handler)(this.clone(), message);
 439
 440                                drop(span_enter);
 441                                let handle_message = async move {
 442                                    handle_message.await;
 443                                    if let Some(mut notifications) = notifications {
 444                                        let _ = notifications.send(()).await;
 445                                    }
 446                                }.instrument(span);
 447
 448                                if is_background {
 449                                    executor.spawn_detached(handle_message);
 450                                } else {
 451                                    foreground_message_handlers.push(handle_message);
 452                                }
 453                            } else {
 454                                tracing::error!(%user_id, %login, %connection_id, %address, "no message handler");
 455                            }
 456                        } else {
 457                            tracing::info!(%user_id, %login, %connection_id, %address, "connection closed");
 458                            break;
 459                        }
 460                    }
 461                }
 462            }
 463
 464            drop(foreground_message_handlers);
 465            tracing::info!(%user_id, %login, %connection_id, %address, "signing out");
 466            if let Err(error) = this.sign_out(connection_id).await {
 467                tracing::error!(%user_id, %login, %connection_id, %address, ?error, "error signing out");
 468            }
 469
 470            Ok(())
 471        }.instrument(span)
 472    }
 473
 474    #[instrument(skip(self), err)]
 475    async fn sign_out(self: &mut Arc<Self>, connection_id: ConnectionId) -> Result<()> {
 476        self.peer.disconnect(connection_id);
 477
 478        let mut projects_to_unshare = Vec::new();
 479        let mut contacts_to_update = HashSet::default();
 480        {
 481            let mut store = self.store().await;
 482
 483            #[cfg(test)]
 484            let removed_connection = store.remove_connection(connection_id).unwrap();
 485            #[cfg(not(test))]
 486            let removed_connection = store.remove_connection(connection_id)?;
 487
 488            for project in removed_connection.hosted_projects {
 489                projects_to_unshare.push(project.id);
 490                broadcast(connection_id, project.guests.keys().copied(), |conn_id| {
 491                    self.peer.send(
 492                        conn_id,
 493                        proto::UnshareProject {
 494                            project_id: project.id.to_proto(),
 495                        },
 496                    )
 497                });
 498            }
 499
 500            for project in removed_connection.guest_projects {
 501                broadcast(connection_id, project.connection_ids, |conn_id| {
 502                    self.peer.send(
 503                        conn_id,
 504                        proto::RemoveProjectCollaborator {
 505                            project_id: project.id.to_proto(),
 506                            peer_id: connection_id.0,
 507                        },
 508                    )
 509                });
 510            }
 511
 512            for connection_id in removed_connection.canceled_call_connection_ids {
 513                self.peer
 514                    .send(connection_id, proto::CallCanceled {})
 515                    .trace_err();
 516                contacts_to_update.extend(store.user_id_for_connection(connection_id).ok());
 517            }
 518
 519            if let Some(room) = removed_connection
 520                .room_id
 521                .and_then(|room_id| store.room(room_id))
 522            {
 523                self.room_updated(room);
 524            }
 525
 526            contacts_to_update.insert(removed_connection.user_id);
 527        };
 528
 529        for user_id in contacts_to_update {
 530            self.update_user_contacts(user_id).await.trace_err();
 531        }
 532
 533        for project_id in projects_to_unshare {
 534            self.app_state
 535                .db
 536                .unregister_project(project_id)
 537                .await
 538                .trace_err();
 539        }
 540
 541        Ok(())
 542    }
 543
 544    pub async fn invite_code_redeemed(
 545        self: &Arc<Self>,
 546        inviter_id: UserId,
 547        invitee_id: UserId,
 548    ) -> Result<()> {
 549        if let Some(user) = self.app_state.db.get_user_by_id(inviter_id).await? {
 550            if let Some(code) = &user.invite_code {
 551                let store = self.store().await;
 552                let invitee_contact = store.contact_for_user(invitee_id, true);
 553                for connection_id in store.connection_ids_for_user(inviter_id) {
 554                    self.peer.send(
 555                        connection_id,
 556                        proto::UpdateContacts {
 557                            contacts: vec![invitee_contact.clone()],
 558                            ..Default::default()
 559                        },
 560                    )?;
 561                    self.peer.send(
 562                        connection_id,
 563                        proto::UpdateInviteInfo {
 564                            url: format!("{}{}", self.app_state.config.invite_link_prefix, &code),
 565                            count: user.invite_count as u32,
 566                        },
 567                    )?;
 568                }
 569            }
 570        }
 571        Ok(())
 572    }
 573
 574    pub async fn invite_count_updated(self: &Arc<Self>, user_id: UserId) -> Result<()> {
 575        if let Some(user) = self.app_state.db.get_user_by_id(user_id).await? {
 576            if let Some(invite_code) = &user.invite_code {
 577                let store = self.store().await;
 578                for connection_id in store.connection_ids_for_user(user_id) {
 579                    self.peer.send(
 580                        connection_id,
 581                        proto::UpdateInviteInfo {
 582                            url: format!(
 583                                "{}{}",
 584                                self.app_state.config.invite_link_prefix, invite_code
 585                            ),
 586                            count: user.invite_count as u32,
 587                        },
 588                    )?;
 589                }
 590            }
 591        }
 592        Ok(())
 593    }
 594
 595    async fn ping(
 596        self: Arc<Server>,
 597        _: TypedEnvelope<proto::Ping>,
 598        response: Response<proto::Ping>,
 599    ) -> Result<()> {
 600        response.send(proto::Ack {})?;
 601        Ok(())
 602    }
 603
 604    async fn create_room(
 605        self: Arc<Server>,
 606        request: TypedEnvelope<proto::CreateRoom>,
 607        response: Response<proto::CreateRoom>,
 608    ) -> Result<()> {
 609        let user_id;
 610        let room_id;
 611        {
 612            let mut store = self.store().await;
 613            user_id = store.user_id_for_connection(request.sender_id)?;
 614            room_id = store.create_room(request.sender_id)?;
 615        }
 616        response.send(proto::CreateRoomResponse { id: room_id })?;
 617        self.update_user_contacts(user_id).await?;
 618        Ok(())
 619    }
 620
 621    async fn join_room(
 622        self: Arc<Server>,
 623        request: TypedEnvelope<proto::JoinRoom>,
 624        response: Response<proto::JoinRoom>,
 625    ) -> Result<()> {
 626        let user_id;
 627        {
 628            let mut store = self.store().await;
 629            user_id = store.user_id_for_connection(request.sender_id)?;
 630            let (room, recipient_connection_ids) =
 631                store.join_room(request.payload.id, request.sender_id)?;
 632            for recipient_id in recipient_connection_ids {
 633                self.peer
 634                    .send(recipient_id, proto::CallCanceled {})
 635                    .trace_err();
 636            }
 637            response.send(proto::JoinRoomResponse {
 638                room: Some(room.clone()),
 639            })?;
 640            self.room_updated(room);
 641        }
 642        self.update_user_contacts(user_id).await?;
 643        Ok(())
 644    }
 645
 646    async fn leave_room(self: Arc<Server>, message: TypedEnvelope<proto::LeaveRoom>) -> Result<()> {
 647        let mut contacts_to_update = HashSet::default();
 648        {
 649            let mut store = self.store().await;
 650            let user_id = store.user_id_for_connection(message.sender_id)?;
 651            let left_room = store.leave_room(message.payload.id, message.sender_id)?;
 652            contacts_to_update.insert(user_id);
 653
 654            for project in left_room.unshared_projects {
 655                for connection_id in project.connection_ids() {
 656                    self.peer.send(
 657                        connection_id,
 658                        proto::UnshareProject {
 659                            project_id: project.id.to_proto(),
 660                        },
 661                    )?;
 662                }
 663            }
 664
 665            for project in left_room.left_projects {
 666                if project.remove_collaborator {
 667                    for connection_id in project.connection_ids {
 668                        self.peer.send(
 669                            connection_id,
 670                            proto::RemoveProjectCollaborator {
 671                                project_id: project.id.to_proto(),
 672                                peer_id: message.sender_id.0,
 673                            },
 674                        )?;
 675                    }
 676
 677                    self.peer.send(
 678                        message.sender_id,
 679                        proto::UnshareProject {
 680                            project_id: project.id.to_proto(),
 681                        },
 682                    )?;
 683                }
 684            }
 685
 686            if let Some(room) = left_room.room {
 687                self.room_updated(room);
 688            }
 689
 690            for connection_id in left_room.canceled_call_connection_ids {
 691                self.peer
 692                    .send(connection_id, proto::CallCanceled {})
 693                    .trace_err();
 694                contacts_to_update.extend(store.user_id_for_connection(connection_id).ok());
 695            }
 696        }
 697
 698        for user_id in contacts_to_update {
 699            self.update_user_contacts(user_id).await?;
 700        }
 701
 702        Ok(())
 703    }
 704
 705    async fn call(
 706        self: Arc<Server>,
 707        request: TypedEnvelope<proto::Call>,
 708        response: Response<proto::Call>,
 709    ) -> Result<()> {
 710        let caller_user_id = self
 711            .store()
 712            .await
 713            .user_id_for_connection(request.sender_id)?;
 714        let recipient_user_id = UserId::from_proto(request.payload.recipient_user_id);
 715        let initial_project_id = request
 716            .payload
 717            .initial_project_id
 718            .map(ProjectId::from_proto);
 719        if !self
 720            .app_state
 721            .db
 722            .has_contact(caller_user_id, recipient_user_id)
 723            .await?
 724        {
 725            return Err(anyhow!("cannot call a user who isn't a contact"))?;
 726        }
 727
 728        let room_id = request.payload.room_id;
 729        let mut calls = {
 730            let mut store = self.store().await;
 731            let (room, recipient_connection_ids, incoming_call) = store.call(
 732                room_id,
 733                recipient_user_id,
 734                initial_project_id,
 735                request.sender_id,
 736            )?;
 737            self.room_updated(room);
 738            recipient_connection_ids
 739                .into_iter()
 740                .map(|recipient_connection_id| {
 741                    self.peer
 742                        .request(recipient_connection_id, incoming_call.clone())
 743                })
 744                .collect::<FuturesUnordered<_>>()
 745        };
 746        self.update_user_contacts(recipient_user_id).await?;
 747
 748        while let Some(call_response) = calls.next().await {
 749            match call_response.as_ref() {
 750                Ok(_) => {
 751                    response.send(proto::Ack {})?;
 752                    return Ok(());
 753                }
 754                Err(_) => {
 755                    call_response.trace_err();
 756                }
 757            }
 758        }
 759
 760        {
 761            let mut store = self.store().await;
 762            let room = store.call_failed(room_id, recipient_user_id)?;
 763            self.room_updated(&room);
 764        }
 765        self.update_user_contacts(recipient_user_id).await?;
 766
 767        Err(anyhow!("failed to ring call recipient"))?
 768    }
 769
 770    async fn cancel_call(
 771        self: Arc<Server>,
 772        request: TypedEnvelope<proto::CancelCall>,
 773        response: Response<proto::CancelCall>,
 774    ) -> Result<()> {
 775        let recipient_user_id = UserId::from_proto(request.payload.recipient_user_id);
 776        {
 777            let mut store = self.store().await;
 778            let (room, recipient_connection_ids) = store.cancel_call(
 779                request.payload.room_id,
 780                recipient_user_id,
 781                request.sender_id,
 782            )?;
 783            for recipient_id in recipient_connection_ids {
 784                self.peer
 785                    .send(recipient_id, proto::CallCanceled {})
 786                    .trace_err();
 787            }
 788            self.room_updated(room);
 789            response.send(proto::Ack {})?;
 790        }
 791        self.update_user_contacts(recipient_user_id).await?;
 792        Ok(())
 793    }
 794
 795    async fn decline_call(
 796        self: Arc<Server>,
 797        message: TypedEnvelope<proto::DeclineCall>,
 798    ) -> Result<()> {
 799        let recipient_user_id;
 800        {
 801            let mut store = self.store().await;
 802            recipient_user_id = store.user_id_for_connection(message.sender_id)?;
 803            let (room, recipient_connection_ids) =
 804                store.decline_call(message.payload.room_id, message.sender_id)?;
 805            for recipient_id in recipient_connection_ids {
 806                self.peer
 807                    .send(recipient_id, proto::CallCanceled {})
 808                    .trace_err();
 809            }
 810            self.room_updated(room);
 811        }
 812        self.update_user_contacts(recipient_user_id).await?;
 813        Ok(())
 814    }
 815
 816    async fn update_participant_location(
 817        self: Arc<Server>,
 818        request: TypedEnvelope<proto::UpdateParticipantLocation>,
 819        response: Response<proto::UpdateParticipantLocation>,
 820    ) -> Result<()> {
 821        let room_id = request.payload.room_id;
 822        let location = request
 823            .payload
 824            .location
 825            .ok_or_else(|| anyhow!("invalid location"))?;
 826        let mut store = self.store().await;
 827        let room = store.update_participant_location(room_id, location, request.sender_id)?;
 828        self.room_updated(room);
 829        response.send(proto::Ack {})?;
 830        Ok(())
 831    }
 832
 833    fn room_updated(&self, room: &proto::Room) {
 834        for participant in &room.participants {
 835            self.peer
 836                .send(
 837                    ConnectionId(participant.peer_id),
 838                    proto::RoomUpdated {
 839                        room: Some(room.clone()),
 840                    },
 841                )
 842                .trace_err();
 843        }
 844    }
 845
 846    async fn share_project(
 847        self: Arc<Server>,
 848        request: TypedEnvelope<proto::ShareProject>,
 849        response: Response<proto::ShareProject>,
 850    ) -> Result<()> {
 851        let user_id = self
 852            .store()
 853            .await
 854            .user_id_for_connection(request.sender_id)?;
 855        let project_id = self.app_state.db.register_project(user_id).await?;
 856        let mut store = self.store().await;
 857        let room = store.share_project(
 858            request.payload.room_id,
 859            project_id,
 860            request.payload.worktrees,
 861            request.sender_id,
 862        )?;
 863        response.send(proto::ShareProjectResponse {
 864            project_id: project_id.to_proto(),
 865        })?;
 866        self.room_updated(room);
 867
 868        Ok(())
 869    }
 870
 871    async fn unshare_project(
 872        self: Arc<Server>,
 873        message: TypedEnvelope<proto::UnshareProject>,
 874    ) -> Result<()> {
 875        let project_id = ProjectId::from_proto(message.payload.project_id);
 876        let mut store = self.store().await;
 877        let (room, project) = store.unshare_project(project_id, message.sender_id)?;
 878        broadcast(
 879            message.sender_id,
 880            project.guest_connection_ids(),
 881            |conn_id| self.peer.send(conn_id, message.payload.clone()),
 882        );
 883        self.room_updated(room);
 884
 885        Ok(())
 886    }
 887
 888    async fn update_user_contacts(self: &Arc<Server>, user_id: UserId) -> Result<()> {
 889        let contacts = self.app_state.db.get_contacts(user_id).await?;
 890        let store = self.store().await;
 891        let updated_contact = store.contact_for_user(user_id, false);
 892        for contact in contacts {
 893            if let db::Contact::Accepted {
 894                user_id: contact_user_id,
 895                ..
 896            } = contact
 897            {
 898                for contact_conn_id in store.connection_ids_for_user(contact_user_id) {
 899                    self.peer
 900                        .send(
 901                            contact_conn_id,
 902                            proto::UpdateContacts {
 903                                contacts: vec![updated_contact.clone()],
 904                                remove_contacts: Default::default(),
 905                                incoming_requests: Default::default(),
 906                                remove_incoming_requests: Default::default(),
 907                                outgoing_requests: Default::default(),
 908                                remove_outgoing_requests: Default::default(),
 909                            },
 910                        )
 911                        .trace_err();
 912                }
 913            }
 914        }
 915        Ok(())
 916    }
 917
 918    async fn join_project(
 919        self: Arc<Server>,
 920        request: TypedEnvelope<proto::JoinProject>,
 921        response: Response<proto::JoinProject>,
 922    ) -> Result<()> {
 923        let project_id = ProjectId::from_proto(request.payload.project_id);
 924
 925        let host_user_id;
 926        let guest_user_id;
 927        let host_connection_id;
 928        {
 929            let state = self.store().await;
 930            let project = state.project(project_id)?;
 931            host_user_id = project.host.user_id;
 932            host_connection_id = project.host_connection_id;
 933            guest_user_id = state.user_id_for_connection(request.sender_id)?;
 934        };
 935
 936        tracing::info!(%project_id, %host_user_id, %host_connection_id, "join project");
 937
 938        let mut store = self.store().await;
 939        let (project, replica_id) = store.join_project(request.sender_id, project_id)?;
 940        let peer_count = project.guests.len();
 941        let mut collaborators = Vec::with_capacity(peer_count);
 942        collaborators.push(proto::Collaborator {
 943            peer_id: project.host_connection_id.0,
 944            replica_id: 0,
 945            user_id: project.host.user_id.to_proto(),
 946        });
 947        let worktrees = project
 948            .worktrees
 949            .iter()
 950            .map(|(id, worktree)| proto::WorktreeMetadata {
 951                id: *id,
 952                root_name: worktree.root_name.clone(),
 953                visible: worktree.visible,
 954                abs_path: worktree.abs_path.as_os_str().as_bytes().to_vec(),
 955            })
 956            .collect::<Vec<_>>();
 957
 958        // Add all guests other than the requesting user's own connections as collaborators
 959        for (guest_conn_id, guest) in &project.guests {
 960            if request.sender_id != *guest_conn_id {
 961                collaborators.push(proto::Collaborator {
 962                    peer_id: guest_conn_id.0,
 963                    replica_id: guest.replica_id as u32,
 964                    user_id: guest.user_id.to_proto(),
 965                });
 966            }
 967        }
 968
 969        for conn_id in project.connection_ids() {
 970            if conn_id != request.sender_id {
 971                self.peer.send(
 972                    conn_id,
 973                    proto::AddProjectCollaborator {
 974                        project_id: project_id.to_proto(),
 975                        collaborator: Some(proto::Collaborator {
 976                            peer_id: request.sender_id.0,
 977                            replica_id: replica_id as u32,
 978                            user_id: guest_user_id.to_proto(),
 979                        }),
 980                    },
 981                )?;
 982            }
 983        }
 984
 985        // First, we send the metadata associated with each worktree.
 986        response.send(proto::JoinProjectResponse {
 987            worktrees: worktrees.clone(),
 988            replica_id: replica_id as u32,
 989            collaborators: collaborators.clone(),
 990            language_servers: project.language_servers.clone(),
 991        })?;
 992
 993        for (worktree_id, worktree) in &project.worktrees {
 994            #[cfg(any(test, feature = "test-support"))]
 995            const MAX_CHUNK_SIZE: usize = 2;
 996            #[cfg(not(any(test, feature = "test-support")))]
 997            const MAX_CHUNK_SIZE: usize = 256;
 998
 999            // Stream this worktree's entries.
1000            let message = proto::UpdateWorktree {
1001                project_id: project_id.to_proto(),
1002                worktree_id: *worktree_id,
1003                abs_path: worktree.abs_path.as_os_str().as_bytes().to_vec(),
1004                root_name: worktree.root_name.clone(),
1005                updated_entries: worktree.entries.values().cloned().collect(),
1006                removed_entries: Default::default(),
1007                scan_id: worktree.scan_id,
1008                is_last_update: worktree.is_complete,
1009            };
1010            for update in proto::split_worktree_update(message, MAX_CHUNK_SIZE) {
1011                self.peer.send(request.sender_id, update.clone())?;
1012            }
1013
1014            // Stream this worktree's diagnostics.
1015            for summary in worktree.diagnostic_summaries.values() {
1016                self.peer.send(
1017                    request.sender_id,
1018                    proto::UpdateDiagnosticSummary {
1019                        project_id: project_id.to_proto(),
1020                        worktree_id: *worktree_id,
1021                        summary: Some(summary.clone()),
1022                    },
1023                )?;
1024            }
1025        }
1026
1027        for language_server in &project.language_servers {
1028            self.peer.send(
1029                request.sender_id,
1030                proto::UpdateLanguageServer {
1031                    project_id: project_id.to_proto(),
1032                    language_server_id: language_server.id,
1033                    variant: Some(
1034                        proto::update_language_server::Variant::DiskBasedDiagnosticsUpdated(
1035                            proto::LspDiskBasedDiagnosticsUpdated {},
1036                        ),
1037                    ),
1038                },
1039            )?;
1040        }
1041
1042        Ok(())
1043    }
1044
1045    async fn leave_project(
1046        self: Arc<Server>,
1047        request: TypedEnvelope<proto::LeaveProject>,
1048    ) -> Result<()> {
1049        let sender_id = request.sender_id;
1050        let project_id = ProjectId::from_proto(request.payload.project_id);
1051        let project;
1052        {
1053            let mut store = self.store().await;
1054            project = store.leave_project(project_id, sender_id)?;
1055            tracing::info!(
1056                %project_id,
1057                host_user_id = %project.host_user_id,
1058                host_connection_id = %project.host_connection_id,
1059                "leave project"
1060            );
1061
1062            if project.remove_collaborator {
1063                broadcast(sender_id, project.connection_ids, |conn_id| {
1064                    self.peer.send(
1065                        conn_id,
1066                        proto::RemoveProjectCollaborator {
1067                            project_id: project_id.to_proto(),
1068                            peer_id: sender_id.0,
1069                        },
1070                    )
1071                });
1072            }
1073        }
1074
1075        Ok(())
1076    }
1077
1078    async fn update_project(
1079        self: Arc<Server>,
1080        request: TypedEnvelope<proto::UpdateProject>,
1081    ) -> Result<()> {
1082        let project_id = ProjectId::from_proto(request.payload.project_id);
1083        {
1084            let mut state = self.store().await;
1085            let guest_connection_ids = state
1086                .read_project(project_id, request.sender_id)?
1087                .guest_connection_ids();
1088            let room =
1089                state.update_project(project_id, &request.payload.worktrees, request.sender_id)?;
1090            broadcast(request.sender_id, guest_connection_ids, |connection_id| {
1091                self.peer
1092                    .forward_send(request.sender_id, connection_id, request.payload.clone())
1093            });
1094            self.room_updated(room);
1095        };
1096
1097        Ok(())
1098    }
1099
1100    async fn register_project_activity(
1101        self: Arc<Server>,
1102        request: TypedEnvelope<proto::RegisterProjectActivity>,
1103    ) -> Result<()> {
1104        self.store().await.register_project_activity(
1105            ProjectId::from_proto(request.payload.project_id),
1106            request.sender_id,
1107        )?;
1108        Ok(())
1109    }
1110
1111    async fn update_worktree(
1112        self: Arc<Server>,
1113        request: TypedEnvelope<proto::UpdateWorktree>,
1114        response: Response<proto::UpdateWorktree>,
1115    ) -> Result<()> {
1116        let project_id = ProjectId::from_proto(request.payload.project_id);
1117        let worktree_id = request.payload.worktree_id;
1118        let connection_ids = self.store().await.update_worktree(
1119            request.sender_id,
1120            project_id,
1121            worktree_id,
1122            &request.payload.root_name,
1123            &request.payload.removed_entries,
1124            &request.payload.updated_entries,
1125            request.payload.scan_id,
1126            request.payload.is_last_update,
1127        )?;
1128
1129        broadcast(request.sender_id, connection_ids, |connection_id| {
1130            self.peer
1131                .forward_send(request.sender_id, connection_id, request.payload.clone())
1132        });
1133        response.send(proto::Ack {})?;
1134        Ok(())
1135    }
1136
1137    async fn update_worktree_extensions(
1138        self: Arc<Server>,
1139        request: TypedEnvelope<proto::UpdateWorktreeExtensions>,
1140    ) -> Result<()> {
1141        let project_id = ProjectId::from_proto(request.payload.project_id);
1142        let worktree_id = request.payload.worktree_id;
1143        let extensions = request
1144            .payload
1145            .extensions
1146            .into_iter()
1147            .zip(request.payload.counts)
1148            .collect();
1149        self.app_state
1150            .db
1151            .update_worktree_extensions(project_id, worktree_id, extensions)
1152            .await?;
1153        Ok(())
1154    }
1155
1156    async fn update_diagnostic_summary(
1157        self: Arc<Server>,
1158        request: TypedEnvelope<proto::UpdateDiagnosticSummary>,
1159    ) -> Result<()> {
1160        let summary = request
1161            .payload
1162            .summary
1163            .clone()
1164            .ok_or_else(|| anyhow!("invalid summary"))?;
1165        let receiver_ids = self.store().await.update_diagnostic_summary(
1166            ProjectId::from_proto(request.payload.project_id),
1167            request.payload.worktree_id,
1168            request.sender_id,
1169            summary,
1170        )?;
1171
1172        broadcast(request.sender_id, receiver_ids, |connection_id| {
1173            self.peer
1174                .forward_send(request.sender_id, connection_id, request.payload.clone())
1175        });
1176        Ok(())
1177    }
1178
1179    async fn start_language_server(
1180        self: Arc<Server>,
1181        request: TypedEnvelope<proto::StartLanguageServer>,
1182    ) -> Result<()> {
1183        let receiver_ids = self.store().await.start_language_server(
1184            ProjectId::from_proto(request.payload.project_id),
1185            request.sender_id,
1186            request
1187                .payload
1188                .server
1189                .clone()
1190                .ok_or_else(|| anyhow!("invalid language server"))?,
1191        )?;
1192        broadcast(request.sender_id, receiver_ids, |connection_id| {
1193            self.peer
1194                .forward_send(request.sender_id, connection_id, request.payload.clone())
1195        });
1196        Ok(())
1197    }
1198
1199    async fn update_language_server(
1200        self: Arc<Server>,
1201        request: TypedEnvelope<proto::UpdateLanguageServer>,
1202    ) -> Result<()> {
1203        let receiver_ids = self.store().await.project_connection_ids(
1204            ProjectId::from_proto(request.payload.project_id),
1205            request.sender_id,
1206        )?;
1207        broadcast(request.sender_id, receiver_ids, |connection_id| {
1208            self.peer
1209                .forward_send(request.sender_id, connection_id, request.payload.clone())
1210        });
1211        Ok(())
1212    }
1213
1214    async fn forward_project_request<T>(
1215        self: Arc<Server>,
1216        request: TypedEnvelope<T>,
1217        response: Response<T>,
1218    ) -> Result<()>
1219    where
1220        T: EntityMessage + RequestMessage,
1221    {
1222        let project_id = ProjectId::from_proto(request.payload.remote_entity_id());
1223        let host_connection_id = self
1224            .store()
1225            .await
1226            .read_project(project_id, request.sender_id)?
1227            .host_connection_id;
1228        let payload = self
1229            .peer
1230            .forward_request(request.sender_id, host_connection_id, request.payload)
1231            .await?;
1232
1233        // Ensure project still exists by the time we get the response from the host.
1234        self.store()
1235            .await
1236            .read_project(project_id, request.sender_id)?;
1237
1238        response.send(payload)?;
1239        Ok(())
1240    }
1241
1242    async fn save_buffer(
1243        self: Arc<Server>,
1244        request: TypedEnvelope<proto::SaveBuffer>,
1245        response: Response<proto::SaveBuffer>,
1246    ) -> Result<()> {
1247        let project_id = ProjectId::from_proto(request.payload.project_id);
1248        let host = self
1249            .store()
1250            .await
1251            .read_project(project_id, request.sender_id)?
1252            .host_connection_id;
1253        let response_payload = self
1254            .peer
1255            .forward_request(request.sender_id, host, request.payload.clone())
1256            .await?;
1257
1258        let mut guests = self
1259            .store()
1260            .await
1261            .read_project(project_id, request.sender_id)?
1262            .connection_ids();
1263        guests.retain(|guest_connection_id| *guest_connection_id != request.sender_id);
1264        broadcast(host, guests, |conn_id| {
1265            self.peer
1266                .forward_send(host, conn_id, response_payload.clone())
1267        });
1268        response.send(response_payload)?;
1269        Ok(())
1270    }
1271
1272    async fn create_buffer_for_peer(
1273        self: Arc<Server>,
1274        request: TypedEnvelope<proto::CreateBufferForPeer>,
1275    ) -> Result<()> {
1276        self.peer.forward_send(
1277            request.sender_id,
1278            ConnectionId(request.payload.peer_id),
1279            request.payload,
1280        )?;
1281        Ok(())
1282    }
1283
1284    async fn update_buffer(
1285        self: Arc<Server>,
1286        request: TypedEnvelope<proto::UpdateBuffer>,
1287        response: Response<proto::UpdateBuffer>,
1288    ) -> Result<()> {
1289        let project_id = ProjectId::from_proto(request.payload.project_id);
1290        let receiver_ids = {
1291            let mut store = self.store().await;
1292            store.register_project_activity(project_id, request.sender_id)?;
1293            store.project_connection_ids(project_id, request.sender_id)?
1294        };
1295
1296        broadcast(request.sender_id, receiver_ids, |connection_id| {
1297            self.peer
1298                .forward_send(request.sender_id, connection_id, request.payload.clone())
1299        });
1300        response.send(proto::Ack {})?;
1301        Ok(())
1302    }
1303
1304    async fn update_buffer_file(
1305        self: Arc<Server>,
1306        request: TypedEnvelope<proto::UpdateBufferFile>,
1307    ) -> Result<()> {
1308        let receiver_ids = self.store().await.project_connection_ids(
1309            ProjectId::from_proto(request.payload.project_id),
1310            request.sender_id,
1311        )?;
1312        broadcast(request.sender_id, receiver_ids, |connection_id| {
1313            self.peer
1314                .forward_send(request.sender_id, connection_id, request.payload.clone())
1315        });
1316        Ok(())
1317    }
1318
1319    async fn buffer_reloaded(
1320        self: Arc<Server>,
1321        request: TypedEnvelope<proto::BufferReloaded>,
1322    ) -> Result<()> {
1323        let receiver_ids = self.store().await.project_connection_ids(
1324            ProjectId::from_proto(request.payload.project_id),
1325            request.sender_id,
1326        )?;
1327        broadcast(request.sender_id, receiver_ids, |connection_id| {
1328            self.peer
1329                .forward_send(request.sender_id, connection_id, request.payload.clone())
1330        });
1331        Ok(())
1332    }
1333
1334    async fn buffer_saved(
1335        self: Arc<Server>,
1336        request: TypedEnvelope<proto::BufferSaved>,
1337    ) -> Result<()> {
1338        let receiver_ids = self.store().await.project_connection_ids(
1339            ProjectId::from_proto(request.payload.project_id),
1340            request.sender_id,
1341        )?;
1342        broadcast(request.sender_id, receiver_ids, |connection_id| {
1343            self.peer
1344                .forward_send(request.sender_id, connection_id, request.payload.clone())
1345        });
1346        Ok(())
1347    }
1348
1349    async fn follow(
1350        self: Arc<Self>,
1351        request: TypedEnvelope<proto::Follow>,
1352        response: Response<proto::Follow>,
1353    ) -> Result<()> {
1354        let project_id = ProjectId::from_proto(request.payload.project_id);
1355        let leader_id = ConnectionId(request.payload.leader_id);
1356        let follower_id = request.sender_id;
1357        {
1358            let mut store = self.store().await;
1359            if !store
1360                .project_connection_ids(project_id, follower_id)?
1361                .contains(&leader_id)
1362            {
1363                Err(anyhow!("no such peer"))?;
1364            }
1365
1366            store.register_project_activity(project_id, follower_id)?;
1367        }
1368
1369        let mut response_payload = self
1370            .peer
1371            .forward_request(request.sender_id, leader_id, request.payload)
1372            .await?;
1373        response_payload
1374            .views
1375            .retain(|view| view.leader_id != Some(follower_id.0));
1376        response.send(response_payload)?;
1377        Ok(())
1378    }
1379
1380    async fn unfollow(self: Arc<Self>, request: TypedEnvelope<proto::Unfollow>) -> Result<()> {
1381        let project_id = ProjectId::from_proto(request.payload.project_id);
1382        let leader_id = ConnectionId(request.payload.leader_id);
1383        let mut store = self.store().await;
1384        if !store
1385            .project_connection_ids(project_id, request.sender_id)?
1386            .contains(&leader_id)
1387        {
1388            Err(anyhow!("no such peer"))?;
1389        }
1390        store.register_project_activity(project_id, request.sender_id)?;
1391        self.peer
1392            .forward_send(request.sender_id, leader_id, request.payload)?;
1393        Ok(())
1394    }
1395
1396    async fn update_followers(
1397        self: Arc<Self>,
1398        request: TypedEnvelope<proto::UpdateFollowers>,
1399    ) -> Result<()> {
1400        let project_id = ProjectId::from_proto(request.payload.project_id);
1401        let mut store = self.store().await;
1402        store.register_project_activity(project_id, request.sender_id)?;
1403        let connection_ids = store.project_connection_ids(project_id, request.sender_id)?;
1404        let leader_id = request
1405            .payload
1406            .variant
1407            .as_ref()
1408            .and_then(|variant| match variant {
1409                proto::update_followers::Variant::CreateView(payload) => payload.leader_id,
1410                proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
1411                proto::update_followers::Variant::UpdateActiveView(payload) => payload.leader_id,
1412            });
1413        for follower_id in &request.payload.follower_ids {
1414            let follower_id = ConnectionId(*follower_id);
1415            if connection_ids.contains(&follower_id) && Some(follower_id.0) != leader_id {
1416                self.peer
1417                    .forward_send(request.sender_id, follower_id, request.payload.clone())?;
1418            }
1419        }
1420        Ok(())
1421    }
1422
1423    async fn get_channels(
1424        self: Arc<Server>,
1425        request: TypedEnvelope<proto::GetChannels>,
1426        response: Response<proto::GetChannels>,
1427    ) -> Result<()> {
1428        let user_id = self
1429            .store()
1430            .await
1431            .user_id_for_connection(request.sender_id)?;
1432        let channels = self.app_state.db.get_accessible_channels(user_id).await?;
1433        response.send(proto::GetChannelsResponse {
1434            channels: channels
1435                .into_iter()
1436                .map(|chan| proto::Channel {
1437                    id: chan.id.to_proto(),
1438                    name: chan.name,
1439                })
1440                .collect(),
1441        })?;
1442        Ok(())
1443    }
1444
1445    async fn get_users(
1446        self: Arc<Server>,
1447        request: TypedEnvelope<proto::GetUsers>,
1448        response: Response<proto::GetUsers>,
1449    ) -> Result<()> {
1450        let user_ids = request
1451            .payload
1452            .user_ids
1453            .into_iter()
1454            .map(UserId::from_proto)
1455            .collect();
1456        let users = self
1457            .app_state
1458            .db
1459            .get_users_by_ids(user_ids)
1460            .await?
1461            .into_iter()
1462            .map(|user| proto::User {
1463                id: user.id.to_proto(),
1464                avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1465                github_login: user.github_login,
1466            })
1467            .collect();
1468        response.send(proto::UsersResponse { users })?;
1469        Ok(())
1470    }
1471
1472    async fn fuzzy_search_users(
1473        self: Arc<Server>,
1474        request: TypedEnvelope<proto::FuzzySearchUsers>,
1475        response: Response<proto::FuzzySearchUsers>,
1476    ) -> Result<()> {
1477        let user_id = self
1478            .store()
1479            .await
1480            .user_id_for_connection(request.sender_id)?;
1481        let query = request.payload.query;
1482        let db = &self.app_state.db;
1483        let users = match query.len() {
1484            0 => vec![],
1485            1 | 2 => db
1486                .get_user_by_github_account(&query, None)
1487                .await?
1488                .into_iter()
1489                .collect(),
1490            _ => db.fuzzy_search_users(&query, 10).await?,
1491        };
1492        let users = users
1493            .into_iter()
1494            .filter(|user| user.id != user_id)
1495            .map(|user| proto::User {
1496                id: user.id.to_proto(),
1497                avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1498                github_login: user.github_login,
1499            })
1500            .collect();
1501        response.send(proto::UsersResponse { users })?;
1502        Ok(())
1503    }
1504
1505    async fn request_contact(
1506        self: Arc<Server>,
1507        request: TypedEnvelope<proto::RequestContact>,
1508        response: Response<proto::RequestContact>,
1509    ) -> Result<()> {
1510        let requester_id = self
1511            .store()
1512            .await
1513            .user_id_for_connection(request.sender_id)?;
1514        let responder_id = UserId::from_proto(request.payload.responder_id);
1515        if requester_id == responder_id {
1516            return Err(anyhow!("cannot add yourself as a contact"))?;
1517        }
1518
1519        self.app_state
1520            .db
1521            .send_contact_request(requester_id, responder_id)
1522            .await?;
1523
1524        // Update outgoing contact requests of requester
1525        let mut update = proto::UpdateContacts::default();
1526        update.outgoing_requests.push(responder_id.to_proto());
1527        for connection_id in self.store().await.connection_ids_for_user(requester_id) {
1528            self.peer.send(connection_id, update.clone())?;
1529        }
1530
1531        // Update incoming contact requests of responder
1532        let mut update = proto::UpdateContacts::default();
1533        update
1534            .incoming_requests
1535            .push(proto::IncomingContactRequest {
1536                requester_id: requester_id.to_proto(),
1537                should_notify: true,
1538            });
1539        for connection_id in self.store().await.connection_ids_for_user(responder_id) {
1540            self.peer.send(connection_id, update.clone())?;
1541        }
1542
1543        response.send(proto::Ack {})?;
1544        Ok(())
1545    }
1546
1547    async fn respond_to_contact_request(
1548        self: Arc<Server>,
1549        request: TypedEnvelope<proto::RespondToContactRequest>,
1550        response: Response<proto::RespondToContactRequest>,
1551    ) -> Result<()> {
1552        let responder_id = self
1553            .store()
1554            .await
1555            .user_id_for_connection(request.sender_id)?;
1556        let requester_id = UserId::from_proto(request.payload.requester_id);
1557        if request.payload.response == proto::ContactRequestResponse::Dismiss as i32 {
1558            self.app_state
1559                .db
1560                .dismiss_contact_notification(responder_id, requester_id)
1561                .await?;
1562        } else {
1563            let accept = request.payload.response == proto::ContactRequestResponse::Accept as i32;
1564            self.app_state
1565                .db
1566                .respond_to_contact_request(responder_id, requester_id, accept)
1567                .await?;
1568
1569            let store = self.store().await;
1570            // Update responder with new contact
1571            let mut update = proto::UpdateContacts::default();
1572            if accept {
1573                update
1574                    .contacts
1575                    .push(store.contact_for_user(requester_id, false));
1576            }
1577            update
1578                .remove_incoming_requests
1579                .push(requester_id.to_proto());
1580            for connection_id in store.connection_ids_for_user(responder_id) {
1581                self.peer.send(connection_id, update.clone())?;
1582            }
1583
1584            // Update requester with new contact
1585            let mut update = proto::UpdateContacts::default();
1586            if accept {
1587                update
1588                    .contacts
1589                    .push(store.contact_for_user(responder_id, true));
1590            }
1591            update
1592                .remove_outgoing_requests
1593                .push(responder_id.to_proto());
1594            for connection_id in store.connection_ids_for_user(requester_id) {
1595                self.peer.send(connection_id, update.clone())?;
1596            }
1597        }
1598
1599        response.send(proto::Ack {})?;
1600        Ok(())
1601    }
1602
1603    async fn remove_contact(
1604        self: Arc<Server>,
1605        request: TypedEnvelope<proto::RemoveContact>,
1606        response: Response<proto::RemoveContact>,
1607    ) -> Result<()> {
1608        let requester_id = self
1609            .store()
1610            .await
1611            .user_id_for_connection(request.sender_id)?;
1612        let responder_id = UserId::from_proto(request.payload.user_id);
1613        self.app_state
1614            .db
1615            .remove_contact(requester_id, responder_id)
1616            .await?;
1617
1618        // Update outgoing contact requests of requester
1619        let mut update = proto::UpdateContacts::default();
1620        update
1621            .remove_outgoing_requests
1622            .push(responder_id.to_proto());
1623        for connection_id in self.store().await.connection_ids_for_user(requester_id) {
1624            self.peer.send(connection_id, update.clone())?;
1625        }
1626
1627        // Update incoming contact requests of responder
1628        let mut update = proto::UpdateContacts::default();
1629        update
1630            .remove_incoming_requests
1631            .push(requester_id.to_proto());
1632        for connection_id in self.store().await.connection_ids_for_user(responder_id) {
1633            self.peer.send(connection_id, update.clone())?;
1634        }
1635
1636        response.send(proto::Ack {})?;
1637        Ok(())
1638    }
1639
1640    async fn join_channel(
1641        self: Arc<Self>,
1642        request: TypedEnvelope<proto::JoinChannel>,
1643        response: Response<proto::JoinChannel>,
1644    ) -> Result<()> {
1645        let user_id = self
1646            .store()
1647            .await
1648            .user_id_for_connection(request.sender_id)?;
1649        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1650        if !self
1651            .app_state
1652            .db
1653            .can_user_access_channel(user_id, channel_id)
1654            .await?
1655        {
1656            Err(anyhow!("access denied"))?;
1657        }
1658
1659        self.store()
1660            .await
1661            .join_channel(request.sender_id, channel_id);
1662        let messages = self
1663            .app_state
1664            .db
1665            .get_channel_messages(channel_id, MESSAGE_COUNT_PER_PAGE, None)
1666            .await?
1667            .into_iter()
1668            .map(|msg| proto::ChannelMessage {
1669                id: msg.id.to_proto(),
1670                body: msg.body,
1671                timestamp: msg.sent_at.unix_timestamp() as u64,
1672                sender_id: msg.sender_id.to_proto(),
1673                nonce: Some(msg.nonce.as_u128().into()),
1674            })
1675            .collect::<Vec<_>>();
1676        response.send(proto::JoinChannelResponse {
1677            done: messages.len() < MESSAGE_COUNT_PER_PAGE,
1678            messages,
1679        })?;
1680        Ok(())
1681    }
1682
1683    async fn leave_channel(
1684        self: Arc<Self>,
1685        request: TypedEnvelope<proto::LeaveChannel>,
1686    ) -> Result<()> {
1687        let user_id = self
1688            .store()
1689            .await
1690            .user_id_for_connection(request.sender_id)?;
1691        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1692        if !self
1693            .app_state
1694            .db
1695            .can_user_access_channel(user_id, channel_id)
1696            .await?
1697        {
1698            Err(anyhow!("access denied"))?;
1699        }
1700
1701        self.store()
1702            .await
1703            .leave_channel(request.sender_id, channel_id);
1704
1705        Ok(())
1706    }
1707
1708    async fn send_channel_message(
1709        self: Arc<Self>,
1710        request: TypedEnvelope<proto::SendChannelMessage>,
1711        response: Response<proto::SendChannelMessage>,
1712    ) -> Result<()> {
1713        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1714        let user_id;
1715        let connection_ids;
1716        {
1717            let state = self.store().await;
1718            user_id = state.user_id_for_connection(request.sender_id)?;
1719            connection_ids = state.channel_connection_ids(channel_id)?;
1720        }
1721
1722        // Validate the message body.
1723        let body = request.payload.body.trim().to_string();
1724        if body.len() > MAX_MESSAGE_LEN {
1725            return Err(anyhow!("message is too long"))?;
1726        }
1727        if body.is_empty() {
1728            return Err(anyhow!("message can't be blank"))?;
1729        }
1730
1731        let timestamp = OffsetDateTime::now_utc();
1732        let nonce = request
1733            .payload
1734            .nonce
1735            .ok_or_else(|| anyhow!("nonce can't be blank"))?;
1736
1737        let message_id = self
1738            .app_state
1739            .db
1740            .create_channel_message(channel_id, user_id, &body, timestamp, nonce.clone().into())
1741            .await?
1742            .to_proto();
1743        let message = proto::ChannelMessage {
1744            sender_id: user_id.to_proto(),
1745            id: message_id,
1746            body,
1747            timestamp: timestamp.unix_timestamp() as u64,
1748            nonce: Some(nonce),
1749        };
1750        broadcast(request.sender_id, connection_ids, |conn_id| {
1751            self.peer.send(
1752                conn_id,
1753                proto::ChannelMessageSent {
1754                    channel_id: channel_id.to_proto(),
1755                    message: Some(message.clone()),
1756                },
1757            )
1758        });
1759        response.send(proto::SendChannelMessageResponse {
1760            message: Some(message),
1761        })?;
1762        Ok(())
1763    }
1764
1765    async fn get_channel_messages(
1766        self: Arc<Self>,
1767        request: TypedEnvelope<proto::GetChannelMessages>,
1768        response: Response<proto::GetChannelMessages>,
1769    ) -> Result<()> {
1770        let user_id = self
1771            .store()
1772            .await
1773            .user_id_for_connection(request.sender_id)?;
1774        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1775        if !self
1776            .app_state
1777            .db
1778            .can_user_access_channel(user_id, channel_id)
1779            .await?
1780        {
1781            Err(anyhow!("access denied"))?;
1782        }
1783
1784        let messages = self
1785            .app_state
1786            .db
1787            .get_channel_messages(
1788                channel_id,
1789                MESSAGE_COUNT_PER_PAGE,
1790                Some(MessageId::from_proto(request.payload.before_message_id)),
1791            )
1792            .await?
1793            .into_iter()
1794            .map(|msg| proto::ChannelMessage {
1795                id: msg.id.to_proto(),
1796                body: msg.body,
1797                timestamp: msg.sent_at.unix_timestamp() as u64,
1798                sender_id: msg.sender_id.to_proto(),
1799                nonce: Some(msg.nonce.as_u128().into()),
1800            })
1801            .collect::<Vec<_>>();
1802        response.send(proto::GetChannelMessagesResponse {
1803            done: messages.len() < MESSAGE_COUNT_PER_PAGE,
1804            messages,
1805        })?;
1806        Ok(())
1807    }
1808
1809    async fn update_diff_base(
1810        self: Arc<Server>,
1811        request: TypedEnvelope<proto::UpdateDiffBase>,
1812    ) -> Result<()> {
1813        let receiver_ids = self.store().await.project_connection_ids(
1814            ProjectId::from_proto(request.payload.project_id),
1815            request.sender_id,
1816        )?;
1817        broadcast(request.sender_id, receiver_ids, |connection_id| {
1818            self.peer
1819                .forward_send(request.sender_id, connection_id, request.payload.clone())
1820        });
1821        Ok(())
1822    }
1823
1824    async fn get_private_user_info(
1825        self: Arc<Self>,
1826        request: TypedEnvelope<proto::GetPrivateUserInfo>,
1827        response: Response<proto::GetPrivateUserInfo>,
1828    ) -> Result<()> {
1829        let user_id = self
1830            .store()
1831            .await
1832            .user_id_for_connection(request.sender_id)?;
1833        let metrics_id = self.app_state.db.get_user_metrics_id(user_id).await?;
1834        let user = self
1835            .app_state
1836            .db
1837            .get_user_by_id(user_id)
1838            .await?
1839            .ok_or_else(|| anyhow!("user not found"))?;
1840        response.send(proto::GetPrivateUserInfoResponse {
1841            metrics_id,
1842            staff: user.admin,
1843        })?;
1844        Ok(())
1845    }
1846
1847    pub(crate) async fn store(&self) -> StoreGuard<'_> {
1848        #[cfg(test)]
1849        tokio::task::yield_now().await;
1850        let guard = self.store.lock().await;
1851        #[cfg(test)]
1852        tokio::task::yield_now().await;
1853        StoreGuard {
1854            guard,
1855            _not_send: PhantomData,
1856        }
1857    }
1858
1859    pub async fn snapshot<'a>(self: &'a Arc<Self>) -> ServerSnapshot<'a> {
1860        ServerSnapshot {
1861            store: self.store().await,
1862            peer: &self.peer,
1863        }
1864    }
1865}
1866
1867impl<'a> Deref for StoreGuard<'a> {
1868    type Target = Store;
1869
1870    fn deref(&self) -> &Self::Target {
1871        &*self.guard
1872    }
1873}
1874
1875impl<'a> DerefMut for StoreGuard<'a> {
1876    fn deref_mut(&mut self) -> &mut Self::Target {
1877        &mut *self.guard
1878    }
1879}
1880
1881impl<'a> Drop for StoreGuard<'a> {
1882    fn drop(&mut self) {
1883        #[cfg(test)]
1884        self.check_invariants();
1885    }
1886}
1887
1888impl Executor for RealExecutor {
1889    type Sleep = Sleep;
1890
1891    fn spawn_detached<F: 'static + Send + Future<Output = ()>>(&self, future: F) {
1892        tokio::task::spawn(future);
1893    }
1894
1895    fn sleep(&self, duration: Duration) -> Self::Sleep {
1896        tokio::time::sleep(duration)
1897    }
1898}
1899
1900fn broadcast<F>(
1901    sender_id: ConnectionId,
1902    receiver_ids: impl IntoIterator<Item = ConnectionId>,
1903    mut f: F,
1904) where
1905    F: FnMut(ConnectionId) -> anyhow::Result<()>,
1906{
1907    for receiver_id in receiver_ids {
1908        if receiver_id != sender_id {
1909            f(receiver_id).trace_err();
1910        }
1911    }
1912}
1913
1914lazy_static! {
1915    static ref ZED_PROTOCOL_VERSION: HeaderName = HeaderName::from_static("x-zed-protocol-version");
1916}
1917
1918pub struct ProtocolVersion(u32);
1919
1920impl Header for ProtocolVersion {
1921    fn name() -> &'static HeaderName {
1922        &ZED_PROTOCOL_VERSION
1923    }
1924
1925    fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1926    where
1927        Self: Sized,
1928        I: Iterator<Item = &'i axum::http::HeaderValue>,
1929    {
1930        let version = values
1931            .next()
1932            .ok_or_else(axum::headers::Error::invalid)?
1933            .to_str()
1934            .map_err(|_| axum::headers::Error::invalid())?
1935            .parse()
1936            .map_err(|_| axum::headers::Error::invalid())?;
1937        Ok(Self(version))
1938    }
1939
1940    fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1941        values.extend([self.0.to_string().parse().unwrap()]);
1942    }
1943}
1944
1945pub fn routes(server: Arc<Server>) -> Router<Body> {
1946    Router::new()
1947        .route("/rpc", get(handle_websocket_request))
1948        .layer(
1949            ServiceBuilder::new()
1950                .layer(Extension(server.app_state.clone()))
1951                .layer(middleware::from_fn(auth::validate_header)),
1952        )
1953        .route("/metrics", get(handle_metrics))
1954        .layer(Extension(server))
1955}
1956
1957pub async fn handle_websocket_request(
1958    TypedHeader(ProtocolVersion(protocol_version)): TypedHeader<ProtocolVersion>,
1959    ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
1960    Extension(server): Extension<Arc<Server>>,
1961    Extension(user): Extension<User>,
1962    ws: WebSocketUpgrade,
1963) -> axum::response::Response {
1964    if protocol_version != rpc::PROTOCOL_VERSION {
1965        return (
1966            StatusCode::UPGRADE_REQUIRED,
1967            "client must be upgraded".to_string(),
1968        )
1969            .into_response();
1970    }
1971    let socket_address = socket_address.to_string();
1972    ws.on_upgrade(move |socket| {
1973        use util::ResultExt;
1974        let socket = socket
1975            .map_ok(to_tungstenite_message)
1976            .err_into()
1977            .with(|message| async move { Ok(to_axum_message(message)) });
1978        let connection = Connection::new(Box::pin(socket));
1979        async move {
1980            server
1981                .handle_connection(connection, socket_address, user, None, RealExecutor)
1982                .await
1983                .log_err();
1984        }
1985    })
1986}
1987
1988pub async fn handle_metrics(Extension(server): Extension<Arc<Server>>) -> axum::response::Response {
1989    // We call `store_mut` here for its side effects of updating metrics.
1990    let metrics = server.store().await.metrics();
1991    METRIC_CONNECTIONS.set(metrics.connections as _);
1992    METRIC_REGISTERED_PROJECTS.set(metrics.registered_projects as _);
1993    METRIC_ACTIVE_PROJECTS.set(metrics.active_projects as _);
1994    METRIC_SHARED_PROJECTS.set(metrics.shared_projects as _);
1995
1996    let encoder = prometheus::TextEncoder::new();
1997    let metric_families = prometheus::gather();
1998    match encoder.encode_to_string(&metric_families) {
1999        Ok(string) => (StatusCode::OK, string).into_response(),
2000        Err(error) => (
2001            StatusCode::INTERNAL_SERVER_ERROR,
2002            format!("failed to encode metrics {:?}", error),
2003        )
2004            .into_response(),
2005    }
2006}
2007
2008fn to_axum_message(message: TungsteniteMessage) -> AxumMessage {
2009    match message {
2010        TungsteniteMessage::Text(payload) => AxumMessage::Text(payload),
2011        TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload),
2012        TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload),
2013        TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload),
2014        TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
2015            code: frame.code.into(),
2016            reason: frame.reason,
2017        })),
2018    }
2019}
2020
2021fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
2022    match message {
2023        AxumMessage::Text(payload) => TungsteniteMessage::Text(payload),
2024        AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload),
2025        AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload),
2026        AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload),
2027        AxumMessage::Close(frame) => {
2028            TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
2029                code: frame.code.into(),
2030                reason: frame.reason,
2031            }))
2032        }
2033    }
2034}
2035
2036pub trait ResultExt {
2037    type Ok;
2038
2039    fn trace_err(self) -> Option<Self::Ok>;
2040}
2041
2042impl<T, E> ResultExt for Result<T, E>
2043where
2044    E: std::fmt::Debug,
2045{
2046    type Ok = T;
2047
2048    fn trace_err(self) -> Option<T> {
2049        match self {
2050            Ok(value) => Some(value),
2051            Err(error) => {
2052                tracing::error!("{:?}", error);
2053                None
2054            }
2055        }
2056    }
2057}