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