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