rpc.rs

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