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