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