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