1mod connection_pool;
2
3use crate::{
4 auth,
5 db::{self, DefaultDb, 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<DefaultDb>);
132
133impl Deref for DbHandle {
134 type Target = DefaultDb;
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 as i32)
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.to_proto(),
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: worktree_id.to_proto(),
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.to_proto(),
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
1317 .retain(|collaborator| collaborator.connection_id != session.connection_id.0 as i32);
1318 let project_connection_ids = collaborators
1319 .into_iter()
1320 .map(|collaborator| ConnectionId(collaborator.connection_id as u32));
1321 broadcast(host_connection_id, project_connection_ids, |conn_id| {
1322 session
1323 .peer
1324 .forward_send(host_connection_id, conn_id, response_payload.clone())
1325 });
1326 response.send(response_payload)?;
1327 Ok(())
1328}
1329
1330async fn create_buffer_for_peer(
1331 request: proto::CreateBufferForPeer,
1332 session: Session,
1333) -> Result<()> {
1334 session.peer.forward_send(
1335 session.connection_id,
1336 ConnectionId(request.peer_id),
1337 request,
1338 )?;
1339 Ok(())
1340}
1341
1342async fn update_buffer(
1343 request: proto::UpdateBuffer,
1344 response: Response<proto::UpdateBuffer>,
1345 session: Session,
1346) -> Result<()> {
1347 let project_id = ProjectId::from_proto(request.project_id);
1348 let project_connection_ids = session
1349 .db()
1350 .await
1351 .project_connection_ids(project_id, session.connection_id)
1352 .await?;
1353
1354 broadcast(
1355 session.connection_id,
1356 project_connection_ids,
1357 |connection_id| {
1358 session
1359 .peer
1360 .forward_send(session.connection_id, connection_id, request.clone())
1361 },
1362 );
1363 response.send(proto::Ack {})?;
1364 Ok(())
1365}
1366
1367async fn update_buffer_file(request: proto::UpdateBufferFile, session: Session) -> Result<()> {
1368 let project_id = ProjectId::from_proto(request.project_id);
1369 let project_connection_ids = session
1370 .db()
1371 .await
1372 .project_connection_ids(project_id, session.connection_id)
1373 .await?;
1374
1375 broadcast(
1376 session.connection_id,
1377 project_connection_ids,
1378 |connection_id| {
1379 session
1380 .peer
1381 .forward_send(session.connection_id, connection_id, request.clone())
1382 },
1383 );
1384 Ok(())
1385}
1386
1387async fn buffer_reloaded(request: proto::BufferReloaded, session: Session) -> Result<()> {
1388 let project_id = ProjectId::from_proto(request.project_id);
1389 let project_connection_ids = session
1390 .db()
1391 .await
1392 .project_connection_ids(project_id, session.connection_id)
1393 .await?;
1394 broadcast(
1395 session.connection_id,
1396 project_connection_ids,
1397 |connection_id| {
1398 session
1399 .peer
1400 .forward_send(session.connection_id, connection_id, request.clone())
1401 },
1402 );
1403 Ok(())
1404}
1405
1406async fn buffer_saved(request: proto::BufferSaved, session: Session) -> Result<()> {
1407 let project_id = ProjectId::from_proto(request.project_id);
1408 let project_connection_ids = session
1409 .db()
1410 .await
1411 .project_connection_ids(project_id, session.connection_id)
1412 .await?;
1413 broadcast(
1414 session.connection_id,
1415 project_connection_ids,
1416 |connection_id| {
1417 session
1418 .peer
1419 .forward_send(session.connection_id, connection_id, request.clone())
1420 },
1421 );
1422 Ok(())
1423}
1424
1425async fn follow(
1426 request: proto::Follow,
1427 response: Response<proto::Follow>,
1428 session: Session,
1429) -> Result<()> {
1430 let project_id = ProjectId::from_proto(request.project_id);
1431 let leader_id = ConnectionId(request.leader_id);
1432 let follower_id = session.connection_id;
1433 let project_connection_ids = session
1434 .db()
1435 .await
1436 .project_connection_ids(project_id, session.connection_id)
1437 .await?;
1438
1439 if !project_connection_ids.contains(&leader_id) {
1440 Err(anyhow!("no such peer"))?;
1441 }
1442
1443 let mut response_payload = session
1444 .peer
1445 .forward_request(session.connection_id, leader_id, request)
1446 .await?;
1447 response_payload
1448 .views
1449 .retain(|view| view.leader_id != Some(follower_id.0));
1450 response.send(response_payload)?;
1451 Ok(())
1452}
1453
1454async fn unfollow(request: proto::Unfollow, session: Session) -> Result<()> {
1455 let project_id = ProjectId::from_proto(request.project_id);
1456 let leader_id = ConnectionId(request.leader_id);
1457 let project_connection_ids = session
1458 .db()
1459 .await
1460 .project_connection_ids(project_id, session.connection_id)
1461 .await?;
1462 if !project_connection_ids.contains(&leader_id) {
1463 Err(anyhow!("no such peer"))?;
1464 }
1465 session
1466 .peer
1467 .forward_send(session.connection_id, leader_id, request)?;
1468 Ok(())
1469}
1470
1471async fn update_followers(request: proto::UpdateFollowers, session: Session) -> Result<()> {
1472 let project_id = ProjectId::from_proto(request.project_id);
1473 let project_connection_ids = session
1474 .db
1475 .lock()
1476 .await
1477 .project_connection_ids(project_id, session.connection_id)
1478 .await?;
1479
1480 let leader_id = request.variant.as_ref().and_then(|variant| match variant {
1481 proto::update_followers::Variant::CreateView(payload) => payload.leader_id,
1482 proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
1483 proto::update_followers::Variant::UpdateActiveView(payload) => payload.leader_id,
1484 });
1485 for follower_id in &request.follower_ids {
1486 let follower_id = ConnectionId(*follower_id);
1487 if project_connection_ids.contains(&follower_id) && Some(follower_id.0) != leader_id {
1488 session
1489 .peer
1490 .forward_send(session.connection_id, follower_id, request.clone())?;
1491 }
1492 }
1493 Ok(())
1494}
1495
1496async fn get_users(
1497 request: proto::GetUsers,
1498 response: Response<proto::GetUsers>,
1499 session: Session,
1500) -> Result<()> {
1501 let user_ids = request
1502 .user_ids
1503 .into_iter()
1504 .map(UserId::from_proto)
1505 .collect();
1506 let users = session
1507 .db()
1508 .await
1509 .get_users_by_ids(user_ids)
1510 .await?
1511 .into_iter()
1512 .map(|user| proto::User {
1513 id: user.id.to_proto(),
1514 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1515 github_login: user.github_login,
1516 })
1517 .collect();
1518 response.send(proto::UsersResponse { users })?;
1519 Ok(())
1520}
1521
1522async fn fuzzy_search_users(
1523 request: proto::FuzzySearchUsers,
1524 response: Response<proto::FuzzySearchUsers>,
1525 session: Session,
1526) -> Result<()> {
1527 let query = request.query;
1528 let users = match query.len() {
1529 0 => vec![],
1530 1 | 2 => session
1531 .db()
1532 .await
1533 .get_user_by_github_account(&query, None)
1534 .await?
1535 .into_iter()
1536 .collect(),
1537 _ => session.db().await.fuzzy_search_users(&query, 10).await?,
1538 };
1539 let users = users
1540 .into_iter()
1541 .filter(|user| user.id != session.user_id)
1542 .map(|user| proto::User {
1543 id: user.id.to_proto(),
1544 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1545 github_login: user.github_login,
1546 })
1547 .collect();
1548 response.send(proto::UsersResponse { users })?;
1549 Ok(())
1550}
1551
1552async fn request_contact(
1553 request: proto::RequestContact,
1554 response: Response<proto::RequestContact>,
1555 session: Session,
1556) -> Result<()> {
1557 let requester_id = session.user_id;
1558 let responder_id = UserId::from_proto(request.responder_id);
1559 if requester_id == responder_id {
1560 return Err(anyhow!("cannot add yourself as a contact"))?;
1561 }
1562
1563 session
1564 .db()
1565 .await
1566 .send_contact_request(requester_id, responder_id)
1567 .await?;
1568
1569 // Update outgoing contact requests of requester
1570 let mut update = proto::UpdateContacts::default();
1571 update.outgoing_requests.push(responder_id.to_proto());
1572 for connection_id in session
1573 .connection_pool()
1574 .await
1575 .user_connection_ids(requester_id)
1576 {
1577 session.peer.send(connection_id, update.clone())?;
1578 }
1579
1580 // Update incoming contact requests of responder
1581 let mut update = proto::UpdateContacts::default();
1582 update
1583 .incoming_requests
1584 .push(proto::IncomingContactRequest {
1585 requester_id: requester_id.to_proto(),
1586 should_notify: true,
1587 });
1588 for connection_id in session
1589 .connection_pool()
1590 .await
1591 .user_connection_ids(responder_id)
1592 {
1593 session.peer.send(connection_id, update.clone())?;
1594 }
1595
1596 response.send(proto::Ack {})?;
1597 Ok(())
1598}
1599
1600async fn respond_to_contact_request(
1601 request: proto::RespondToContactRequest,
1602 response: Response<proto::RespondToContactRequest>,
1603 session: Session,
1604) -> Result<()> {
1605 let responder_id = session.user_id;
1606 let requester_id = UserId::from_proto(request.requester_id);
1607 let db = session.db().await;
1608 if request.response == proto::ContactRequestResponse::Dismiss as i32 {
1609 db.dismiss_contact_notification(responder_id, requester_id)
1610 .await?;
1611 } else {
1612 let accept = request.response == proto::ContactRequestResponse::Accept as i32;
1613
1614 db.respond_to_contact_request(responder_id, requester_id, accept)
1615 .await?;
1616 let busy = db.is_user_busy(requester_id).await?;
1617
1618 let pool = session.connection_pool().await;
1619 // Update responder with new contact
1620 let mut update = proto::UpdateContacts::default();
1621 if accept {
1622 update
1623 .contacts
1624 .push(contact_for_user(requester_id, false, busy, &pool));
1625 }
1626 update
1627 .remove_incoming_requests
1628 .push(requester_id.to_proto());
1629 for connection_id in pool.user_connection_ids(responder_id) {
1630 session.peer.send(connection_id, update.clone())?;
1631 }
1632
1633 // Update requester with new contact
1634 let mut update = proto::UpdateContacts::default();
1635 if accept {
1636 update
1637 .contacts
1638 .push(contact_for_user(responder_id, true, busy, &pool));
1639 }
1640 update
1641 .remove_outgoing_requests
1642 .push(responder_id.to_proto());
1643 for connection_id in pool.user_connection_ids(requester_id) {
1644 session.peer.send(connection_id, update.clone())?;
1645 }
1646 }
1647
1648 response.send(proto::Ack {})?;
1649 Ok(())
1650}
1651
1652async fn remove_contact(
1653 request: proto::RemoveContact,
1654 response: Response<proto::RemoveContact>,
1655 session: Session,
1656) -> Result<()> {
1657 let requester_id = session.user_id;
1658 let responder_id = UserId::from_proto(request.user_id);
1659 let db = session.db().await;
1660 db.remove_contact(requester_id, responder_id).await?;
1661
1662 let pool = session.connection_pool().await;
1663 // Update outgoing contact requests of requester
1664 let mut update = proto::UpdateContacts::default();
1665 update
1666 .remove_outgoing_requests
1667 .push(responder_id.to_proto());
1668 for connection_id in pool.user_connection_ids(requester_id) {
1669 session.peer.send(connection_id, update.clone())?;
1670 }
1671
1672 // Update incoming contact requests of responder
1673 let mut update = proto::UpdateContacts::default();
1674 update
1675 .remove_incoming_requests
1676 .push(requester_id.to_proto());
1677 for connection_id in pool.user_connection_ids(responder_id) {
1678 session.peer.send(connection_id, update.clone())?;
1679 }
1680
1681 response.send(proto::Ack {})?;
1682 Ok(())
1683}
1684
1685async fn update_diff_base(request: proto::UpdateDiffBase, session: Session) -> Result<()> {
1686 let project_id = ProjectId::from_proto(request.project_id);
1687 let project_connection_ids = session
1688 .db()
1689 .await
1690 .project_connection_ids(project_id, session.connection_id)
1691 .await?;
1692 broadcast(
1693 session.connection_id,
1694 project_connection_ids,
1695 |connection_id| {
1696 session
1697 .peer
1698 .forward_send(session.connection_id, connection_id, request.clone())
1699 },
1700 );
1701 Ok(())
1702}
1703
1704async fn get_private_user_info(
1705 _request: proto::GetPrivateUserInfo,
1706 response: Response<proto::GetPrivateUserInfo>,
1707 session: Session,
1708) -> Result<()> {
1709 let metrics_id = session
1710 .db()
1711 .await
1712 .get_user_metrics_id(session.user_id)
1713 .await?;
1714 let user = session
1715 .db()
1716 .await
1717 .get_user_by_id(session.user_id)
1718 .await?
1719 .ok_or_else(|| anyhow!("user not found"))?;
1720 response.send(proto::GetPrivateUserInfoResponse {
1721 metrics_id,
1722 staff: user.admin,
1723 })?;
1724 Ok(())
1725}
1726
1727fn to_axum_message(message: TungsteniteMessage) -> AxumMessage {
1728 match message {
1729 TungsteniteMessage::Text(payload) => AxumMessage::Text(payload),
1730 TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload),
1731 TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload),
1732 TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload),
1733 TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
1734 code: frame.code.into(),
1735 reason: frame.reason,
1736 })),
1737 }
1738}
1739
1740fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
1741 match message {
1742 AxumMessage::Text(payload) => TungsteniteMessage::Text(payload),
1743 AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload),
1744 AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload),
1745 AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload),
1746 AxumMessage::Close(frame) => {
1747 TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
1748 code: frame.code.into(),
1749 reason: frame.reason,
1750 }))
1751 }
1752 }
1753}
1754
1755fn build_initial_contacts_update(
1756 contacts: Vec<db::Contact>,
1757 pool: &ConnectionPool,
1758) -> proto::UpdateContacts {
1759 let mut update = proto::UpdateContacts::default();
1760
1761 for contact in contacts {
1762 match contact {
1763 db::Contact::Accepted {
1764 user_id,
1765 should_notify,
1766 busy,
1767 } => {
1768 update
1769 .contacts
1770 .push(contact_for_user(user_id, should_notify, busy, &pool));
1771 }
1772 db::Contact::Outgoing { user_id } => update.outgoing_requests.push(user_id.to_proto()),
1773 db::Contact::Incoming {
1774 user_id,
1775 should_notify,
1776 } => update
1777 .incoming_requests
1778 .push(proto::IncomingContactRequest {
1779 requester_id: user_id.to_proto(),
1780 should_notify,
1781 }),
1782 }
1783 }
1784
1785 update
1786}
1787
1788fn contact_for_user(
1789 user_id: UserId,
1790 should_notify: bool,
1791 busy: bool,
1792 pool: &ConnectionPool,
1793) -> proto::Contact {
1794 proto::Contact {
1795 user_id: user_id.to_proto(),
1796 online: pool.is_user_online(user_id),
1797 busy,
1798 should_notify,
1799 }
1800}
1801
1802fn room_updated(room: &proto::Room, session: &Session) {
1803 for participant in &room.participants {
1804 session
1805 .peer
1806 .send(
1807 ConnectionId(participant.peer_id),
1808 proto::RoomUpdated {
1809 room: Some(room.clone()),
1810 },
1811 )
1812 .trace_err();
1813 }
1814}
1815
1816async fn update_user_contacts(user_id: UserId, session: &Session) -> Result<()> {
1817 let db = session.db().await;
1818 let contacts = db.get_contacts(user_id).await?;
1819 let busy = db.is_user_busy(user_id).await?;
1820
1821 let pool = session.connection_pool().await;
1822 let updated_contact = contact_for_user(user_id, false, busy, &pool);
1823 for contact in contacts {
1824 if let db::Contact::Accepted {
1825 user_id: contact_user_id,
1826 ..
1827 } = contact
1828 {
1829 for contact_conn_id in pool.user_connection_ids(contact_user_id) {
1830 session
1831 .peer
1832 .send(
1833 contact_conn_id,
1834 proto::UpdateContacts {
1835 contacts: vec![updated_contact.clone()],
1836 remove_contacts: Default::default(),
1837 incoming_requests: Default::default(),
1838 remove_incoming_requests: Default::default(),
1839 outgoing_requests: Default::default(),
1840 remove_outgoing_requests: Default::default(),
1841 },
1842 )
1843 .trace_err();
1844 }
1845 }
1846 }
1847 Ok(())
1848}
1849
1850async fn leave_room_for_session(session: &Session) -> Result<()> {
1851 let mut contacts_to_update = HashSet::default();
1852
1853 let canceled_calls_to_user_ids;
1854 let live_kit_room;
1855 let delete_live_kit_room;
1856 {
1857 let Some(mut left_room) = session.db().await.leave_room(session.connection_id).await? else {
1858 return Err(anyhow!("no room to leave"))?;
1859 };
1860 contacts_to_update.insert(session.user_id);
1861
1862 for project in left_room.left_projects.values() {
1863 for connection_id in &project.connection_ids {
1864 if project.host_user_id == session.user_id {
1865 session
1866 .peer
1867 .send(
1868 *connection_id,
1869 proto::UnshareProject {
1870 project_id: project.id.to_proto(),
1871 },
1872 )
1873 .trace_err();
1874 } else {
1875 session
1876 .peer
1877 .send(
1878 *connection_id,
1879 proto::RemoveProjectCollaborator {
1880 project_id: project.id.to_proto(),
1881 peer_id: session.connection_id.0,
1882 },
1883 )
1884 .trace_err();
1885 }
1886 }
1887
1888 session
1889 .peer
1890 .send(
1891 session.connection_id,
1892 proto::UnshareProject {
1893 project_id: project.id.to_proto(),
1894 },
1895 )
1896 .trace_err();
1897 }
1898
1899 room_updated(&left_room.room, &session);
1900 canceled_calls_to_user_ids = mem::take(&mut left_room.canceled_calls_to_user_ids);
1901 live_kit_room = mem::take(&mut left_room.room.live_kit_room);
1902 delete_live_kit_room = left_room.room.participants.is_empty();
1903 }
1904
1905 {
1906 let pool = session.connection_pool().await;
1907 for canceled_user_id in canceled_calls_to_user_ids {
1908 for connection_id in pool.user_connection_ids(canceled_user_id) {
1909 session
1910 .peer
1911 .send(connection_id, proto::CallCanceled {})
1912 .trace_err();
1913 }
1914 contacts_to_update.insert(canceled_user_id);
1915 }
1916 }
1917
1918 for contact_user_id in contacts_to_update {
1919 update_user_contacts(contact_user_id, &session).await?;
1920 }
1921
1922 if let Some(live_kit) = session.live_kit_client.as_ref() {
1923 live_kit
1924 .remove_participant(live_kit_room.clone(), session.connection_id.to_string())
1925 .await
1926 .trace_err();
1927
1928 if delete_live_kit_room {
1929 live_kit.delete_room(live_kit_room).await.trace_err();
1930 }
1931 }
1932
1933 Ok(())
1934}
1935
1936pub trait ResultExt {
1937 type Ok;
1938
1939 fn trace_err(self) -> Option<Self::Ok>;
1940}
1941
1942impl<T, E> ResultExt for Result<T, E>
1943where
1944 E: std::fmt::Debug,
1945{
1946 type Ok = T;
1947
1948 fn trace_err(self) -> Option<T> {
1949 match self {
1950 Ok(value) => Some(value),
1951 Err(error) => {
1952 tracing::error!("{:?}", error);
1953 None
1954 }
1955 }
1956 }
1957}