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                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.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_unregister = Vec::new();
 479        let removed_user_id;
 480        {
 481            let mut store = self.store().await;
 482            let removed_connection = store.remove_connection(connection_id)?;
 483
 484            for (project_id, project) in removed_connection.hosted_projects {
 485                projects_to_unregister.push(project_id);
 486                broadcast(connection_id, project.guests.keys().copied(), |conn_id| {
 487                    self.peer.send(
 488                        conn_id,
 489                        proto::UnregisterProject {
 490                            project_id: project_id.to_proto(),
 491                        },
 492                    )
 493                });
 494
 495                for (_, receipts) in project.join_requests {
 496                    for receipt in receipts {
 497                        self.peer.respond(
 498                            receipt,
 499                            proto::JoinProjectResponse {
 500                                variant: Some(proto::join_project_response::Variant::Decline(
 501                                    proto::join_project_response::Decline {
 502                                        reason: proto::join_project_response::decline::Reason::WentOffline as i32
 503                                    },
 504                                )),
 505                            },
 506                        )?;
 507                    }
 508                }
 509            }
 510
 511            for project_id in removed_connection.guest_project_ids {
 512                if let Some(project) = store.project(project_id).trace_err() {
 513                    broadcast(connection_id, project.connection_ids(), |conn_id| {
 514                        self.peer.send(
 515                            conn_id,
 516                            proto::RemoveProjectCollaborator {
 517                                project_id: project_id.to_proto(),
 518                                peer_id: connection_id.0,
 519                            },
 520                        )
 521                    });
 522                    if project.guests.is_empty() {
 523                        self.peer
 524                            .send(
 525                                project.host_connection_id,
 526                                proto::ProjectUnshared {
 527                                    project_id: project_id.to_proto(),
 528                                },
 529                            )
 530                            .trace_err();
 531                    }
 532                }
 533            }
 534
 535            if let Some(room) = removed_connection
 536                .room_id
 537                .and_then(|room_id| store.room(room_id))
 538            {
 539                self.room_updated(room);
 540            }
 541
 542            removed_user_id = removed_connection.user_id;
 543        };
 544
 545        self.update_user_contacts(removed_user_id).await.trace_err();
 546
 547        for project_id in projects_to_unregister {
 548            self.app_state
 549                .db
 550                .unregister_project(project_id)
 551                .await
 552                .trace_err();
 553        }
 554
 555        Ok(())
 556    }
 557
 558    pub async fn invite_code_redeemed(
 559        self: &Arc<Self>,
 560        code: &str,
 561        invitee_id: UserId,
 562    ) -> Result<()> {
 563        let user = self.app_state.db.get_user_for_invite_code(code).await?;
 564        let store = self.store().await;
 565        let invitee_contact = store.contact_for_user(invitee_id, true);
 566        for connection_id in store.connection_ids_for_user(user.id) {
 567            self.peer.send(
 568                connection_id,
 569                proto::UpdateContacts {
 570                    contacts: vec![invitee_contact.clone()],
 571                    ..Default::default()
 572                },
 573            )?;
 574            self.peer.send(
 575                connection_id,
 576                proto::UpdateInviteInfo {
 577                    url: format!("{}{}", self.app_state.invite_link_prefix, code),
 578                    count: user.invite_count as u32,
 579                },
 580            )?;
 581        }
 582        Ok(())
 583    }
 584
 585    pub async fn invite_count_updated(self: &Arc<Self>, user_id: UserId) -> Result<()> {
 586        if let Some(user) = self.app_state.db.get_user_by_id(user_id).await? {
 587            if let Some(invite_code) = &user.invite_code {
 588                let store = self.store().await;
 589                for connection_id in store.connection_ids_for_user(user_id) {
 590                    self.peer.send(
 591                        connection_id,
 592                        proto::UpdateInviteInfo {
 593                            url: format!("{}{}", self.app_state.invite_link_prefix, invite_code),
 594                            count: user.invite_count as u32,
 595                        },
 596                    )?;
 597                }
 598            }
 599        }
 600        Ok(())
 601    }
 602
 603    async fn ping(
 604        self: Arc<Server>,
 605        _: TypedEnvelope<proto::Ping>,
 606        response: Response<proto::Ping>,
 607    ) -> Result<()> {
 608        response.send(proto::Ack {})?;
 609        Ok(())
 610    }
 611
 612    async fn create_room(
 613        self: Arc<Server>,
 614        request: TypedEnvelope<proto::CreateRoom>,
 615        response: Response<proto::CreateRoom>,
 616    ) -> Result<()> {
 617        let room_id = self.store().await.create_room(request.sender_id)?;
 618        response.send(proto::CreateRoomResponse { id: room_id })?;
 619        Ok(())
 620    }
 621
 622    async fn join_room(
 623        self: Arc<Server>,
 624        request: TypedEnvelope<proto::JoinRoom>,
 625        response: Response<proto::JoinRoom>,
 626    ) -> Result<()> {
 627        let room_id = request.payload.id;
 628        let mut store = self.store().await;
 629        let (room, recipient_connection_ids) = store.join_room(room_id, request.sender_id)?;
 630        for recipient_id in recipient_connection_ids {
 631            self.peer
 632                .send(recipient_id, proto::CancelCall {})
 633                .trace_err();
 634        }
 635        response.send(proto::JoinRoomResponse {
 636            room: Some(room.clone()),
 637        })?;
 638        self.room_updated(room);
 639        Ok(())
 640    }
 641
 642    async fn leave_room(self: Arc<Server>, message: TypedEnvelope<proto::LeaveRoom>) -> Result<()> {
 643        let room_id = message.payload.id;
 644        let mut store = self.store().await;
 645        let room = store.leave_room(room_id, message.sender_id)?;
 646        self.room_updated(room);
 647        Ok(())
 648    }
 649
 650    async fn call(
 651        self: Arc<Server>,
 652        request: TypedEnvelope<proto::Call>,
 653        response: Response<proto::Call>,
 654    ) -> Result<()> {
 655        let recipient_user_id = UserId::from_proto(request.payload.recipient_user_id);
 656        let room_id = request.payload.room_id;
 657        let mut calls = {
 658            let mut store = self.store().await;
 659            let (room, recipient_connection_ids, incoming_call) =
 660                store.call(room_id, request.sender_id, recipient_user_id)?;
 661            self.room_updated(room);
 662            recipient_connection_ids
 663                .into_iter()
 664                .map(|recipient_connection_id| {
 665                    self.peer
 666                        .request(recipient_connection_id, incoming_call.clone())
 667                })
 668                .collect::<FuturesUnordered<_>>()
 669        };
 670
 671        while let Some(call_response) = calls.next().await {
 672            match call_response.as_ref() {
 673                Ok(_) => {
 674                    response.send(proto::Ack {})?;
 675                    return Ok(());
 676                }
 677                Err(_) => {
 678                    call_response.trace_err();
 679                }
 680            }
 681        }
 682
 683        {
 684            let mut store = self.store().await;
 685            let room = store.call_failed(room_id, recipient_user_id)?;
 686            self.room_updated(&room);
 687        }
 688
 689        Err(anyhow!("failed to ring call recipient"))?
 690    }
 691
 692    async fn decline_call(
 693        self: Arc<Server>,
 694        message: TypedEnvelope<proto::DeclineCall>,
 695    ) -> Result<()> {
 696        let mut store = self.store().await;
 697        let (room, recipient_connection_ids) = store.call_declined(message.sender_id)?;
 698        for recipient_id in recipient_connection_ids {
 699            self.peer
 700                .send(recipient_id, proto::CancelCall {})
 701                .trace_err();
 702        }
 703        self.room_updated(room);
 704        Ok(())
 705    }
 706
 707    fn room_updated(&self, room: &proto::Room) {
 708        for participant in &room.participants {
 709            self.peer
 710                .send(
 711                    ConnectionId(participant.peer_id),
 712                    proto::RoomUpdated {
 713                        room: Some(room.clone()),
 714                    },
 715                )
 716                .trace_err();
 717        }
 718    }
 719
 720    async fn register_project(
 721        self: Arc<Server>,
 722        request: TypedEnvelope<proto::RegisterProject>,
 723        response: Response<proto::RegisterProject>,
 724    ) -> Result<()> {
 725        let user_id = self
 726            .store()
 727            .await
 728            .user_id_for_connection(request.sender_id)?;
 729        let project_id = self.app_state.db.register_project(user_id).await?;
 730        self.store().await.register_project(
 731            request.sender_id,
 732            project_id,
 733            request.payload.online,
 734        )?;
 735
 736        response.send(proto::RegisterProjectResponse {
 737            project_id: project_id.to_proto(),
 738        })?;
 739
 740        Ok(())
 741    }
 742
 743    async fn unregister_project(
 744        self: Arc<Server>,
 745        request: TypedEnvelope<proto::UnregisterProject>,
 746        response: Response<proto::UnregisterProject>,
 747    ) -> Result<()> {
 748        let project_id = ProjectId::from_proto(request.payload.project_id);
 749        let (user_id, project) = {
 750            let mut state = self.store().await;
 751            let project = state.unregister_project(project_id, request.sender_id)?;
 752            (state.user_id_for_connection(request.sender_id)?, project)
 753        };
 754        self.app_state.db.unregister_project(project_id).await?;
 755
 756        broadcast(
 757            request.sender_id,
 758            project.guests.keys().copied(),
 759            |conn_id| {
 760                self.peer.send(
 761                    conn_id,
 762                    proto::UnregisterProject {
 763                        project_id: project_id.to_proto(),
 764                    },
 765                )
 766            },
 767        );
 768        for (_, receipts) in project.join_requests {
 769            for receipt in receipts {
 770                self.peer.respond(
 771                    receipt,
 772                    proto::JoinProjectResponse {
 773                        variant: Some(proto::join_project_response::Variant::Decline(
 774                            proto::join_project_response::Decline {
 775                                reason: proto::join_project_response::decline::Reason::Closed
 776                                    as i32,
 777                            },
 778                        )),
 779                    },
 780                )?;
 781            }
 782        }
 783
 784        // Send out the `UpdateContacts` message before responding to the unregister
 785        // request. This way, when the project's host can keep track of the project's
 786        // remote id until after they've received the `UpdateContacts` message for
 787        // themself.
 788        self.update_user_contacts(user_id).await?;
 789        response.send(proto::Ack {})?;
 790
 791        Ok(())
 792    }
 793
 794    async fn update_user_contacts(self: &Arc<Server>, user_id: UserId) -> Result<()> {
 795        let contacts = self.app_state.db.get_contacts(user_id).await?;
 796        let store = self.store().await;
 797        let updated_contact = store.contact_for_user(user_id, false);
 798        for contact in contacts {
 799            if let db::Contact::Accepted {
 800                user_id: contact_user_id,
 801                ..
 802            } = contact
 803            {
 804                for contact_conn_id in store.connection_ids_for_user(contact_user_id) {
 805                    self.peer
 806                        .send(
 807                            contact_conn_id,
 808                            proto::UpdateContacts {
 809                                contacts: vec![updated_contact.clone()],
 810                                remove_contacts: Default::default(),
 811                                incoming_requests: Default::default(),
 812                                remove_incoming_requests: Default::default(),
 813                                outgoing_requests: Default::default(),
 814                                remove_outgoing_requests: Default::default(),
 815                            },
 816                        )
 817                        .trace_err();
 818                }
 819            }
 820        }
 821        Ok(())
 822    }
 823
 824    async fn join_project(
 825        self: Arc<Server>,
 826        request: TypedEnvelope<proto::JoinProject>,
 827        response: Response<proto::JoinProject>,
 828    ) -> Result<()> {
 829        let project_id = ProjectId::from_proto(request.payload.project_id);
 830
 831        let host_user_id;
 832        let guest_user_id;
 833        let host_connection_id;
 834        {
 835            let state = self.store().await;
 836            let project = state.project(project_id)?;
 837            host_user_id = project.host.user_id;
 838            host_connection_id = project.host_connection_id;
 839            guest_user_id = state.user_id_for_connection(request.sender_id)?;
 840        };
 841
 842        tracing::info!(%project_id, %host_user_id, %host_connection_id, "join project");
 843        let has_contact = self
 844            .app_state
 845            .db
 846            .has_contact(guest_user_id, host_user_id)
 847            .await?;
 848        if !has_contact {
 849            return Err(anyhow!("no such project"))?;
 850        }
 851
 852        self.store().await.request_join_project(
 853            guest_user_id,
 854            project_id,
 855            response.into_receipt(),
 856        )?;
 857        self.peer.send(
 858            host_connection_id,
 859            proto::RequestJoinProject {
 860                project_id: project_id.to_proto(),
 861                requester_id: guest_user_id.to_proto(),
 862            },
 863        )?;
 864        Ok(())
 865    }
 866
 867    async fn respond_to_join_project_request(
 868        self: Arc<Server>,
 869        request: TypedEnvelope<proto::RespondToJoinProjectRequest>,
 870    ) -> Result<()> {
 871        let host_user_id;
 872
 873        {
 874            let mut state = self.store().await;
 875            let project_id = ProjectId::from_proto(request.payload.project_id);
 876            let project = state.project(project_id)?;
 877            if project.host_connection_id != request.sender_id {
 878                Err(anyhow!("no such connection"))?;
 879            }
 880
 881            host_user_id = project.host.user_id;
 882            let guest_user_id = UserId::from_proto(request.payload.requester_id);
 883
 884            if !request.payload.allow {
 885                let receipts = state
 886                    .deny_join_project_request(request.sender_id, guest_user_id, project_id)
 887                    .ok_or_else(|| anyhow!("no such request"))?;
 888                for receipt in receipts {
 889                    self.peer.respond(
 890                        receipt,
 891                        proto::JoinProjectResponse {
 892                            variant: Some(proto::join_project_response::Variant::Decline(
 893                                proto::join_project_response::Decline {
 894                                    reason: proto::join_project_response::decline::Reason::Declined
 895                                        as i32,
 896                                },
 897                            )),
 898                        },
 899                    )?;
 900                }
 901                return Ok(());
 902            }
 903
 904            let (receipts_with_replica_ids, project) = state
 905                .accept_join_project_request(request.sender_id, guest_user_id, project_id)
 906                .ok_or_else(|| anyhow!("no such request"))?;
 907
 908            let peer_count = project.guests.len();
 909            let mut collaborators = Vec::with_capacity(peer_count);
 910            collaborators.push(proto::Collaborator {
 911                peer_id: project.host_connection_id.0,
 912                replica_id: 0,
 913                user_id: project.host.user_id.to_proto(),
 914            });
 915            let worktrees = project
 916                .worktrees
 917                .iter()
 918                .map(|(id, worktree)| proto::WorktreeMetadata {
 919                    id: *id,
 920                    root_name: worktree.root_name.clone(),
 921                    visible: worktree.visible,
 922                })
 923                .collect::<Vec<_>>();
 924
 925            // Add all guests other than the requesting user's own connections as collaborators
 926            for (guest_conn_id, guest) in &project.guests {
 927                if receipts_with_replica_ids
 928                    .iter()
 929                    .all(|(receipt, _)| receipt.sender_id != *guest_conn_id)
 930                {
 931                    collaborators.push(proto::Collaborator {
 932                        peer_id: guest_conn_id.0,
 933                        replica_id: guest.replica_id as u32,
 934                        user_id: guest.user_id.to_proto(),
 935                    });
 936                }
 937            }
 938
 939            for conn_id in project.connection_ids() {
 940                for (receipt, replica_id) in &receipts_with_replica_ids {
 941                    if conn_id != receipt.sender_id {
 942                        self.peer.send(
 943                            conn_id,
 944                            proto::AddProjectCollaborator {
 945                                project_id: project_id.to_proto(),
 946                                collaborator: Some(proto::Collaborator {
 947                                    peer_id: receipt.sender_id.0,
 948                                    replica_id: *replica_id as u32,
 949                                    user_id: guest_user_id.to_proto(),
 950                                }),
 951                            },
 952                        )?;
 953                    }
 954                }
 955            }
 956
 957            // First, we send the metadata associated with each worktree.
 958            for (receipt, replica_id) in &receipts_with_replica_ids {
 959                self.peer.respond(
 960                    *receipt,
 961                    proto::JoinProjectResponse {
 962                        variant: Some(proto::join_project_response::Variant::Accept(
 963                            proto::join_project_response::Accept {
 964                                worktrees: worktrees.clone(),
 965                                replica_id: *replica_id as u32,
 966                                collaborators: collaborators.clone(),
 967                                language_servers: project.language_servers.clone(),
 968                            },
 969                        )),
 970                    },
 971                )?;
 972            }
 973
 974            for (worktree_id, worktree) in &project.worktrees {
 975                #[cfg(any(test, feature = "test-support"))]
 976                const MAX_CHUNK_SIZE: usize = 2;
 977                #[cfg(not(any(test, feature = "test-support")))]
 978                const MAX_CHUNK_SIZE: usize = 256;
 979
 980                // Stream this worktree's entries.
 981                let message = proto::UpdateWorktree {
 982                    project_id: project_id.to_proto(),
 983                    worktree_id: *worktree_id,
 984                    root_name: worktree.root_name.clone(),
 985                    updated_entries: worktree.entries.values().cloned().collect(),
 986                    removed_entries: Default::default(),
 987                    scan_id: worktree.scan_id,
 988                    is_last_update: worktree.is_complete,
 989                };
 990                for update in proto::split_worktree_update(message, MAX_CHUNK_SIZE) {
 991                    for (receipt, _) in &receipts_with_replica_ids {
 992                        self.peer.send(receipt.sender_id, update.clone())?;
 993                    }
 994                }
 995
 996                // Stream this worktree's diagnostics.
 997                for summary in worktree.diagnostic_summaries.values() {
 998                    for (receipt, _) in &receipts_with_replica_ids {
 999                        self.peer.send(
1000                            receipt.sender_id,
1001                            proto::UpdateDiagnosticSummary {
1002                                project_id: project_id.to_proto(),
1003                                worktree_id: *worktree_id,
1004                                summary: Some(summary.clone()),
1005                            },
1006                        )?;
1007                    }
1008                }
1009            }
1010        }
1011
1012        self.update_user_contacts(host_user_id).await?;
1013        Ok(())
1014    }
1015
1016    async fn leave_project(
1017        self: Arc<Server>,
1018        request: TypedEnvelope<proto::LeaveProject>,
1019    ) -> Result<()> {
1020        let sender_id = request.sender_id;
1021        let project_id = ProjectId::from_proto(request.payload.project_id);
1022        let project;
1023        {
1024            let mut store = self.store().await;
1025            project = store.leave_project(sender_id, project_id)?;
1026            tracing::info!(
1027                %project_id,
1028                host_user_id = %project.host_user_id,
1029                host_connection_id = %project.host_connection_id,
1030                "leave project"
1031            );
1032
1033            if project.remove_collaborator {
1034                broadcast(sender_id, project.connection_ids, |conn_id| {
1035                    self.peer.send(
1036                        conn_id,
1037                        proto::RemoveProjectCollaborator {
1038                            project_id: project_id.to_proto(),
1039                            peer_id: sender_id.0,
1040                        },
1041                    )
1042                });
1043            }
1044
1045            if let Some(requester_id) = project.cancel_request {
1046                self.peer.send(
1047                    project.host_connection_id,
1048                    proto::JoinProjectRequestCancelled {
1049                        project_id: project_id.to_proto(),
1050                        requester_id: requester_id.to_proto(),
1051                    },
1052                )?;
1053            }
1054
1055            if project.unshare {
1056                self.peer.send(
1057                    project.host_connection_id,
1058                    proto::ProjectUnshared {
1059                        project_id: project_id.to_proto(),
1060                    },
1061                )?;
1062            }
1063        }
1064        self.update_user_contacts(project.host_user_id).await?;
1065        Ok(())
1066    }
1067
1068    async fn update_project(
1069        self: Arc<Server>,
1070        request: TypedEnvelope<proto::UpdateProject>,
1071    ) -> Result<()> {
1072        let project_id = ProjectId::from_proto(request.payload.project_id);
1073        let user_id;
1074        {
1075            let mut state = self.store().await;
1076            user_id = state.user_id_for_connection(request.sender_id)?;
1077            let guest_connection_ids = state
1078                .read_project(project_id, request.sender_id)?
1079                .guest_connection_ids();
1080            let unshared_project = state.update_project(
1081                project_id,
1082                &request.payload.worktrees,
1083                request.payload.online,
1084                request.sender_id,
1085            )?;
1086
1087            if let Some(unshared_project) = unshared_project {
1088                broadcast(
1089                    request.sender_id,
1090                    unshared_project.guests.keys().copied(),
1091                    |conn_id| {
1092                        self.peer.send(
1093                            conn_id,
1094                            proto::UnregisterProject {
1095                                project_id: project_id.to_proto(),
1096                            },
1097                        )
1098                    },
1099                );
1100                for (_, receipts) in unshared_project.pending_join_requests {
1101                    for receipt in receipts {
1102                        self.peer.respond(
1103                            receipt,
1104                            proto::JoinProjectResponse {
1105                                variant: Some(proto::join_project_response::Variant::Decline(
1106                                    proto::join_project_response::Decline {
1107                                        reason:
1108                                            proto::join_project_response::decline::Reason::Closed
1109                                                as i32,
1110                                    },
1111                                )),
1112                            },
1113                        )?;
1114                    }
1115                }
1116            } else {
1117                broadcast(request.sender_id, guest_connection_ids, |connection_id| {
1118                    self.peer.forward_send(
1119                        request.sender_id,
1120                        connection_id,
1121                        request.payload.clone(),
1122                    )
1123                });
1124            }
1125        };
1126
1127        self.update_user_contacts(user_id).await?;
1128        Ok(())
1129    }
1130
1131    async fn register_project_activity(
1132        self: Arc<Server>,
1133        request: TypedEnvelope<proto::RegisterProjectActivity>,
1134    ) -> Result<()> {
1135        self.store().await.register_project_activity(
1136            ProjectId::from_proto(request.payload.project_id),
1137            request.sender_id,
1138        )?;
1139        Ok(())
1140    }
1141
1142    async fn update_worktree(
1143        self: Arc<Server>,
1144        request: TypedEnvelope<proto::UpdateWorktree>,
1145        response: Response<proto::UpdateWorktree>,
1146    ) -> Result<()> {
1147        let project_id = ProjectId::from_proto(request.payload.project_id);
1148        let worktree_id = request.payload.worktree_id;
1149        let (connection_ids, metadata_changed) = {
1150            let mut store = self.store().await;
1151            let (connection_ids, metadata_changed) = store.update_worktree(
1152                request.sender_id,
1153                project_id,
1154                worktree_id,
1155                &request.payload.root_name,
1156                &request.payload.removed_entries,
1157                &request.payload.updated_entries,
1158                request.payload.scan_id,
1159                request.payload.is_last_update,
1160            )?;
1161            (connection_ids, metadata_changed)
1162        };
1163
1164        broadcast(request.sender_id, connection_ids, |connection_id| {
1165            self.peer
1166                .forward_send(request.sender_id, connection_id, request.payload.clone())
1167        });
1168        if metadata_changed {
1169            let user_id = self
1170                .store()
1171                .await
1172                .user_id_for_connection(request.sender_id)?;
1173            self.update_user_contacts(user_id).await?;
1174        }
1175        response.send(proto::Ack {})?;
1176        Ok(())
1177    }
1178
1179    async fn update_worktree_extensions(
1180        self: Arc<Server>,
1181        request: TypedEnvelope<proto::UpdateWorktreeExtensions>,
1182    ) -> Result<()> {
1183        let project_id = ProjectId::from_proto(request.payload.project_id);
1184        let worktree_id = request.payload.worktree_id;
1185        let extensions = request
1186            .payload
1187            .extensions
1188            .into_iter()
1189            .zip(request.payload.counts)
1190            .collect();
1191        self.app_state
1192            .db
1193            .update_worktree_extensions(project_id, worktree_id, extensions)
1194            .await?;
1195        Ok(())
1196    }
1197
1198    async fn update_diagnostic_summary(
1199        self: Arc<Server>,
1200        request: TypedEnvelope<proto::UpdateDiagnosticSummary>,
1201    ) -> Result<()> {
1202        let summary = request
1203            .payload
1204            .summary
1205            .clone()
1206            .ok_or_else(|| anyhow!("invalid summary"))?;
1207        let receiver_ids = self.store().await.update_diagnostic_summary(
1208            ProjectId::from_proto(request.payload.project_id),
1209            request.payload.worktree_id,
1210            request.sender_id,
1211            summary,
1212        )?;
1213
1214        broadcast(request.sender_id, receiver_ids, |connection_id| {
1215            self.peer
1216                .forward_send(request.sender_id, connection_id, request.payload.clone())
1217        });
1218        Ok(())
1219    }
1220
1221    async fn start_language_server(
1222        self: Arc<Server>,
1223        request: TypedEnvelope<proto::StartLanguageServer>,
1224    ) -> Result<()> {
1225        let receiver_ids = self.store().await.start_language_server(
1226            ProjectId::from_proto(request.payload.project_id),
1227            request.sender_id,
1228            request
1229                .payload
1230                .server
1231                .clone()
1232                .ok_or_else(|| anyhow!("invalid language server"))?,
1233        )?;
1234        broadcast(request.sender_id, receiver_ids, |connection_id| {
1235            self.peer
1236                .forward_send(request.sender_id, connection_id, request.payload.clone())
1237        });
1238        Ok(())
1239    }
1240
1241    async fn update_language_server(
1242        self: Arc<Server>,
1243        request: TypedEnvelope<proto::UpdateLanguageServer>,
1244    ) -> Result<()> {
1245        let receiver_ids = self.store().await.project_connection_ids(
1246            ProjectId::from_proto(request.payload.project_id),
1247            request.sender_id,
1248        )?;
1249        broadcast(request.sender_id, receiver_ids, |connection_id| {
1250            self.peer
1251                .forward_send(request.sender_id, connection_id, request.payload.clone())
1252        });
1253        Ok(())
1254    }
1255
1256    async fn forward_project_request<T>(
1257        self: Arc<Server>,
1258        request: TypedEnvelope<T>,
1259        response: Response<T>,
1260    ) -> Result<()>
1261    where
1262        T: EntityMessage + RequestMessage,
1263    {
1264        let project_id = ProjectId::from_proto(request.payload.remote_entity_id());
1265        let host_connection_id = self
1266            .store()
1267            .await
1268            .read_project(project_id, request.sender_id)?
1269            .host_connection_id;
1270        let payload = self
1271            .peer
1272            .forward_request(request.sender_id, host_connection_id, request.payload)
1273            .await?;
1274
1275        // Ensure project still exists by the time we get the response from the host.
1276        self.store()
1277            .await
1278            .read_project(project_id, request.sender_id)?;
1279
1280        response.send(payload)?;
1281        Ok(())
1282    }
1283
1284    async fn save_buffer(
1285        self: Arc<Server>,
1286        request: TypedEnvelope<proto::SaveBuffer>,
1287        response: Response<proto::SaveBuffer>,
1288    ) -> Result<()> {
1289        let project_id = ProjectId::from_proto(request.payload.project_id);
1290        let host = self
1291            .store()
1292            .await
1293            .read_project(project_id, request.sender_id)?
1294            .host_connection_id;
1295        let response_payload = self
1296            .peer
1297            .forward_request(request.sender_id, host, request.payload.clone())
1298            .await?;
1299
1300        let mut guests = self
1301            .store()
1302            .await
1303            .read_project(project_id, request.sender_id)?
1304            .connection_ids();
1305        guests.retain(|guest_connection_id| *guest_connection_id != request.sender_id);
1306        broadcast(host, guests, |conn_id| {
1307            self.peer
1308                .forward_send(host, conn_id, response_payload.clone())
1309        });
1310        response.send(response_payload)?;
1311        Ok(())
1312    }
1313
1314    async fn create_buffer_for_peer(
1315        self: Arc<Server>,
1316        request: TypedEnvelope<proto::CreateBufferForPeer>,
1317    ) -> Result<()> {
1318        self.peer.forward_send(
1319            request.sender_id,
1320            ConnectionId(request.payload.peer_id),
1321            request.payload,
1322        )?;
1323        Ok(())
1324    }
1325
1326    async fn update_buffer(
1327        self: Arc<Server>,
1328        request: TypedEnvelope<proto::UpdateBuffer>,
1329        response: Response<proto::UpdateBuffer>,
1330    ) -> Result<()> {
1331        let project_id = ProjectId::from_proto(request.payload.project_id);
1332        let receiver_ids = {
1333            let mut store = self.store().await;
1334            store.register_project_activity(project_id, request.sender_id)?;
1335            store.project_connection_ids(project_id, request.sender_id)?
1336        };
1337
1338        broadcast(request.sender_id, receiver_ids, |connection_id| {
1339            self.peer
1340                .forward_send(request.sender_id, connection_id, request.payload.clone())
1341        });
1342        response.send(proto::Ack {})?;
1343        Ok(())
1344    }
1345
1346    async fn update_buffer_file(
1347        self: Arc<Server>,
1348        request: TypedEnvelope<proto::UpdateBufferFile>,
1349    ) -> Result<()> {
1350        let receiver_ids = self.store().await.project_connection_ids(
1351            ProjectId::from_proto(request.payload.project_id),
1352            request.sender_id,
1353        )?;
1354        broadcast(request.sender_id, receiver_ids, |connection_id| {
1355            self.peer
1356                .forward_send(request.sender_id, connection_id, request.payload.clone())
1357        });
1358        Ok(())
1359    }
1360
1361    async fn buffer_reloaded(
1362        self: Arc<Server>,
1363        request: TypedEnvelope<proto::BufferReloaded>,
1364    ) -> Result<()> {
1365        let receiver_ids = self.store().await.project_connection_ids(
1366            ProjectId::from_proto(request.payload.project_id),
1367            request.sender_id,
1368        )?;
1369        broadcast(request.sender_id, receiver_ids, |connection_id| {
1370            self.peer
1371                .forward_send(request.sender_id, connection_id, request.payload.clone())
1372        });
1373        Ok(())
1374    }
1375
1376    async fn buffer_saved(
1377        self: Arc<Server>,
1378        request: TypedEnvelope<proto::BufferSaved>,
1379    ) -> Result<()> {
1380        let receiver_ids = self.store().await.project_connection_ids(
1381            ProjectId::from_proto(request.payload.project_id),
1382            request.sender_id,
1383        )?;
1384        broadcast(request.sender_id, receiver_ids, |connection_id| {
1385            self.peer
1386                .forward_send(request.sender_id, connection_id, request.payload.clone())
1387        });
1388        Ok(())
1389    }
1390
1391    async fn follow(
1392        self: Arc<Self>,
1393        request: TypedEnvelope<proto::Follow>,
1394        response: Response<proto::Follow>,
1395    ) -> Result<()> {
1396        let project_id = ProjectId::from_proto(request.payload.project_id);
1397        let leader_id = ConnectionId(request.payload.leader_id);
1398        let follower_id = request.sender_id;
1399        {
1400            let mut store = self.store().await;
1401            if !store
1402                .project_connection_ids(project_id, follower_id)?
1403                .contains(&leader_id)
1404            {
1405                Err(anyhow!("no such peer"))?;
1406            }
1407
1408            store.register_project_activity(project_id, follower_id)?;
1409        }
1410
1411        let mut response_payload = self
1412            .peer
1413            .forward_request(request.sender_id, leader_id, request.payload)
1414            .await?;
1415        response_payload
1416            .views
1417            .retain(|view| view.leader_id != Some(follower_id.0));
1418        response.send(response_payload)?;
1419        Ok(())
1420    }
1421
1422    async fn unfollow(self: Arc<Self>, request: TypedEnvelope<proto::Unfollow>) -> Result<()> {
1423        let project_id = ProjectId::from_proto(request.payload.project_id);
1424        let leader_id = ConnectionId(request.payload.leader_id);
1425        let mut store = self.store().await;
1426        if !store
1427            .project_connection_ids(project_id, request.sender_id)?
1428            .contains(&leader_id)
1429        {
1430            Err(anyhow!("no such peer"))?;
1431        }
1432        store.register_project_activity(project_id, request.sender_id)?;
1433        self.peer
1434            .forward_send(request.sender_id, leader_id, request.payload)?;
1435        Ok(())
1436    }
1437
1438    async fn update_followers(
1439        self: Arc<Self>,
1440        request: TypedEnvelope<proto::UpdateFollowers>,
1441    ) -> Result<()> {
1442        let project_id = ProjectId::from_proto(request.payload.project_id);
1443        let mut store = self.store().await;
1444        store.register_project_activity(project_id, request.sender_id)?;
1445        let connection_ids = store.project_connection_ids(project_id, request.sender_id)?;
1446        let leader_id = request
1447            .payload
1448            .variant
1449            .as_ref()
1450            .and_then(|variant| match variant {
1451                proto::update_followers::Variant::CreateView(payload) => payload.leader_id,
1452                proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
1453                proto::update_followers::Variant::UpdateActiveView(payload) => payload.leader_id,
1454            });
1455        for follower_id in &request.payload.follower_ids {
1456            let follower_id = ConnectionId(*follower_id);
1457            if connection_ids.contains(&follower_id) && Some(follower_id.0) != leader_id {
1458                self.peer
1459                    .forward_send(request.sender_id, follower_id, request.payload.clone())?;
1460            }
1461        }
1462        Ok(())
1463    }
1464
1465    async fn get_channels(
1466        self: Arc<Server>,
1467        request: TypedEnvelope<proto::GetChannels>,
1468        response: Response<proto::GetChannels>,
1469    ) -> Result<()> {
1470        let user_id = self
1471            .store()
1472            .await
1473            .user_id_for_connection(request.sender_id)?;
1474        let channels = self.app_state.db.get_accessible_channels(user_id).await?;
1475        response.send(proto::GetChannelsResponse {
1476            channels: channels
1477                .into_iter()
1478                .map(|chan| proto::Channel {
1479                    id: chan.id.to_proto(),
1480                    name: chan.name,
1481                })
1482                .collect(),
1483        })?;
1484        Ok(())
1485    }
1486
1487    async fn get_users(
1488        self: Arc<Server>,
1489        request: TypedEnvelope<proto::GetUsers>,
1490        response: Response<proto::GetUsers>,
1491    ) -> Result<()> {
1492        let user_ids = request
1493            .payload
1494            .user_ids
1495            .into_iter()
1496            .map(UserId::from_proto)
1497            .collect();
1498        let users = self
1499            .app_state
1500            .db
1501            .get_users_by_ids(user_ids)
1502            .await?
1503            .into_iter()
1504            .map(|user| proto::User {
1505                id: user.id.to_proto(),
1506                avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1507                github_login: user.github_login,
1508            })
1509            .collect();
1510        response.send(proto::UsersResponse { users })?;
1511        Ok(())
1512    }
1513
1514    async fn fuzzy_search_users(
1515        self: Arc<Server>,
1516        request: TypedEnvelope<proto::FuzzySearchUsers>,
1517        response: Response<proto::FuzzySearchUsers>,
1518    ) -> Result<()> {
1519        let user_id = self
1520            .store()
1521            .await
1522            .user_id_for_connection(request.sender_id)?;
1523        let query = request.payload.query;
1524        let db = &self.app_state.db;
1525        let users = match query.len() {
1526            0 => vec![],
1527            1 | 2 => db
1528                .get_user_by_github_login(&query)
1529                .await?
1530                .into_iter()
1531                .collect(),
1532            _ => db.fuzzy_search_users(&query, 10).await?,
1533        };
1534        let users = users
1535            .into_iter()
1536            .filter(|user| user.id != user_id)
1537            .map(|user| proto::User {
1538                id: user.id.to_proto(),
1539                avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1540                github_login: user.github_login,
1541            })
1542            .collect();
1543        response.send(proto::UsersResponse { users })?;
1544        Ok(())
1545    }
1546
1547    async fn request_contact(
1548        self: Arc<Server>,
1549        request: TypedEnvelope<proto::RequestContact>,
1550        response: Response<proto::RequestContact>,
1551    ) -> Result<()> {
1552        let requester_id = self
1553            .store()
1554            .await
1555            .user_id_for_connection(request.sender_id)?;
1556        let responder_id = UserId::from_proto(request.payload.responder_id);
1557        if requester_id == responder_id {
1558            return Err(anyhow!("cannot add yourself as a contact"))?;
1559        }
1560
1561        self.app_state
1562            .db
1563            .send_contact_request(requester_id, responder_id)
1564            .await?;
1565
1566        // Update outgoing contact requests of requester
1567        let mut update = proto::UpdateContacts::default();
1568        update.outgoing_requests.push(responder_id.to_proto());
1569        for connection_id in self.store().await.connection_ids_for_user(requester_id) {
1570            self.peer.send(connection_id, update.clone())?;
1571        }
1572
1573        // Update incoming contact requests of responder
1574        let mut update = proto::UpdateContacts::default();
1575        update
1576            .incoming_requests
1577            .push(proto::IncomingContactRequest {
1578                requester_id: requester_id.to_proto(),
1579                should_notify: true,
1580            });
1581        for connection_id in self.store().await.connection_ids_for_user(responder_id) {
1582            self.peer.send(connection_id, update.clone())?;
1583        }
1584
1585        response.send(proto::Ack {})?;
1586        Ok(())
1587    }
1588
1589    async fn respond_to_contact_request(
1590        self: Arc<Server>,
1591        request: TypedEnvelope<proto::RespondToContactRequest>,
1592        response: Response<proto::RespondToContactRequest>,
1593    ) -> Result<()> {
1594        let responder_id = self
1595            .store()
1596            .await
1597            .user_id_for_connection(request.sender_id)?;
1598        let requester_id = UserId::from_proto(request.payload.requester_id);
1599        if request.payload.response == proto::ContactRequestResponse::Dismiss as i32 {
1600            self.app_state
1601                .db
1602                .dismiss_contact_notification(responder_id, requester_id)
1603                .await?;
1604        } else {
1605            let accept = request.payload.response == proto::ContactRequestResponse::Accept as i32;
1606            self.app_state
1607                .db
1608                .respond_to_contact_request(responder_id, requester_id, accept)
1609                .await?;
1610
1611            let store = self.store().await;
1612            // Update responder with new contact
1613            let mut update = proto::UpdateContacts::default();
1614            if accept {
1615                update
1616                    .contacts
1617                    .push(store.contact_for_user(requester_id, false));
1618            }
1619            update
1620                .remove_incoming_requests
1621                .push(requester_id.to_proto());
1622            for connection_id in store.connection_ids_for_user(responder_id) {
1623                self.peer.send(connection_id, update.clone())?;
1624            }
1625
1626            // Update requester with new contact
1627            let mut update = proto::UpdateContacts::default();
1628            if accept {
1629                update
1630                    .contacts
1631                    .push(store.contact_for_user(responder_id, true));
1632            }
1633            update
1634                .remove_outgoing_requests
1635                .push(responder_id.to_proto());
1636            for connection_id in store.connection_ids_for_user(requester_id) {
1637                self.peer.send(connection_id, update.clone())?;
1638            }
1639        }
1640
1641        response.send(proto::Ack {})?;
1642        Ok(())
1643    }
1644
1645    async fn remove_contact(
1646        self: Arc<Server>,
1647        request: TypedEnvelope<proto::RemoveContact>,
1648        response: Response<proto::RemoveContact>,
1649    ) -> Result<()> {
1650        let requester_id = self
1651            .store()
1652            .await
1653            .user_id_for_connection(request.sender_id)?;
1654        let responder_id = UserId::from_proto(request.payload.user_id);
1655        self.app_state
1656            .db
1657            .remove_contact(requester_id, responder_id)
1658            .await?;
1659
1660        // Update outgoing contact requests of requester
1661        let mut update = proto::UpdateContacts::default();
1662        update
1663            .remove_outgoing_requests
1664            .push(responder_id.to_proto());
1665        for connection_id in self.store().await.connection_ids_for_user(requester_id) {
1666            self.peer.send(connection_id, update.clone())?;
1667        }
1668
1669        // Update incoming contact requests of responder
1670        let mut update = proto::UpdateContacts::default();
1671        update
1672            .remove_incoming_requests
1673            .push(requester_id.to_proto());
1674        for connection_id in self.store().await.connection_ids_for_user(responder_id) {
1675            self.peer.send(connection_id, update.clone())?;
1676        }
1677
1678        response.send(proto::Ack {})?;
1679        Ok(())
1680    }
1681
1682    async fn join_channel(
1683        self: Arc<Self>,
1684        request: TypedEnvelope<proto::JoinChannel>,
1685        response: Response<proto::JoinChannel>,
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            .join_channel(request.sender_id, channel_id);
1704        let messages = self
1705            .app_state
1706            .db
1707            .get_channel_messages(channel_id, MESSAGE_COUNT_PER_PAGE, None)
1708            .await?
1709            .into_iter()
1710            .map(|msg| proto::ChannelMessage {
1711                id: msg.id.to_proto(),
1712                body: msg.body,
1713                timestamp: msg.sent_at.unix_timestamp() as u64,
1714                sender_id: msg.sender_id.to_proto(),
1715                nonce: Some(msg.nonce.as_u128().into()),
1716            })
1717            .collect::<Vec<_>>();
1718        response.send(proto::JoinChannelResponse {
1719            done: messages.len() < MESSAGE_COUNT_PER_PAGE,
1720            messages,
1721        })?;
1722        Ok(())
1723    }
1724
1725    async fn leave_channel(
1726        self: Arc<Self>,
1727        request: TypedEnvelope<proto::LeaveChannel>,
1728    ) -> Result<()> {
1729        let user_id = self
1730            .store()
1731            .await
1732            .user_id_for_connection(request.sender_id)?;
1733        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1734        if !self
1735            .app_state
1736            .db
1737            .can_user_access_channel(user_id, channel_id)
1738            .await?
1739        {
1740            Err(anyhow!("access denied"))?;
1741        }
1742
1743        self.store()
1744            .await
1745            .leave_channel(request.sender_id, channel_id);
1746
1747        Ok(())
1748    }
1749
1750    async fn send_channel_message(
1751        self: Arc<Self>,
1752        request: TypedEnvelope<proto::SendChannelMessage>,
1753        response: Response<proto::SendChannelMessage>,
1754    ) -> Result<()> {
1755        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1756        let user_id;
1757        let connection_ids;
1758        {
1759            let state = self.store().await;
1760            user_id = state.user_id_for_connection(request.sender_id)?;
1761            connection_ids = state.channel_connection_ids(channel_id)?;
1762        }
1763
1764        // Validate the message body.
1765        let body = request.payload.body.trim().to_string();
1766        if body.len() > MAX_MESSAGE_LEN {
1767            return Err(anyhow!("message is too long"))?;
1768        }
1769        if body.is_empty() {
1770            return Err(anyhow!("message can't be blank"))?;
1771        }
1772
1773        let timestamp = OffsetDateTime::now_utc();
1774        let nonce = request
1775            .payload
1776            .nonce
1777            .ok_or_else(|| anyhow!("nonce can't be blank"))?;
1778
1779        let message_id = self
1780            .app_state
1781            .db
1782            .create_channel_message(channel_id, user_id, &body, timestamp, nonce.clone().into())
1783            .await?
1784            .to_proto();
1785        let message = proto::ChannelMessage {
1786            sender_id: user_id.to_proto(),
1787            id: message_id,
1788            body,
1789            timestamp: timestamp.unix_timestamp() as u64,
1790            nonce: Some(nonce),
1791        };
1792        broadcast(request.sender_id, connection_ids, |conn_id| {
1793            self.peer.send(
1794                conn_id,
1795                proto::ChannelMessageSent {
1796                    channel_id: channel_id.to_proto(),
1797                    message: Some(message.clone()),
1798                },
1799            )
1800        });
1801        response.send(proto::SendChannelMessageResponse {
1802            message: Some(message),
1803        })?;
1804        Ok(())
1805    }
1806
1807    async fn get_channel_messages(
1808        self: Arc<Self>,
1809        request: TypedEnvelope<proto::GetChannelMessages>,
1810        response: Response<proto::GetChannelMessages>,
1811    ) -> Result<()> {
1812        let user_id = self
1813            .store()
1814            .await
1815            .user_id_for_connection(request.sender_id)?;
1816        let channel_id = ChannelId::from_proto(request.payload.channel_id);
1817        if !self
1818            .app_state
1819            .db
1820            .can_user_access_channel(user_id, channel_id)
1821            .await?
1822        {
1823            Err(anyhow!("access denied"))?;
1824        }
1825
1826        let messages = self
1827            .app_state
1828            .db
1829            .get_channel_messages(
1830                channel_id,
1831                MESSAGE_COUNT_PER_PAGE,
1832                Some(MessageId::from_proto(request.payload.before_message_id)),
1833            )
1834            .await?
1835            .into_iter()
1836            .map(|msg| proto::ChannelMessage {
1837                id: msg.id.to_proto(),
1838                body: msg.body,
1839                timestamp: msg.sent_at.unix_timestamp() as u64,
1840                sender_id: msg.sender_id.to_proto(),
1841                nonce: Some(msg.nonce.as_u128().into()),
1842            })
1843            .collect::<Vec<_>>();
1844        response.send(proto::GetChannelMessagesResponse {
1845            done: messages.len() < MESSAGE_COUNT_PER_PAGE,
1846            messages,
1847        })?;
1848        Ok(())
1849    }
1850
1851    pub(crate) async fn store(&self) -> StoreGuard<'_> {
1852        #[cfg(test)]
1853        tokio::task::yield_now().await;
1854        let guard = self.store.lock().await;
1855        #[cfg(test)]
1856        tokio::task::yield_now().await;
1857        StoreGuard {
1858            guard,
1859            _not_send: PhantomData,
1860        }
1861    }
1862
1863    pub async fn snapshot<'a>(self: &'a Arc<Self>) -> ServerSnapshot<'a> {
1864        ServerSnapshot {
1865            store: self.store().await,
1866            peer: &self.peer,
1867        }
1868    }
1869}
1870
1871impl<'a> Deref for StoreGuard<'a> {
1872    type Target = Store;
1873
1874    fn deref(&self) -> &Self::Target {
1875        &*self.guard
1876    }
1877}
1878
1879impl<'a> DerefMut for StoreGuard<'a> {
1880    fn deref_mut(&mut self) -> &mut Self::Target {
1881        &mut *self.guard
1882    }
1883}
1884
1885impl<'a> Drop for StoreGuard<'a> {
1886    fn drop(&mut self) {
1887        #[cfg(test)]
1888        self.check_invariants();
1889    }
1890}
1891
1892impl Executor for RealExecutor {
1893    type Sleep = Sleep;
1894
1895    fn spawn_detached<F: 'static + Send + Future<Output = ()>>(&self, future: F) {
1896        tokio::task::spawn(future);
1897    }
1898
1899    fn sleep(&self, duration: Duration) -> Self::Sleep {
1900        tokio::time::sleep(duration)
1901    }
1902}
1903
1904fn broadcast<F>(
1905    sender_id: ConnectionId,
1906    receiver_ids: impl IntoIterator<Item = ConnectionId>,
1907    mut f: F,
1908) where
1909    F: FnMut(ConnectionId) -> anyhow::Result<()>,
1910{
1911    for receiver_id in receiver_ids {
1912        if receiver_id != sender_id {
1913            f(receiver_id).trace_err();
1914        }
1915    }
1916}
1917
1918lazy_static! {
1919    static ref ZED_PROTOCOL_VERSION: HeaderName = HeaderName::from_static("x-zed-protocol-version");
1920}
1921
1922pub struct ProtocolVersion(u32);
1923
1924impl Header for ProtocolVersion {
1925    fn name() -> &'static HeaderName {
1926        &ZED_PROTOCOL_VERSION
1927    }
1928
1929    fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1930    where
1931        Self: Sized,
1932        I: Iterator<Item = &'i axum::http::HeaderValue>,
1933    {
1934        let version = values
1935            .next()
1936            .ok_or_else(axum::headers::Error::invalid)?
1937            .to_str()
1938            .map_err(|_| axum::headers::Error::invalid())?
1939            .parse()
1940            .map_err(|_| axum::headers::Error::invalid())?;
1941        Ok(Self(version))
1942    }
1943
1944    fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1945        values.extend([self.0.to_string().parse().unwrap()]);
1946    }
1947}
1948
1949pub fn routes(server: Arc<Server>) -> Router<Body> {
1950    Router::new()
1951        .route("/rpc", get(handle_websocket_request))
1952        .layer(
1953            ServiceBuilder::new()
1954                .layer(Extension(server.app_state.clone()))
1955                .layer(middleware::from_fn(auth::validate_header)),
1956        )
1957        .route("/metrics", get(handle_metrics))
1958        .layer(Extension(server))
1959}
1960
1961pub async fn handle_websocket_request(
1962    TypedHeader(ProtocolVersion(protocol_version)): TypedHeader<ProtocolVersion>,
1963    ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
1964    Extension(server): Extension<Arc<Server>>,
1965    Extension(user): Extension<User>,
1966    ws: WebSocketUpgrade,
1967) -> axum::response::Response {
1968    if protocol_version != rpc::PROTOCOL_VERSION {
1969        return (
1970            StatusCode::UPGRADE_REQUIRED,
1971            "client must be upgraded".to_string(),
1972        )
1973            .into_response();
1974    }
1975    let socket_address = socket_address.to_string();
1976    ws.on_upgrade(move |socket| {
1977        use util::ResultExt;
1978        let socket = socket
1979            .map_ok(to_tungstenite_message)
1980            .err_into()
1981            .with(|message| async move { Ok(to_axum_message(message)) });
1982        let connection = Connection::new(Box::pin(socket));
1983        async move {
1984            server
1985                .handle_connection(connection, socket_address, user, None, RealExecutor)
1986                .await
1987                .log_err();
1988        }
1989    })
1990}
1991
1992pub async fn handle_metrics(Extension(server): Extension<Arc<Server>>) -> axum::response::Response {
1993    // We call `store_mut` here for its side effects of updating metrics.
1994    let metrics = server.store().await.metrics();
1995    METRIC_CONNECTIONS.set(metrics.connections as _);
1996    METRIC_REGISTERED_PROJECTS.set(metrics.registered_projects as _);
1997    METRIC_ACTIVE_PROJECTS.set(metrics.active_projects as _);
1998    METRIC_SHARED_PROJECTS.set(metrics.shared_projects as _);
1999
2000    let encoder = prometheus::TextEncoder::new();
2001    let metric_families = prometheus::gather();
2002    match encoder.encode_to_string(&metric_families) {
2003        Ok(string) => (StatusCode::OK, string).into_response(),
2004        Err(error) => (
2005            StatusCode::INTERNAL_SERVER_ERROR,
2006            format!("failed to encode metrics {:?}", error),
2007        )
2008            .into_response(),
2009    }
2010}
2011
2012fn to_axum_message(message: TungsteniteMessage) -> AxumMessage {
2013    match message {
2014        TungsteniteMessage::Text(payload) => AxumMessage::Text(payload),
2015        TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload),
2016        TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload),
2017        TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload),
2018        TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
2019            code: frame.code.into(),
2020            reason: frame.reason,
2021        })),
2022    }
2023}
2024
2025fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
2026    match message {
2027        AxumMessage::Text(payload) => TungsteniteMessage::Text(payload),
2028        AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload),
2029        AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload),
2030        AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload),
2031        AxumMessage::Close(frame) => {
2032            TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
2033                code: frame.code.into(),
2034                reason: frame.reason,
2035            }))
2036        }
2037    }
2038}
2039
2040pub trait ResultExt {
2041    type Ok;
2042
2043    fn trace_err(self) -> Option<Self::Ok>;
2044}
2045
2046impl<T, E> ResultExt for Result<T, E>
2047where
2048    E: std::fmt::Debug,
2049{
2050    type Ok = T;
2051
2052    fn trace_err(self) -> Option<T> {
2053        match self {
2054            Ok(value) => Some(value),
2055            Err(error) => {
2056                tracing::error!("{:?}", error);
2057                None
2058            }
2059        }
2060    }
2061}