1mod connection_pool;
2
3use crate::{
4 auth,
5 db::{
6 self, BufferId, ChannelId, ChannelRole, ChannelsForUser, CreateChannelResult,
7 CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId,
8 MoveChannelResult, NotificationId, ProjectId, RemoveChannelMemberResult,
9 RenameChannelResult, RespondToChannelInvite, RoomId, ServerId, SetChannelVisibilityResult,
10 User, UserId,
11 },
12 executor::Executor,
13 AppState, Result,
14};
15use anyhow::anyhow;
16use async_tungstenite::tungstenite::{
17 protocol::CloseFrame as TungsteniteCloseFrame, Message as TungsteniteMessage,
18};
19use axum::{
20 body::Body,
21 extract::{
22 ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage},
23 ConnectInfo, WebSocketUpgrade,
24 },
25 headers::{Header, HeaderName},
26 http::StatusCode,
27 middleware,
28 response::IntoResponse,
29 routing::get,
30 Extension, Router, TypedHeader,
31};
32use collections::{HashMap, HashSet};
33pub use connection_pool::ConnectionPool;
34use futures::{
35 channel::oneshot,
36 future::{self, BoxFuture},
37 stream::FuturesUnordered,
38 FutureExt, SinkExt, StreamExt, TryStreamExt,
39};
40use lazy_static::lazy_static;
41use prometheus::{register_int_gauge, IntGauge};
42use rpc::{
43 proto::{
44 self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo,
45 RequestMessage, ShareProject, UpdateChannelBufferCollaborators,
46 },
47 Connection, ConnectionId, Peer, Receipt, TypedEnvelope,
48};
49use serde::{Serialize, Serializer};
50use std::{
51 any::TypeId,
52 fmt,
53 future::Future,
54 marker::PhantomData,
55 mem,
56 net::SocketAddr,
57 ops::{Deref, DerefMut},
58 rc::Rc,
59 sync::{
60 atomic::{AtomicBool, Ordering::SeqCst},
61 Arc,
62 },
63 time::{Duration, Instant},
64};
65use time::OffsetDateTime;
66use tokio::sync::{watch, Semaphore};
67use tower::ServiceBuilder;
68use tracing::{info_span, instrument, Instrument};
69use util::channel::RELEASE_CHANNEL_NAME;
70
71pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30);
72pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(10);
73
74const MESSAGE_COUNT_PER_PAGE: usize = 100;
75const MAX_MESSAGE_LEN: usize = 1024;
76const NOTIFICATION_COUNT_PER_PAGE: usize = 50;
77
78lazy_static! {
79 static ref METRIC_CONNECTIONS: IntGauge =
80 register_int_gauge!("connections", "number of connections").unwrap();
81 static ref METRIC_SHARED_PROJECTS: IntGauge = register_int_gauge!(
82 "shared_projects",
83 "number of open projects with one or more guests"
84 )
85 .unwrap();
86}
87
88type MessageHandler =
89 Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>;
90
91struct Response<R> {
92 peer: Arc<Peer>,
93 receipt: Receipt<R>,
94 responded: Arc<AtomicBool>,
95}
96
97impl<R: RequestMessage> Response<R> {
98 fn send(self, payload: R::Response) -> Result<()> {
99 self.responded.store(true, SeqCst);
100 self.peer.respond(self.receipt, payload)?;
101 Ok(())
102 }
103}
104
105#[derive(Clone)]
106struct Session {
107 user_id: UserId,
108 connection_id: ConnectionId,
109 db: Arc<tokio::sync::Mutex<DbHandle>>,
110 peer: Arc<Peer>,
111 connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
112 live_kit_client: Option<Arc<dyn live_kit_server::api::Client>>,
113 _executor: Executor,
114}
115
116impl Session {
117 async fn db(&self) -> tokio::sync::MutexGuard<DbHandle> {
118 #[cfg(test)]
119 tokio::task::yield_now().await;
120 let guard = self.db.lock().await;
121 #[cfg(test)]
122 tokio::task::yield_now().await;
123 guard
124 }
125
126 async fn connection_pool(&self) -> ConnectionPoolGuard<'_> {
127 #[cfg(test)]
128 tokio::task::yield_now().await;
129 let guard = self.connection_pool.lock();
130 ConnectionPoolGuard {
131 guard,
132 _not_send: PhantomData,
133 }
134 }
135}
136
137impl fmt::Debug for Session {
138 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139 f.debug_struct("Session")
140 .field("user_id", &self.user_id)
141 .field("connection_id", &self.connection_id)
142 .finish()
143 }
144}
145
146struct DbHandle(Arc<Database>);
147
148impl Deref for DbHandle {
149 type Target = Database;
150
151 fn deref(&self) -> &Self::Target {
152 self.0.as_ref()
153 }
154}
155
156pub struct Server {
157 id: parking_lot::Mutex<ServerId>,
158 peer: Arc<Peer>,
159 pub(crate) connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
160 app_state: Arc<AppState>,
161 executor: Executor,
162 handlers: HashMap<TypeId, MessageHandler>,
163 teardown: watch::Sender<()>,
164}
165
166pub(crate) struct ConnectionPoolGuard<'a> {
167 guard: parking_lot::MutexGuard<'a, ConnectionPool>,
168 _not_send: PhantomData<Rc<()>>,
169}
170
171#[derive(Serialize)]
172pub struct ServerSnapshot<'a> {
173 peer: &'a Peer,
174 #[serde(serialize_with = "serialize_deref")]
175 connection_pool: ConnectionPoolGuard<'a>,
176}
177
178pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
179where
180 S: Serializer,
181 T: Deref<Target = U>,
182 U: Serialize,
183{
184 Serialize::serialize(value.deref(), serializer)
185}
186
187impl Server {
188 pub fn new(id: ServerId, app_state: Arc<AppState>, executor: Executor) -> Arc<Self> {
189 let mut server = Self {
190 id: parking_lot::Mutex::new(id),
191 peer: Peer::new(id.0 as u32),
192 app_state,
193 executor,
194 connection_pool: Default::default(),
195 handlers: Default::default(),
196 teardown: watch::channel(()).0,
197 };
198
199 server
200 .add_request_handler(ping)
201 .add_request_handler(create_room)
202 .add_request_handler(join_room)
203 .add_request_handler(rejoin_room)
204 .add_request_handler(leave_room)
205 .add_request_handler(call)
206 .add_request_handler(cancel_call)
207 .add_message_handler(decline_call)
208 .add_request_handler(update_participant_location)
209 .add_request_handler(share_project)
210 .add_message_handler(unshare_project)
211 .add_request_handler(join_project)
212 .add_message_handler(leave_project)
213 .add_request_handler(update_project)
214 .add_request_handler(update_worktree)
215 .add_message_handler(start_language_server)
216 .add_message_handler(update_language_server)
217 .add_message_handler(update_diagnostic_summary)
218 .add_message_handler(update_worktree_settings)
219 .add_request_handler(forward_read_only_project_request::<proto::GetHover>)
220 .add_request_handler(forward_read_only_project_request::<proto::GetDefinition>)
221 .add_request_handler(forward_read_only_project_request::<proto::GetTypeDefinition>)
222 .add_request_handler(forward_read_only_project_request::<proto::GetReferences>)
223 .add_request_handler(forward_read_only_project_request::<proto::SearchProject>)
224 .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>)
225 .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>)
226 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)
227 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)
228 .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)
229 .add_request_handler(forward_read_only_project_request::<proto::InlayHints>)
230 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)
231 .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>)
232 .add_request_handler(
233 forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>,
234 )
235 .add_request_handler(
236 forward_mutating_project_request::<proto::ResolveCompletionDocumentation>,
237 )
238 .add_request_handler(forward_mutating_project_request::<proto::GetCodeActions>)
239 .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>)
240 .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>)
241 .add_request_handler(forward_mutating_project_request::<proto::PerformRename>)
242 .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>)
243 .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>)
244 .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>)
245 .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>)
246 .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>)
247 .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>)
248 .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>)
249 .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>)
250 .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)
251 .add_message_handler(create_buffer_for_peer)
252 .add_request_handler(update_buffer)
253 .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>)
254 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>)
255 .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>)
256 .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>)
257 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBase>)
258 .add_request_handler(get_users)
259 .add_request_handler(fuzzy_search_users)
260 .add_request_handler(request_contact)
261 .add_request_handler(remove_contact)
262 .add_request_handler(respond_to_contact_request)
263 .add_request_handler(create_channel)
264 .add_request_handler(delete_channel)
265 .add_request_handler(invite_channel_member)
266 .add_request_handler(remove_channel_member)
267 .add_request_handler(set_channel_member_role)
268 .add_request_handler(set_channel_visibility)
269 .add_request_handler(rename_channel)
270 .add_request_handler(join_channel_buffer)
271 .add_request_handler(leave_channel_buffer)
272 .add_message_handler(update_channel_buffer)
273 .add_request_handler(rejoin_channel_buffers)
274 .add_request_handler(get_channel_members)
275 .add_request_handler(respond_to_channel_invite)
276 .add_request_handler(join_channel)
277 .add_request_handler(join_channel_chat)
278 .add_message_handler(leave_channel_chat)
279 .add_request_handler(send_channel_message)
280 .add_request_handler(remove_channel_message)
281 .add_request_handler(get_channel_messages)
282 .add_request_handler(get_channel_messages_by_id)
283 .add_request_handler(get_notifications)
284 .add_request_handler(mark_notification_as_read)
285 .add_request_handler(move_channel)
286 .add_request_handler(follow)
287 .add_message_handler(unfollow)
288 .add_message_handler(update_followers)
289 .add_request_handler(get_private_user_info)
290 .add_message_handler(acknowledge_channel_message)
291 .add_message_handler(acknowledge_buffer_version);
292
293 Arc::new(server)
294 }
295
296 pub async fn start(&self) -> Result<()> {
297 let server_id = *self.id.lock();
298 let app_state = self.app_state.clone();
299 let peer = self.peer.clone();
300 let timeout = self.executor.sleep(CLEANUP_TIMEOUT);
301 let pool = self.connection_pool.clone();
302 let live_kit_client = self.app_state.live_kit_client.clone();
303
304 let span = info_span!("start server");
305 self.executor.spawn_detached(
306 async move {
307 tracing::info!("waiting for cleanup timeout");
308 timeout.await;
309 tracing::info!("cleanup timeout expired, retrieving stale rooms");
310 if let Some((room_ids, channel_ids)) = app_state
311 .db
312 .stale_server_resource_ids(&app_state.config.zed_environment, server_id)
313 .await
314 .trace_err()
315 {
316 tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms");
317 tracing::info!(
318 stale_channel_buffer_count = channel_ids.len(),
319 "retrieved stale channel buffers"
320 );
321
322 for channel_id in channel_ids {
323 if let Some(refreshed_channel_buffer) = app_state
324 .db
325 .clear_stale_channel_buffer_collaborators(channel_id, server_id)
326 .await
327 .trace_err()
328 {
329 for connection_id in refreshed_channel_buffer.connection_ids {
330 peer.send(
331 connection_id,
332 proto::UpdateChannelBufferCollaborators {
333 channel_id: channel_id.to_proto(),
334 collaborators: refreshed_channel_buffer
335 .collaborators
336 .clone(),
337 },
338 )
339 .trace_err();
340 }
341 }
342 }
343
344 for room_id in room_ids {
345 let mut contacts_to_update = HashSet::default();
346 let mut canceled_calls_to_user_ids = Vec::new();
347 let mut live_kit_room = String::new();
348 let mut delete_live_kit_room = false;
349
350 if let Some(mut refreshed_room) = app_state
351 .db
352 .clear_stale_room_participants(room_id, server_id)
353 .await
354 .trace_err()
355 {
356 tracing::info!(
357 room_id = room_id.0,
358 new_participant_count = refreshed_room.room.participants.len(),
359 "refreshed room"
360 );
361 room_updated(&refreshed_room.room, &peer);
362 if let Some(channel_id) = refreshed_room.channel_id {
363 channel_updated(
364 channel_id,
365 &refreshed_room.room,
366 &refreshed_room.channel_members,
367 &peer,
368 &*pool.lock(),
369 );
370 }
371 contacts_to_update
372 .extend(refreshed_room.stale_participant_user_ids.iter().copied());
373 contacts_to_update
374 .extend(refreshed_room.canceled_calls_to_user_ids.iter().copied());
375 canceled_calls_to_user_ids =
376 mem::take(&mut refreshed_room.canceled_calls_to_user_ids);
377 live_kit_room = mem::take(&mut refreshed_room.room.live_kit_room);
378 delete_live_kit_room = refreshed_room.room.participants.is_empty();
379 }
380
381 {
382 let pool = pool.lock();
383 for canceled_user_id in canceled_calls_to_user_ids {
384 for connection_id in pool.user_connection_ids(canceled_user_id) {
385 peer.send(
386 connection_id,
387 proto::CallCanceled {
388 room_id: room_id.to_proto(),
389 },
390 )
391 .trace_err();
392 }
393 }
394 }
395
396 for user_id in contacts_to_update {
397 let busy = app_state.db.is_user_busy(user_id).await.trace_err();
398 let contacts = app_state.db.get_contacts(user_id).await.trace_err();
399 if let Some((busy, contacts)) = busy.zip(contacts) {
400 let pool = pool.lock();
401 let updated_contact = contact_for_user(user_id, busy, &pool);
402 for contact in contacts {
403 if let db::Contact::Accepted {
404 user_id: contact_user_id,
405 ..
406 } = contact
407 {
408 for contact_conn_id in
409 pool.user_connection_ids(contact_user_id)
410 {
411 peer.send(
412 contact_conn_id,
413 proto::UpdateContacts {
414 contacts: vec![updated_contact.clone()],
415 remove_contacts: Default::default(),
416 incoming_requests: Default::default(),
417 remove_incoming_requests: Default::default(),
418 outgoing_requests: Default::default(),
419 remove_outgoing_requests: Default::default(),
420 },
421 )
422 .trace_err();
423 }
424 }
425 }
426 }
427 }
428
429 if let Some(live_kit) = live_kit_client.as_ref() {
430 if delete_live_kit_room {
431 live_kit.delete_room(live_kit_room).await.trace_err();
432 }
433 }
434 }
435 }
436
437 app_state
438 .db
439 .delete_stale_servers(&app_state.config.zed_environment, server_id)
440 .await
441 .trace_err();
442 }
443 .instrument(span),
444 );
445 Ok(())
446 }
447
448 pub fn teardown(&self) {
449 self.peer.teardown();
450 self.connection_pool.lock().reset();
451 let _ = self.teardown.send(());
452 }
453
454 #[cfg(test)]
455 pub fn reset(&self, id: ServerId) {
456 self.teardown();
457 *self.id.lock() = id;
458 self.peer.reset(id.0 as u32);
459 }
460
461 #[cfg(test)]
462 pub fn id(&self) -> ServerId {
463 *self.id.lock()
464 }
465
466 fn add_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
467 where
468 F: 'static + Send + Sync + Fn(TypedEnvelope<M>, Session) -> Fut,
469 Fut: 'static + Send + Future<Output = Result<()>>,
470 M: EnvelopedMessage,
471 {
472 let prev_handler = self.handlers.insert(
473 TypeId::of::<M>(),
474 Box::new(move |envelope, session| {
475 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
476 let span = info_span!(
477 "handle message",
478 payload_type = envelope.payload_type_name()
479 );
480 span.in_scope(|| {
481 tracing::info!(
482 payload_type = envelope.payload_type_name(),
483 "message received"
484 );
485 });
486 let start_time = Instant::now();
487 let future = (handler)(*envelope, session);
488 async move {
489 let result = future.await;
490 let duration_ms = start_time.elapsed().as_micros() as f64 / 1000.0;
491 match result {
492 Err(error) => {
493 tracing::error!(%error, ?duration_ms, "error handling message")
494 }
495 Ok(()) => tracing::info!(?duration_ms, "finished handling message"),
496 }
497 }
498 .instrument(span)
499 .boxed()
500 }),
501 );
502 if prev_handler.is_some() {
503 panic!("registered a handler for the same message twice");
504 }
505 self
506 }
507
508 fn add_message_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
509 where
510 F: 'static + Send + Sync + Fn(M, Session) -> Fut,
511 Fut: 'static + Send + Future<Output = Result<()>>,
512 M: EnvelopedMessage,
513 {
514 self.add_handler(move |envelope, session| handler(envelope.payload, session));
515 self
516 }
517
518 fn add_request_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
519 where
520 F: 'static + Send + Sync + Fn(M, Response<M>, Session) -> Fut,
521 Fut: Send + Future<Output = Result<()>>,
522 M: RequestMessage,
523 {
524 let handler = Arc::new(handler);
525 self.add_handler(move |envelope, session| {
526 let receipt = envelope.receipt();
527 let handler = handler.clone();
528 async move {
529 let peer = session.peer.clone();
530 let responded = Arc::new(AtomicBool::default());
531 let response = Response {
532 peer: peer.clone(),
533 responded: responded.clone(),
534 receipt,
535 };
536 match (handler)(envelope.payload, response, session).await {
537 Ok(()) => {
538 if responded.load(std::sync::atomic::Ordering::SeqCst) {
539 Ok(())
540 } else {
541 Err(anyhow!("handler did not send a response"))?
542 }
543 }
544 Err(error) => {
545 peer.respond_with_error(
546 receipt,
547 proto::Error {
548 message: error.to_string(),
549 },
550 )?;
551 Err(error)
552 }
553 }
554 }
555 })
556 }
557
558 pub fn handle_connection(
559 self: &Arc<Self>,
560 connection: Connection,
561 address: String,
562 user: User,
563 mut send_connection_id: Option<oneshot::Sender<ConnectionId>>,
564 executor: Executor,
565 ) -> impl Future<Output = Result<()>> {
566 let this = self.clone();
567 let user_id = user.id;
568 let login = user.github_login;
569 let span = info_span!("handle connection", %user_id, %login, %address);
570 let mut teardown = self.teardown.subscribe();
571 async move {
572 let (connection_id, handle_io, mut incoming_rx) = this
573 .peer
574 .add_connection(connection, {
575 let executor = executor.clone();
576 move |duration| executor.sleep(duration)
577 });
578
579 tracing::info!(%user_id, %login, %connection_id, %address, "connection opened");
580 this.peer.send(connection_id, proto::Hello { peer_id: Some(connection_id.into()) })?;
581 tracing::info!(%user_id, %login, %connection_id, %address, "sent hello message");
582
583 if let Some(send_connection_id) = send_connection_id.take() {
584 let _ = send_connection_id.send(connection_id);
585 }
586
587 if !user.connected_once {
588 this.peer.send(connection_id, proto::ShowContacts {})?;
589 this.app_state.db.set_user_connected_once(user_id, true).await?;
590 }
591
592 let (contacts, channels_for_user, channel_invites) = future::try_join3(
593 this.app_state.db.get_contacts(user_id),
594 this.app_state.db.get_channels_for_user(user_id),
595 this.app_state.db.get_channel_invites_for_user(user_id),
596 ).await?;
597
598 {
599 let mut pool = this.connection_pool.lock();
600 pool.add_connection(connection_id, user_id, user.admin);
601 this.peer.send(connection_id, build_initial_contacts_update(contacts, &pool))?;
602 this.peer.send(connection_id, build_channels_update(
603 channels_for_user,
604 channel_invites
605 ))?;
606 }
607
608 if let Some(incoming_call) = this.app_state.db.incoming_call_for_user(user_id).await? {
609 this.peer.send(connection_id, incoming_call)?;
610 }
611
612 let session = Session {
613 user_id,
614 connection_id,
615 db: Arc::new(tokio::sync::Mutex::new(DbHandle(this.app_state.db.clone()))),
616 peer: this.peer.clone(),
617 connection_pool: this.connection_pool.clone(),
618 live_kit_client: this.app_state.live_kit_client.clone(),
619 _executor: executor.clone()
620 };
621 update_user_contacts(user_id, &session).await?;
622
623 let handle_io = handle_io.fuse();
624 futures::pin_mut!(handle_io);
625
626 // Handlers for foreground messages are pushed into the following `FuturesUnordered`.
627 // This prevents deadlocks when e.g., client A performs a request to client B and
628 // client B performs a request to client A. If both clients stop processing further
629 // messages until their respective request completes, they won't have a chance to
630 // respond to the other client's request and cause a deadlock.
631 //
632 // This arrangement ensures we will attempt to process earlier messages first, but fall
633 // back to processing messages arrived later in the spirit of making progress.
634 let mut foreground_message_handlers = FuturesUnordered::new();
635 let concurrent_handlers = Arc::new(Semaphore::new(256));
636 loop {
637 let next_message = async {
638 let permit = concurrent_handlers.clone().acquire_owned().await.unwrap();
639 let message = incoming_rx.next().await;
640 (permit, message)
641 }.fuse();
642 futures::pin_mut!(next_message);
643 futures::select_biased! {
644 _ = teardown.changed().fuse() => return Ok(()),
645 result = handle_io => {
646 if let Err(error) = result {
647 tracing::error!(?error, %user_id, %login, %connection_id, %address, "error handling I/O");
648 }
649 break;
650 }
651 _ = foreground_message_handlers.next() => {}
652 next_message = next_message => {
653 let (permit, message) = next_message;
654 if let Some(message) = message {
655 let type_name = message.payload_type_name();
656 let span = tracing::info_span!("receive message", %user_id, %login, %connection_id, %address, type_name);
657 let span_enter = span.enter();
658 if let Some(handler) = this.handlers.get(&message.payload_type_id()) {
659 let is_background = message.is_background();
660 let handle_message = (handler)(message, session.clone());
661 drop(span_enter);
662
663 let handle_message = async move {
664 handle_message.await;
665 drop(permit);
666 }.instrument(span);
667 if is_background {
668 executor.spawn_detached(handle_message);
669 } else {
670 foreground_message_handlers.push(handle_message);
671 }
672 } else {
673 tracing::error!(%user_id, %login, %connection_id, %address, "no message handler");
674 }
675 } else {
676 tracing::info!(%user_id, %login, %connection_id, %address, "connection closed");
677 break;
678 }
679 }
680 }
681 }
682
683 drop(foreground_message_handlers);
684 tracing::info!(%user_id, %login, %connection_id, %address, "signing out");
685 if let Err(error) = connection_lost(session, teardown, executor).await {
686 tracing::error!(%user_id, %login, %connection_id, %address, ?error, "error signing out");
687 }
688
689 Ok(())
690 }.instrument(span)
691 }
692
693 pub async fn invite_code_redeemed(
694 self: &Arc<Self>,
695 inviter_id: UserId,
696 invitee_id: UserId,
697 ) -> Result<()> {
698 if let Some(user) = self.app_state.db.get_user_by_id(inviter_id).await? {
699 if let Some(code) = &user.invite_code {
700 let pool = self.connection_pool.lock();
701 let invitee_contact = contact_for_user(invitee_id, false, &pool);
702 for connection_id in pool.user_connection_ids(inviter_id) {
703 self.peer.send(
704 connection_id,
705 proto::UpdateContacts {
706 contacts: vec![invitee_contact.clone()],
707 ..Default::default()
708 },
709 )?;
710 self.peer.send(
711 connection_id,
712 proto::UpdateInviteInfo {
713 url: format!("{}{}", self.app_state.config.invite_link_prefix, &code),
714 count: user.invite_count as u32,
715 },
716 )?;
717 }
718 }
719 }
720 Ok(())
721 }
722
723 pub async fn invite_count_updated(self: &Arc<Self>, user_id: UserId) -> Result<()> {
724 if let Some(user) = self.app_state.db.get_user_by_id(user_id).await? {
725 if let Some(invite_code) = &user.invite_code {
726 let pool = self.connection_pool.lock();
727 for connection_id in pool.user_connection_ids(user_id) {
728 self.peer.send(
729 connection_id,
730 proto::UpdateInviteInfo {
731 url: format!(
732 "{}{}",
733 self.app_state.config.invite_link_prefix, invite_code
734 ),
735 count: user.invite_count as u32,
736 },
737 )?;
738 }
739 }
740 }
741 Ok(())
742 }
743
744 pub async fn snapshot<'a>(self: &'a Arc<Self>) -> ServerSnapshot<'a> {
745 ServerSnapshot {
746 connection_pool: ConnectionPoolGuard {
747 guard: self.connection_pool.lock(),
748 _not_send: PhantomData,
749 },
750 peer: &self.peer,
751 }
752 }
753}
754
755impl<'a> Deref for ConnectionPoolGuard<'a> {
756 type Target = ConnectionPool;
757
758 fn deref(&self) -> &Self::Target {
759 &*self.guard
760 }
761}
762
763impl<'a> DerefMut for ConnectionPoolGuard<'a> {
764 fn deref_mut(&mut self) -> &mut Self::Target {
765 &mut *self.guard
766 }
767}
768
769impl<'a> Drop for ConnectionPoolGuard<'a> {
770 fn drop(&mut self) {
771 #[cfg(test)]
772 self.check_invariants();
773 }
774}
775
776fn broadcast<F>(
777 sender_id: Option<ConnectionId>,
778 receiver_ids: impl IntoIterator<Item = ConnectionId>,
779 mut f: F,
780) where
781 F: FnMut(ConnectionId) -> anyhow::Result<()>,
782{
783 for receiver_id in receiver_ids {
784 if Some(receiver_id) != sender_id {
785 if let Err(error) = f(receiver_id) {
786 tracing::error!("failed to send to {:?} {}", receiver_id, error);
787 }
788 }
789 }
790}
791
792lazy_static! {
793 static ref ZED_PROTOCOL_VERSION: HeaderName = HeaderName::from_static("x-zed-protocol-version");
794}
795
796pub struct ProtocolVersion(u32);
797
798impl Header for ProtocolVersion {
799 fn name() -> &'static HeaderName {
800 &ZED_PROTOCOL_VERSION
801 }
802
803 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
804 where
805 Self: Sized,
806 I: Iterator<Item = &'i axum::http::HeaderValue>,
807 {
808 let version = values
809 .next()
810 .ok_or_else(axum::headers::Error::invalid)?
811 .to_str()
812 .map_err(|_| axum::headers::Error::invalid())?
813 .parse()
814 .map_err(|_| axum::headers::Error::invalid())?;
815 Ok(Self(version))
816 }
817
818 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
819 values.extend([self.0.to_string().parse().unwrap()]);
820 }
821}
822
823pub fn routes(server: Arc<Server>) -> Router<Body> {
824 Router::new()
825 .route("/rpc", get(handle_websocket_request))
826 .layer(
827 ServiceBuilder::new()
828 .layer(Extension(server.app_state.clone()))
829 .layer(middleware::from_fn(auth::validate_header)),
830 )
831 .route("/metrics", get(handle_metrics))
832 .layer(Extension(server))
833}
834
835pub async fn handle_websocket_request(
836 TypedHeader(ProtocolVersion(protocol_version)): TypedHeader<ProtocolVersion>,
837 ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
838 Extension(server): Extension<Arc<Server>>,
839 Extension(user): Extension<User>,
840 ws: WebSocketUpgrade,
841) -> axum::response::Response {
842 if protocol_version != rpc::PROTOCOL_VERSION {
843 return (
844 StatusCode::UPGRADE_REQUIRED,
845 "client must be upgraded".to_string(),
846 )
847 .into_response();
848 }
849 let socket_address = socket_address.to_string();
850 ws.on_upgrade(move |socket| {
851 use util::ResultExt;
852 let socket = socket
853 .map_ok(to_tungstenite_message)
854 .err_into()
855 .with(|message| async move { Ok(to_axum_message(message)) });
856 let connection = Connection::new(Box::pin(socket));
857 async move {
858 server
859 .handle_connection(connection, socket_address, user, None, Executor::Production)
860 .await
861 .log_err();
862 }
863 })
864}
865
866pub async fn handle_metrics(Extension(server): Extension<Arc<Server>>) -> Result<String> {
867 let connections = server
868 .connection_pool
869 .lock()
870 .connections()
871 .filter(|connection| !connection.admin)
872 .count();
873
874 METRIC_CONNECTIONS.set(connections as _);
875
876 let shared_projects = server.app_state.db.project_count_excluding_admins().await?;
877 METRIC_SHARED_PROJECTS.set(shared_projects as _);
878
879 let encoder = prometheus::TextEncoder::new();
880 let metric_families = prometheus::gather();
881 let encoded_metrics = encoder
882 .encode_to_string(&metric_families)
883 .map_err(|err| anyhow!("{}", err))?;
884 Ok(encoded_metrics)
885}
886
887#[instrument(err, skip(executor))]
888async fn connection_lost(
889 session: Session,
890 mut teardown: watch::Receiver<()>,
891 executor: Executor,
892) -> Result<()> {
893 session.peer.disconnect(session.connection_id);
894 session
895 .connection_pool()
896 .await
897 .remove_connection(session.connection_id)?;
898
899 session
900 .db()
901 .await
902 .connection_lost(session.connection_id)
903 .await
904 .trace_err();
905
906 futures::select_biased! {
907 _ = executor.sleep(RECONNECT_TIMEOUT).fuse() => {
908 log::info!("connection lost, removing all resources for user:{}, connection:{:?}", session.user_id, session.connection_id);
909 leave_room_for_session(&session).await.trace_err();
910 leave_channel_buffers_for_session(&session)
911 .await
912 .trace_err();
913
914 if !session
915 .connection_pool()
916 .await
917 .is_user_online(session.user_id)
918 {
919 let db = session.db().await;
920 if let Some(room) = db.decline_call(None, session.user_id).await.trace_err().flatten() {
921 room_updated(&room, &session.peer);
922 }
923 }
924
925 update_user_contacts(session.user_id, &session).await?;
926 }
927 _ = teardown.changed().fuse() => {}
928 }
929
930 Ok(())
931}
932
933async fn ping(_: proto::Ping, response: Response<proto::Ping>, _session: Session) -> Result<()> {
934 response.send(proto::Ack {})?;
935 Ok(())
936}
937
938async fn create_room(
939 _request: proto::CreateRoom,
940 response: Response<proto::CreateRoom>,
941 session: Session,
942) -> Result<()> {
943 let live_kit_room = nanoid::nanoid!(30);
944
945 let live_kit_connection_info = {
946 let live_kit_room = live_kit_room.clone();
947 let live_kit = session.live_kit_client.as_ref();
948
949 util::async_maybe!({
950 let live_kit = live_kit?;
951
952 let token = live_kit
953 .room_token(&live_kit_room, &session.user_id.to_string())
954 .trace_err()?;
955
956 Some(proto::LiveKitConnectionInfo {
957 server_url: live_kit.url().into(),
958 token,
959 can_publish: true,
960 })
961 })
962 }
963 .await;
964
965 let room = session
966 .db()
967 .await
968 .create_room(
969 session.user_id,
970 session.connection_id,
971 &live_kit_room,
972 RELEASE_CHANNEL_NAME.as_str(),
973 )
974 .await?;
975
976 response.send(proto::CreateRoomResponse {
977 room: Some(room.clone()),
978 live_kit_connection_info,
979 })?;
980
981 update_user_contacts(session.user_id, &session).await?;
982 Ok(())
983}
984
985async fn join_room(
986 request: proto::JoinRoom,
987 response: Response<proto::JoinRoom>,
988 session: Session,
989) -> Result<()> {
990 let room_id = RoomId::from_proto(request.id);
991
992 let channel_id = session.db().await.channel_id_for_room(room_id).await?;
993
994 if let Some(channel_id) = channel_id {
995 return join_channel_internal(channel_id, Box::new(response), session).await;
996 }
997
998 let joined_room = {
999 let room = session
1000 .db()
1001 .await
1002 .join_room(
1003 room_id,
1004 session.user_id,
1005 session.connection_id,
1006 RELEASE_CHANNEL_NAME.as_str(),
1007 )
1008 .await?;
1009 room_updated(&room.room, &session.peer);
1010 room.into_inner()
1011 };
1012
1013 for connection_id in session
1014 .connection_pool()
1015 .await
1016 .user_connection_ids(session.user_id)
1017 {
1018 session
1019 .peer
1020 .send(
1021 connection_id,
1022 proto::CallCanceled {
1023 room_id: room_id.to_proto(),
1024 },
1025 )
1026 .trace_err();
1027 }
1028
1029 let live_kit_connection_info = if let Some(live_kit) = session.live_kit_client.as_ref() {
1030 if let Some(token) = live_kit
1031 .room_token(
1032 &joined_room.room.live_kit_room,
1033 &session.user_id.to_string(),
1034 )
1035 .trace_err()
1036 {
1037 Some(proto::LiveKitConnectionInfo {
1038 server_url: live_kit.url().into(),
1039 token,
1040 can_publish: true,
1041 })
1042 } else {
1043 None
1044 }
1045 } else {
1046 None
1047 };
1048
1049 response.send(proto::JoinRoomResponse {
1050 room: Some(joined_room.room),
1051 channel_id: None,
1052 live_kit_connection_info,
1053 })?;
1054
1055 update_user_contacts(session.user_id, &session).await?;
1056 Ok(())
1057}
1058
1059async fn rejoin_room(
1060 request: proto::RejoinRoom,
1061 response: Response<proto::RejoinRoom>,
1062 session: Session,
1063) -> Result<()> {
1064 let room;
1065 let channel_id;
1066 let channel_members;
1067 {
1068 let mut rejoined_room = session
1069 .db()
1070 .await
1071 .rejoin_room(request, session.user_id, session.connection_id)
1072 .await?;
1073
1074 response.send(proto::RejoinRoomResponse {
1075 room: Some(rejoined_room.room.clone()),
1076 reshared_projects: rejoined_room
1077 .reshared_projects
1078 .iter()
1079 .map(|project| proto::ResharedProject {
1080 id: project.id.to_proto(),
1081 collaborators: project
1082 .collaborators
1083 .iter()
1084 .map(|collaborator| collaborator.to_proto())
1085 .collect(),
1086 })
1087 .collect(),
1088 rejoined_projects: rejoined_room
1089 .rejoined_projects
1090 .iter()
1091 .map(|rejoined_project| proto::RejoinedProject {
1092 id: rejoined_project.id.to_proto(),
1093 worktrees: rejoined_project
1094 .worktrees
1095 .iter()
1096 .map(|worktree| proto::WorktreeMetadata {
1097 id: worktree.id,
1098 root_name: worktree.root_name.clone(),
1099 visible: worktree.visible,
1100 abs_path: worktree.abs_path.clone(),
1101 })
1102 .collect(),
1103 collaborators: rejoined_project
1104 .collaborators
1105 .iter()
1106 .map(|collaborator| collaborator.to_proto())
1107 .collect(),
1108 language_servers: rejoined_project.language_servers.clone(),
1109 })
1110 .collect(),
1111 })?;
1112 room_updated(&rejoined_room.room, &session.peer);
1113
1114 for project in &rejoined_room.reshared_projects {
1115 for collaborator in &project.collaborators {
1116 session
1117 .peer
1118 .send(
1119 collaborator.connection_id,
1120 proto::UpdateProjectCollaborator {
1121 project_id: project.id.to_proto(),
1122 old_peer_id: Some(project.old_connection_id.into()),
1123 new_peer_id: Some(session.connection_id.into()),
1124 },
1125 )
1126 .trace_err();
1127 }
1128
1129 broadcast(
1130 Some(session.connection_id),
1131 project
1132 .collaborators
1133 .iter()
1134 .map(|collaborator| collaborator.connection_id),
1135 |connection_id| {
1136 session.peer.forward_send(
1137 session.connection_id,
1138 connection_id,
1139 proto::UpdateProject {
1140 project_id: project.id.to_proto(),
1141 worktrees: project.worktrees.clone(),
1142 },
1143 )
1144 },
1145 );
1146 }
1147
1148 for project in &rejoined_room.rejoined_projects {
1149 for collaborator in &project.collaborators {
1150 session
1151 .peer
1152 .send(
1153 collaborator.connection_id,
1154 proto::UpdateProjectCollaborator {
1155 project_id: project.id.to_proto(),
1156 old_peer_id: Some(project.old_connection_id.into()),
1157 new_peer_id: Some(session.connection_id.into()),
1158 },
1159 )
1160 .trace_err();
1161 }
1162 }
1163
1164 for project in &mut rejoined_room.rejoined_projects {
1165 for worktree in mem::take(&mut project.worktrees) {
1166 #[cfg(any(test, feature = "test-support"))]
1167 const MAX_CHUNK_SIZE: usize = 2;
1168 #[cfg(not(any(test, feature = "test-support")))]
1169 const MAX_CHUNK_SIZE: usize = 256;
1170
1171 // Stream this worktree's entries.
1172 let message = proto::UpdateWorktree {
1173 project_id: project.id.to_proto(),
1174 worktree_id: worktree.id,
1175 abs_path: worktree.abs_path.clone(),
1176 root_name: worktree.root_name,
1177 updated_entries: worktree.updated_entries,
1178 removed_entries: worktree.removed_entries,
1179 scan_id: worktree.scan_id,
1180 is_last_update: worktree.completed_scan_id == worktree.scan_id,
1181 updated_repositories: worktree.updated_repositories,
1182 removed_repositories: worktree.removed_repositories,
1183 };
1184 for update in proto::split_worktree_update(message, MAX_CHUNK_SIZE) {
1185 session.peer.send(session.connection_id, update.clone())?;
1186 }
1187
1188 // Stream this worktree's diagnostics.
1189 for summary in worktree.diagnostic_summaries {
1190 session.peer.send(
1191 session.connection_id,
1192 proto::UpdateDiagnosticSummary {
1193 project_id: project.id.to_proto(),
1194 worktree_id: worktree.id,
1195 summary: Some(summary),
1196 },
1197 )?;
1198 }
1199
1200 for settings_file in worktree.settings_files {
1201 session.peer.send(
1202 session.connection_id,
1203 proto::UpdateWorktreeSettings {
1204 project_id: project.id.to_proto(),
1205 worktree_id: worktree.id,
1206 path: settings_file.path,
1207 content: Some(settings_file.content),
1208 },
1209 )?;
1210 }
1211 }
1212
1213 for language_server in &project.language_servers {
1214 session.peer.send(
1215 session.connection_id,
1216 proto::UpdateLanguageServer {
1217 project_id: project.id.to_proto(),
1218 language_server_id: language_server.id,
1219 variant: Some(
1220 proto::update_language_server::Variant::DiskBasedDiagnosticsUpdated(
1221 proto::LspDiskBasedDiagnosticsUpdated {},
1222 ),
1223 ),
1224 },
1225 )?;
1226 }
1227 }
1228
1229 let rejoined_room = rejoined_room.into_inner();
1230
1231 room = rejoined_room.room;
1232 channel_id = rejoined_room.channel_id;
1233 channel_members = rejoined_room.channel_members;
1234 }
1235
1236 if let Some(channel_id) = channel_id {
1237 channel_updated(
1238 channel_id,
1239 &room,
1240 &channel_members,
1241 &session.peer,
1242 &*session.connection_pool().await,
1243 );
1244 }
1245
1246 update_user_contacts(session.user_id, &session).await?;
1247 Ok(())
1248}
1249
1250async fn leave_room(
1251 _: proto::LeaveRoom,
1252 response: Response<proto::LeaveRoom>,
1253 session: Session,
1254) -> Result<()> {
1255 leave_room_for_session(&session).await?;
1256 response.send(proto::Ack {})?;
1257 Ok(())
1258}
1259
1260async fn call(
1261 request: proto::Call,
1262 response: Response<proto::Call>,
1263 session: Session,
1264) -> Result<()> {
1265 let room_id = RoomId::from_proto(request.room_id);
1266 let calling_user_id = session.user_id;
1267 let calling_connection_id = session.connection_id;
1268 let called_user_id = UserId::from_proto(request.called_user_id);
1269 let initial_project_id = request.initial_project_id.map(ProjectId::from_proto);
1270 if !session
1271 .db()
1272 .await
1273 .has_contact(calling_user_id, called_user_id)
1274 .await?
1275 {
1276 return Err(anyhow!("cannot call a user who isn't a contact"))?;
1277 }
1278
1279 let incoming_call = {
1280 let (room, incoming_call) = &mut *session
1281 .db()
1282 .await
1283 .call(
1284 room_id,
1285 calling_user_id,
1286 calling_connection_id,
1287 called_user_id,
1288 initial_project_id,
1289 )
1290 .await?;
1291 room_updated(&room, &session.peer);
1292 mem::take(incoming_call)
1293 };
1294 update_user_contacts(called_user_id, &session).await?;
1295
1296 let mut calls = session
1297 .connection_pool()
1298 .await
1299 .user_connection_ids(called_user_id)
1300 .map(|connection_id| session.peer.request(connection_id, incoming_call.clone()))
1301 .collect::<FuturesUnordered<_>>();
1302
1303 while let Some(call_response) = calls.next().await {
1304 match call_response.as_ref() {
1305 Ok(_) => {
1306 response.send(proto::Ack {})?;
1307 return Ok(());
1308 }
1309 Err(_) => {
1310 call_response.trace_err();
1311 }
1312 }
1313 }
1314
1315 {
1316 let room = session
1317 .db()
1318 .await
1319 .call_failed(room_id, called_user_id)
1320 .await?;
1321 room_updated(&room, &session.peer);
1322 }
1323 update_user_contacts(called_user_id, &session).await?;
1324
1325 Err(anyhow!("failed to ring user"))?
1326}
1327
1328async fn cancel_call(
1329 request: proto::CancelCall,
1330 response: Response<proto::CancelCall>,
1331 session: Session,
1332) -> Result<()> {
1333 let called_user_id = UserId::from_proto(request.called_user_id);
1334 let room_id = RoomId::from_proto(request.room_id);
1335 {
1336 let room = session
1337 .db()
1338 .await
1339 .cancel_call(room_id, session.connection_id, called_user_id)
1340 .await?;
1341 room_updated(&room, &session.peer);
1342 }
1343
1344 for connection_id in session
1345 .connection_pool()
1346 .await
1347 .user_connection_ids(called_user_id)
1348 {
1349 session
1350 .peer
1351 .send(
1352 connection_id,
1353 proto::CallCanceled {
1354 room_id: room_id.to_proto(),
1355 },
1356 )
1357 .trace_err();
1358 }
1359 response.send(proto::Ack {})?;
1360
1361 update_user_contacts(called_user_id, &session).await?;
1362 Ok(())
1363}
1364
1365async fn decline_call(message: proto::DeclineCall, session: Session) -> Result<()> {
1366 let room_id = RoomId::from_proto(message.room_id);
1367 {
1368 let room = session
1369 .db()
1370 .await
1371 .decline_call(Some(room_id), session.user_id)
1372 .await?
1373 .ok_or_else(|| anyhow!("failed to decline call"))?;
1374 room_updated(&room, &session.peer);
1375 }
1376
1377 for connection_id in session
1378 .connection_pool()
1379 .await
1380 .user_connection_ids(session.user_id)
1381 {
1382 session
1383 .peer
1384 .send(
1385 connection_id,
1386 proto::CallCanceled {
1387 room_id: room_id.to_proto(),
1388 },
1389 )
1390 .trace_err();
1391 }
1392 update_user_contacts(session.user_id, &session).await?;
1393 Ok(())
1394}
1395
1396async fn update_participant_location(
1397 request: proto::UpdateParticipantLocation,
1398 response: Response<proto::UpdateParticipantLocation>,
1399 session: Session,
1400) -> Result<()> {
1401 let room_id = RoomId::from_proto(request.room_id);
1402 let location = request
1403 .location
1404 .ok_or_else(|| anyhow!("invalid location"))?;
1405
1406 let db = session.db().await;
1407 let room = db
1408 .update_room_participant_location(room_id, session.connection_id, location)
1409 .await?;
1410
1411 room_updated(&room, &session.peer);
1412 response.send(proto::Ack {})?;
1413 Ok(())
1414}
1415
1416async fn share_project(
1417 request: proto::ShareProject,
1418 response: Response<proto::ShareProject>,
1419 session: Session,
1420) -> Result<()> {
1421 let (project_id, room) = &*session
1422 .db()
1423 .await
1424 .share_project(
1425 RoomId::from_proto(request.room_id),
1426 session.connection_id,
1427 &request.worktrees,
1428 )
1429 .await?;
1430 response.send(proto::ShareProjectResponse {
1431 project_id: project_id.to_proto(),
1432 })?;
1433 room_updated(&room, &session.peer);
1434
1435 Ok(())
1436}
1437
1438async fn unshare_project(message: proto::UnshareProject, session: Session) -> Result<()> {
1439 let project_id = ProjectId::from_proto(message.project_id);
1440
1441 let (room, guest_connection_ids) = &*session
1442 .db()
1443 .await
1444 .unshare_project(project_id, session.connection_id)
1445 .await?;
1446
1447 broadcast(
1448 Some(session.connection_id),
1449 guest_connection_ids.iter().copied(),
1450 |conn_id| session.peer.send(conn_id, message.clone()),
1451 );
1452 room_updated(&room, &session.peer);
1453
1454 Ok(())
1455}
1456
1457async fn join_project(
1458 request: proto::JoinProject,
1459 response: Response<proto::JoinProject>,
1460 session: Session,
1461) -> Result<()> {
1462 let project_id = ProjectId::from_proto(request.project_id);
1463 let guest_user_id = session.user_id;
1464
1465 tracing::info!(%project_id, "join project");
1466
1467 let (project, replica_id) = &mut *session
1468 .db()
1469 .await
1470 .join_project(project_id, session.connection_id)
1471 .await?;
1472
1473 let collaborators = project
1474 .collaborators
1475 .iter()
1476 .filter(|collaborator| collaborator.connection_id != session.connection_id)
1477 .map(|collaborator| collaborator.to_proto())
1478 .collect::<Vec<_>>();
1479
1480 let worktrees = project
1481 .worktrees
1482 .iter()
1483 .map(|(id, worktree)| proto::WorktreeMetadata {
1484 id: *id,
1485 root_name: worktree.root_name.clone(),
1486 visible: worktree.visible,
1487 abs_path: worktree.abs_path.clone(),
1488 })
1489 .collect::<Vec<_>>();
1490
1491 for collaborator in &collaborators {
1492 session
1493 .peer
1494 .send(
1495 collaborator.peer_id.unwrap().into(),
1496 proto::AddProjectCollaborator {
1497 project_id: project_id.to_proto(),
1498 collaborator: Some(proto::Collaborator {
1499 peer_id: Some(session.connection_id.into()),
1500 replica_id: replica_id.0 as u32,
1501 user_id: guest_user_id.to_proto(),
1502 }),
1503 },
1504 )
1505 .trace_err();
1506 }
1507
1508 // First, we send the metadata associated with each worktree.
1509 response.send(proto::JoinProjectResponse {
1510 worktrees: worktrees.clone(),
1511 replica_id: replica_id.0 as u32,
1512 collaborators: collaborators.clone(),
1513 language_servers: project.language_servers.clone(),
1514 })?;
1515
1516 for (worktree_id, worktree) in mem::take(&mut project.worktrees) {
1517 #[cfg(any(test, feature = "test-support"))]
1518 const MAX_CHUNK_SIZE: usize = 2;
1519 #[cfg(not(any(test, feature = "test-support")))]
1520 const MAX_CHUNK_SIZE: usize = 256;
1521
1522 // Stream this worktree's entries.
1523 let message = proto::UpdateWorktree {
1524 project_id: project_id.to_proto(),
1525 worktree_id,
1526 abs_path: worktree.abs_path.clone(),
1527 root_name: worktree.root_name,
1528 updated_entries: worktree.entries,
1529 removed_entries: Default::default(),
1530 scan_id: worktree.scan_id,
1531 is_last_update: worktree.scan_id == worktree.completed_scan_id,
1532 updated_repositories: worktree.repository_entries.into_values().collect(),
1533 removed_repositories: Default::default(),
1534 };
1535 for update in proto::split_worktree_update(message, MAX_CHUNK_SIZE) {
1536 session.peer.send(session.connection_id, update.clone())?;
1537 }
1538
1539 // Stream this worktree's diagnostics.
1540 for summary in worktree.diagnostic_summaries {
1541 session.peer.send(
1542 session.connection_id,
1543 proto::UpdateDiagnosticSummary {
1544 project_id: project_id.to_proto(),
1545 worktree_id: worktree.id,
1546 summary: Some(summary),
1547 },
1548 )?;
1549 }
1550
1551 for settings_file in worktree.settings_files {
1552 session.peer.send(
1553 session.connection_id,
1554 proto::UpdateWorktreeSettings {
1555 project_id: project_id.to_proto(),
1556 worktree_id: worktree.id,
1557 path: settings_file.path,
1558 content: Some(settings_file.content),
1559 },
1560 )?;
1561 }
1562 }
1563
1564 for language_server in &project.language_servers {
1565 session.peer.send(
1566 session.connection_id,
1567 proto::UpdateLanguageServer {
1568 project_id: project_id.to_proto(),
1569 language_server_id: language_server.id,
1570 variant: Some(
1571 proto::update_language_server::Variant::DiskBasedDiagnosticsUpdated(
1572 proto::LspDiskBasedDiagnosticsUpdated {},
1573 ),
1574 ),
1575 },
1576 )?;
1577 }
1578
1579 Ok(())
1580}
1581
1582async fn leave_project(request: proto::LeaveProject, session: Session) -> Result<()> {
1583 let sender_id = session.connection_id;
1584 let project_id = ProjectId::from_proto(request.project_id);
1585
1586 let (room, project) = &*session
1587 .db()
1588 .await
1589 .leave_project(project_id, sender_id)
1590 .await?;
1591 tracing::info!(
1592 %project_id,
1593 host_user_id = %project.host_user_id,
1594 host_connection_id = %project.host_connection_id,
1595 "leave project"
1596 );
1597
1598 project_left(&project, &session);
1599 room_updated(&room, &session.peer);
1600
1601 Ok(())
1602}
1603
1604async fn update_project(
1605 request: proto::UpdateProject,
1606 response: Response<proto::UpdateProject>,
1607 session: Session,
1608) -> Result<()> {
1609 let project_id = ProjectId::from_proto(request.project_id);
1610 let (room, guest_connection_ids) = &*session
1611 .db()
1612 .await
1613 .update_project(project_id, session.connection_id, &request.worktrees)
1614 .await?;
1615 broadcast(
1616 Some(session.connection_id),
1617 guest_connection_ids.iter().copied(),
1618 |connection_id| {
1619 session
1620 .peer
1621 .forward_send(session.connection_id, connection_id, request.clone())
1622 },
1623 );
1624 room_updated(&room, &session.peer);
1625 response.send(proto::Ack {})?;
1626
1627 Ok(())
1628}
1629
1630async fn update_worktree(
1631 request: proto::UpdateWorktree,
1632 response: Response<proto::UpdateWorktree>,
1633 session: Session,
1634) -> Result<()> {
1635 let guest_connection_ids = session
1636 .db()
1637 .await
1638 .update_worktree(&request, session.connection_id)
1639 .await?;
1640
1641 broadcast(
1642 Some(session.connection_id),
1643 guest_connection_ids.iter().copied(),
1644 |connection_id| {
1645 session
1646 .peer
1647 .forward_send(session.connection_id, connection_id, request.clone())
1648 },
1649 );
1650 response.send(proto::Ack {})?;
1651 Ok(())
1652}
1653
1654async fn update_diagnostic_summary(
1655 message: proto::UpdateDiagnosticSummary,
1656 session: Session,
1657) -> Result<()> {
1658 let guest_connection_ids = session
1659 .db()
1660 .await
1661 .update_diagnostic_summary(&message, session.connection_id)
1662 .await?;
1663
1664 broadcast(
1665 Some(session.connection_id),
1666 guest_connection_ids.iter().copied(),
1667 |connection_id| {
1668 session
1669 .peer
1670 .forward_send(session.connection_id, connection_id, message.clone())
1671 },
1672 );
1673
1674 Ok(())
1675}
1676
1677async fn update_worktree_settings(
1678 message: proto::UpdateWorktreeSettings,
1679 session: Session,
1680) -> Result<()> {
1681 let guest_connection_ids = session
1682 .db()
1683 .await
1684 .update_worktree_settings(&message, session.connection_id)
1685 .await?;
1686
1687 broadcast(
1688 Some(session.connection_id),
1689 guest_connection_ids.iter().copied(),
1690 |connection_id| {
1691 session
1692 .peer
1693 .forward_send(session.connection_id, connection_id, message.clone())
1694 },
1695 );
1696
1697 Ok(())
1698}
1699
1700async fn start_language_server(
1701 request: proto::StartLanguageServer,
1702 session: Session,
1703) -> Result<()> {
1704 let guest_connection_ids = session
1705 .db()
1706 .await
1707 .start_language_server(&request, session.connection_id)
1708 .await?;
1709
1710 broadcast(
1711 Some(session.connection_id),
1712 guest_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 update_language_server(
1723 request: proto::UpdateLanguageServer,
1724 session: Session,
1725) -> Result<()> {
1726 let project_id = ProjectId::from_proto(request.project_id);
1727 let project_connection_ids = session
1728 .db()
1729 .await
1730 .project_connection_ids(project_id, session.connection_id)
1731 .await?;
1732 broadcast(
1733 Some(session.connection_id),
1734 project_connection_ids.iter().copied(),
1735 |connection_id| {
1736 session
1737 .peer
1738 .forward_send(session.connection_id, connection_id, request.clone())
1739 },
1740 );
1741 Ok(())
1742}
1743
1744async fn forward_read_only_project_request<T>(
1745 request: T,
1746 response: Response<T>,
1747 session: Session,
1748) -> Result<()>
1749where
1750 T: EntityMessage + RequestMessage,
1751{
1752 let project_id = ProjectId::from_proto(request.remote_entity_id());
1753 let host_connection_id = session
1754 .db()
1755 .await
1756 .host_for_read_only_project_request(project_id, session.connection_id)
1757 .await?;
1758 let payload = session
1759 .peer
1760 .forward_request(session.connection_id, host_connection_id, request)
1761 .await?;
1762 response.send(payload)?;
1763 Ok(())
1764}
1765
1766async fn forward_mutating_project_request<T>(
1767 request: T,
1768 response: Response<T>,
1769 session: Session,
1770) -> Result<()>
1771where
1772 T: EntityMessage + RequestMessage,
1773{
1774 let project_id = ProjectId::from_proto(request.remote_entity_id());
1775 let host_connection_id = session
1776 .db()
1777 .await
1778 .host_for_mutating_project_request(project_id, session.connection_id)
1779 .await?;
1780 let payload = session
1781 .peer
1782 .forward_request(session.connection_id, host_connection_id, request)
1783 .await?;
1784 response.send(payload)?;
1785 Ok(())
1786}
1787
1788async fn create_buffer_for_peer(
1789 request: proto::CreateBufferForPeer,
1790 session: Session,
1791) -> Result<()> {
1792 session
1793 .db()
1794 .await
1795 .check_user_is_project_host(
1796 ProjectId::from_proto(request.project_id),
1797 session.connection_id,
1798 )
1799 .await?;
1800 let peer_id = request.peer_id.ok_or_else(|| anyhow!("invalid peer id"))?;
1801 session
1802 .peer
1803 .forward_send(session.connection_id, peer_id.into(), request)?;
1804 Ok(())
1805}
1806
1807async fn update_buffer(
1808 request: proto::UpdateBuffer,
1809 response: Response<proto::UpdateBuffer>,
1810 session: Session,
1811) -> Result<()> {
1812 let project_id = ProjectId::from_proto(request.project_id);
1813 let mut guest_connection_ids;
1814 let mut host_connection_id = None;
1815
1816 {
1817 let collaborators = session
1818 .db()
1819 .await
1820 .project_collaborators_for_buffer_update(project_id, session.connection_id)
1821 .await?;
1822 guest_connection_ids = Vec::with_capacity(collaborators.len() - 1);
1823 for collaborator in collaborators.iter() {
1824 if collaborator.is_host {
1825 host_connection_id = Some(collaborator.connection_id);
1826 } else {
1827 guest_connection_ids.push(collaborator.connection_id);
1828 }
1829 }
1830 }
1831 let host_connection_id = host_connection_id.ok_or_else(|| anyhow!("host not found"))?;
1832
1833 broadcast(
1834 Some(session.connection_id),
1835 guest_connection_ids,
1836 |connection_id| {
1837 session
1838 .peer
1839 .forward_send(session.connection_id, connection_id, request.clone())
1840 },
1841 );
1842 if host_connection_id != session.connection_id {
1843 session
1844 .peer
1845 .forward_request(session.connection_id, host_connection_id, request.clone())
1846 .await?;
1847 }
1848
1849 response.send(proto::Ack {})?;
1850 Ok(())
1851}
1852
1853async fn broadcast_project_message_from_host<T: EntityMessage<Entity = ShareProject>>(
1854 request: T,
1855 session: Session,
1856) -> Result<()> {
1857 let project_id = ProjectId::from_proto(request.remote_entity_id());
1858 let project_connection_ids = session
1859 .db()
1860 .await
1861 .project_connection_ids(project_id, session.connection_id)
1862 .await?;
1863
1864 broadcast(
1865 Some(session.connection_id),
1866 project_connection_ids.iter().copied(),
1867 |connection_id| {
1868 session
1869 .peer
1870 .forward_send(session.connection_id, connection_id, request.clone())
1871 },
1872 );
1873 Ok(())
1874}
1875
1876async fn follow(
1877 request: proto::Follow,
1878 response: Response<proto::Follow>,
1879 session: Session,
1880) -> Result<()> {
1881 let room_id = RoomId::from_proto(request.room_id);
1882 let project_id = request.project_id.map(ProjectId::from_proto);
1883 let leader_id = request
1884 .leader_id
1885 .ok_or_else(|| anyhow!("invalid leader id"))?
1886 .into();
1887 let follower_id = session.connection_id;
1888
1889 session
1890 .db()
1891 .await
1892 .check_room_participants(room_id, leader_id, session.connection_id)
1893 .await?;
1894
1895 let response_payload = session
1896 .peer
1897 .forward_request(session.connection_id, leader_id, request)
1898 .await?;
1899 response.send(response_payload)?;
1900
1901 if let Some(project_id) = project_id {
1902 let room = session
1903 .db()
1904 .await
1905 .follow(room_id, project_id, leader_id, follower_id)
1906 .await?;
1907 room_updated(&room, &session.peer);
1908 }
1909
1910 Ok(())
1911}
1912
1913async fn unfollow(request: proto::Unfollow, session: Session) -> Result<()> {
1914 let room_id = RoomId::from_proto(request.room_id);
1915 let project_id = request.project_id.map(ProjectId::from_proto);
1916 let leader_id = request
1917 .leader_id
1918 .ok_or_else(|| anyhow!("invalid leader id"))?
1919 .into();
1920 let follower_id = session.connection_id;
1921
1922 session
1923 .db()
1924 .await
1925 .check_room_participants(room_id, leader_id, session.connection_id)
1926 .await?;
1927
1928 session
1929 .peer
1930 .forward_send(session.connection_id, leader_id, request)?;
1931
1932 if let Some(project_id) = project_id {
1933 let room = session
1934 .db()
1935 .await
1936 .unfollow(room_id, project_id, leader_id, follower_id)
1937 .await?;
1938 room_updated(&room, &session.peer);
1939 }
1940
1941 Ok(())
1942}
1943
1944async fn update_followers(request: proto::UpdateFollowers, session: Session) -> Result<()> {
1945 let room_id = RoomId::from_proto(request.room_id);
1946 let database = session.db.lock().await;
1947
1948 let connection_ids = if let Some(project_id) = request.project_id {
1949 let project_id = ProjectId::from_proto(project_id);
1950 database
1951 .project_connection_ids(project_id, session.connection_id)
1952 .await?
1953 } else {
1954 database
1955 .room_connection_ids(room_id, session.connection_id)
1956 .await?
1957 };
1958
1959 // For now, don't send view update messages back to that view's current leader.
1960 let connection_id_to_omit = request.variant.as_ref().and_then(|variant| match variant {
1961 proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
1962 _ => None,
1963 });
1964
1965 for follower_peer_id in request.follower_ids.iter().copied() {
1966 let follower_connection_id = follower_peer_id.into();
1967 if Some(follower_peer_id) != connection_id_to_omit
1968 && connection_ids.contains(&follower_connection_id)
1969 {
1970 session.peer.forward_send(
1971 session.connection_id,
1972 follower_connection_id,
1973 request.clone(),
1974 )?;
1975 }
1976 }
1977 Ok(())
1978}
1979
1980async fn get_users(
1981 request: proto::GetUsers,
1982 response: Response<proto::GetUsers>,
1983 session: Session,
1984) -> Result<()> {
1985 let user_ids = request
1986 .user_ids
1987 .into_iter()
1988 .map(UserId::from_proto)
1989 .collect();
1990 let users = session
1991 .db()
1992 .await
1993 .get_users_by_ids(user_ids)
1994 .await?
1995 .into_iter()
1996 .map(|user| proto::User {
1997 id: user.id.to_proto(),
1998 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
1999 github_login: user.github_login,
2000 })
2001 .collect();
2002 response.send(proto::UsersResponse { users })?;
2003 Ok(())
2004}
2005
2006async fn fuzzy_search_users(
2007 request: proto::FuzzySearchUsers,
2008 response: Response<proto::FuzzySearchUsers>,
2009 session: Session,
2010) -> Result<()> {
2011 let query = request.query;
2012 let users = match query.len() {
2013 0 => vec![],
2014 1 | 2 => session
2015 .db()
2016 .await
2017 .get_user_by_github_login(&query)
2018 .await?
2019 .into_iter()
2020 .collect(),
2021 _ => session.db().await.fuzzy_search_users(&query, 10).await?,
2022 };
2023 let users = users
2024 .into_iter()
2025 .filter(|user| user.id != session.user_id)
2026 .map(|user| proto::User {
2027 id: user.id.to_proto(),
2028 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
2029 github_login: user.github_login,
2030 })
2031 .collect();
2032 response.send(proto::UsersResponse { users })?;
2033 Ok(())
2034}
2035
2036async fn request_contact(
2037 request: proto::RequestContact,
2038 response: Response<proto::RequestContact>,
2039 session: Session,
2040) -> Result<()> {
2041 let requester_id = session.user_id;
2042 let responder_id = UserId::from_proto(request.responder_id);
2043 if requester_id == responder_id {
2044 return Err(anyhow!("cannot add yourself as a contact"))?;
2045 }
2046
2047 let notifications = session
2048 .db()
2049 .await
2050 .send_contact_request(requester_id, responder_id)
2051 .await?;
2052
2053 // Update outgoing contact requests of requester
2054 let mut update = proto::UpdateContacts::default();
2055 update.outgoing_requests.push(responder_id.to_proto());
2056 for connection_id in session
2057 .connection_pool()
2058 .await
2059 .user_connection_ids(requester_id)
2060 {
2061 session.peer.send(connection_id, update.clone())?;
2062 }
2063
2064 // Update incoming contact requests of responder
2065 let mut update = proto::UpdateContacts::default();
2066 update
2067 .incoming_requests
2068 .push(proto::IncomingContactRequest {
2069 requester_id: requester_id.to_proto(),
2070 });
2071 let connection_pool = session.connection_pool().await;
2072 for connection_id in connection_pool.user_connection_ids(responder_id) {
2073 session.peer.send(connection_id, update.clone())?;
2074 }
2075
2076 send_notifications(&*connection_pool, &session.peer, notifications);
2077
2078 response.send(proto::Ack {})?;
2079 Ok(())
2080}
2081
2082async fn respond_to_contact_request(
2083 request: proto::RespondToContactRequest,
2084 response: Response<proto::RespondToContactRequest>,
2085 session: Session,
2086) -> Result<()> {
2087 let responder_id = session.user_id;
2088 let requester_id = UserId::from_proto(request.requester_id);
2089 let db = session.db().await;
2090 if request.response == proto::ContactRequestResponse::Dismiss as i32 {
2091 db.dismiss_contact_notification(responder_id, requester_id)
2092 .await?;
2093 } else {
2094 let accept = request.response == proto::ContactRequestResponse::Accept as i32;
2095
2096 let notifications = db
2097 .respond_to_contact_request(responder_id, requester_id, accept)
2098 .await?;
2099 let requester_busy = db.is_user_busy(requester_id).await?;
2100 let responder_busy = db.is_user_busy(responder_id).await?;
2101
2102 let pool = session.connection_pool().await;
2103 // Update responder with new contact
2104 let mut update = proto::UpdateContacts::default();
2105 if accept {
2106 update
2107 .contacts
2108 .push(contact_for_user(requester_id, requester_busy, &pool));
2109 }
2110 update
2111 .remove_incoming_requests
2112 .push(requester_id.to_proto());
2113 for connection_id in pool.user_connection_ids(responder_id) {
2114 session.peer.send(connection_id, update.clone())?;
2115 }
2116
2117 // Update requester with new contact
2118 let mut update = proto::UpdateContacts::default();
2119 if accept {
2120 update
2121 .contacts
2122 .push(contact_for_user(responder_id, responder_busy, &pool));
2123 }
2124 update
2125 .remove_outgoing_requests
2126 .push(responder_id.to_proto());
2127
2128 for connection_id in pool.user_connection_ids(requester_id) {
2129 session.peer.send(connection_id, update.clone())?;
2130 }
2131
2132 send_notifications(&*pool, &session.peer, notifications);
2133 }
2134
2135 response.send(proto::Ack {})?;
2136 Ok(())
2137}
2138
2139async fn remove_contact(
2140 request: proto::RemoveContact,
2141 response: Response<proto::RemoveContact>,
2142 session: Session,
2143) -> Result<()> {
2144 let requester_id = session.user_id;
2145 let responder_id = UserId::from_proto(request.user_id);
2146 let db = session.db().await;
2147 let (contact_accepted, deleted_notification_id) =
2148 db.remove_contact(requester_id, responder_id).await?;
2149
2150 let pool = session.connection_pool().await;
2151 // Update outgoing contact requests of requester
2152 let mut update = proto::UpdateContacts::default();
2153 if contact_accepted {
2154 update.remove_contacts.push(responder_id.to_proto());
2155 } else {
2156 update
2157 .remove_outgoing_requests
2158 .push(responder_id.to_proto());
2159 }
2160 for connection_id in pool.user_connection_ids(requester_id) {
2161 session.peer.send(connection_id, update.clone())?;
2162 }
2163
2164 // Update incoming contact requests of responder
2165 let mut update = proto::UpdateContacts::default();
2166 if contact_accepted {
2167 update.remove_contacts.push(requester_id.to_proto());
2168 } else {
2169 update
2170 .remove_incoming_requests
2171 .push(requester_id.to_proto());
2172 }
2173 for connection_id in pool.user_connection_ids(responder_id) {
2174 session.peer.send(connection_id, update.clone())?;
2175 if let Some(notification_id) = deleted_notification_id {
2176 session.peer.send(
2177 connection_id,
2178 proto::DeleteNotification {
2179 notification_id: notification_id.to_proto(),
2180 },
2181 )?;
2182 }
2183 }
2184
2185 response.send(proto::Ack {})?;
2186 Ok(())
2187}
2188
2189async fn create_channel(
2190 request: proto::CreateChannel,
2191 response: Response<proto::CreateChannel>,
2192 session: Session,
2193) -> Result<()> {
2194 let db = session.db().await;
2195
2196 let parent_id = request.parent_id.map(|id| ChannelId::from_proto(id));
2197 let CreateChannelResult {
2198 channel,
2199 participants_to_update,
2200 } = db
2201 .create_channel(&request.name, parent_id, session.user_id)
2202 .await?;
2203
2204 response.send(proto::CreateChannelResponse {
2205 channel: Some(channel.to_proto()),
2206 parent_id: request.parent_id,
2207 })?;
2208
2209 let connection_pool = session.connection_pool().await;
2210 for (user_id, channels) in participants_to_update {
2211 let update = build_channels_update(channels, vec![]);
2212 for connection_id in connection_pool.user_connection_ids(user_id) {
2213 if user_id == session.user_id {
2214 continue;
2215 }
2216 session.peer.send(connection_id, update.clone())?;
2217 }
2218 }
2219
2220 Ok(())
2221}
2222
2223async fn delete_channel(
2224 request: proto::DeleteChannel,
2225 response: Response<proto::DeleteChannel>,
2226 session: Session,
2227) -> Result<()> {
2228 let db = session.db().await;
2229
2230 let channel_id = request.channel_id;
2231 let (removed_channels, member_ids) = db
2232 .delete_channel(ChannelId::from_proto(channel_id), session.user_id)
2233 .await?;
2234 response.send(proto::Ack {})?;
2235
2236 // Notify members of removed channels
2237 let mut update = proto::UpdateChannels::default();
2238 update
2239 .delete_channels
2240 .extend(removed_channels.into_iter().map(|id| id.to_proto()));
2241
2242 let connection_pool = session.connection_pool().await;
2243 for member_id in member_ids {
2244 for connection_id in connection_pool.user_connection_ids(member_id) {
2245 session.peer.send(connection_id, update.clone())?;
2246 }
2247 }
2248
2249 Ok(())
2250}
2251
2252async fn invite_channel_member(
2253 request: proto::InviteChannelMember,
2254 response: Response<proto::InviteChannelMember>,
2255 session: Session,
2256) -> Result<()> {
2257 let db = session.db().await;
2258 let channel_id = ChannelId::from_proto(request.channel_id);
2259 let invitee_id = UserId::from_proto(request.user_id);
2260 let InviteMemberResult {
2261 channel,
2262 notifications,
2263 } = db
2264 .invite_channel_member(
2265 channel_id,
2266 invitee_id,
2267 session.user_id,
2268 request.role().into(),
2269 )
2270 .await?;
2271
2272 let update = proto::UpdateChannels {
2273 channel_invitations: vec![channel.to_proto()],
2274 ..Default::default()
2275 };
2276
2277 let connection_pool = session.connection_pool().await;
2278 for connection_id in connection_pool.user_connection_ids(invitee_id) {
2279 session.peer.send(connection_id, update.clone())?;
2280 }
2281
2282 send_notifications(&*connection_pool, &session.peer, notifications);
2283
2284 response.send(proto::Ack {})?;
2285 Ok(())
2286}
2287
2288async fn remove_channel_member(
2289 request: proto::RemoveChannelMember,
2290 response: Response<proto::RemoveChannelMember>,
2291 session: Session,
2292) -> Result<()> {
2293 let db = session.db().await;
2294 let channel_id = ChannelId::from_proto(request.channel_id);
2295 let member_id = UserId::from_proto(request.user_id);
2296
2297 let RemoveChannelMemberResult {
2298 membership_update,
2299 notification_id,
2300 } = db
2301 .remove_channel_member(channel_id, member_id, session.user_id)
2302 .await?;
2303
2304 let connection_pool = &session.connection_pool().await;
2305 notify_membership_updated(
2306 &connection_pool,
2307 membership_update,
2308 member_id,
2309 &session.peer,
2310 );
2311 for connection_id in connection_pool.user_connection_ids(member_id) {
2312 if let Some(notification_id) = notification_id {
2313 session
2314 .peer
2315 .send(
2316 connection_id,
2317 proto::DeleteNotification {
2318 notification_id: notification_id.to_proto(),
2319 },
2320 )
2321 .trace_err();
2322 }
2323 }
2324
2325 response.send(proto::Ack {})?;
2326 Ok(())
2327}
2328
2329async fn set_channel_visibility(
2330 request: proto::SetChannelVisibility,
2331 response: Response<proto::SetChannelVisibility>,
2332 session: Session,
2333) -> Result<()> {
2334 let db = session.db().await;
2335 let channel_id = ChannelId::from_proto(request.channel_id);
2336 let visibility = request.visibility().into();
2337
2338 let SetChannelVisibilityResult {
2339 participants_to_update,
2340 participants_to_remove,
2341 channels_to_remove,
2342 } = db
2343 .set_channel_visibility(channel_id, visibility, session.user_id)
2344 .await?;
2345
2346 let connection_pool = session.connection_pool().await;
2347 for (user_id, channels) in participants_to_update {
2348 let update = build_channels_update(channels, vec![]);
2349 for connection_id in connection_pool.user_connection_ids(user_id) {
2350 session.peer.send(connection_id, update.clone())?;
2351 }
2352 }
2353 for user_id in participants_to_remove {
2354 let update = proto::UpdateChannels {
2355 delete_channels: channels_to_remove.iter().map(|id| id.to_proto()).collect(),
2356 ..Default::default()
2357 };
2358 for connection_id in connection_pool.user_connection_ids(user_id) {
2359 session.peer.send(connection_id, update.clone())?;
2360 }
2361 }
2362
2363 response.send(proto::Ack {})?;
2364 Ok(())
2365}
2366
2367async fn set_channel_member_role(
2368 request: proto::SetChannelMemberRole,
2369 response: Response<proto::SetChannelMemberRole>,
2370 session: Session,
2371) -> Result<()> {
2372 let db = session.db().await;
2373 let channel_id = ChannelId::from_proto(request.channel_id);
2374 let member_id = UserId::from_proto(request.user_id);
2375 let result = db
2376 .set_channel_member_role(
2377 channel_id,
2378 session.user_id,
2379 member_id,
2380 request.role().into(),
2381 )
2382 .await?;
2383
2384 match result {
2385 db::SetMemberRoleResult::MembershipUpdated(membership_update) => {
2386 let connection_pool = session.connection_pool().await;
2387 notify_membership_updated(
2388 &connection_pool,
2389 membership_update,
2390 member_id,
2391 &session.peer,
2392 )
2393 }
2394 db::SetMemberRoleResult::InviteUpdated(channel) => {
2395 let update = proto::UpdateChannels {
2396 channel_invitations: vec![channel.to_proto()],
2397 ..Default::default()
2398 };
2399
2400 for connection_id in session
2401 .connection_pool()
2402 .await
2403 .user_connection_ids(member_id)
2404 {
2405 session.peer.send(connection_id, update.clone())?;
2406 }
2407 }
2408 }
2409
2410 response.send(proto::Ack {})?;
2411 Ok(())
2412}
2413
2414async fn rename_channel(
2415 request: proto::RenameChannel,
2416 response: Response<proto::RenameChannel>,
2417 session: Session,
2418) -> Result<()> {
2419 let db = session.db().await;
2420 let channel_id = ChannelId::from_proto(request.channel_id);
2421 let RenameChannelResult {
2422 channel,
2423 participants_to_update,
2424 } = db
2425 .rename_channel(channel_id, session.user_id, &request.name)
2426 .await?;
2427
2428 response.send(proto::RenameChannelResponse {
2429 channel: Some(channel.to_proto()),
2430 })?;
2431
2432 let connection_pool = session.connection_pool().await;
2433 for (user_id, channel) in participants_to_update {
2434 for connection_id in connection_pool.user_connection_ids(user_id) {
2435 let update = proto::UpdateChannels {
2436 channels: vec![channel.to_proto()],
2437 ..Default::default()
2438 };
2439
2440 session.peer.send(connection_id, update.clone())?;
2441 }
2442 }
2443
2444 Ok(())
2445}
2446
2447async fn move_channel(
2448 request: proto::MoveChannel,
2449 response: Response<proto::MoveChannel>,
2450 session: Session,
2451) -> Result<()> {
2452 let channel_id = ChannelId::from_proto(request.channel_id);
2453 let to = request.to.map(ChannelId::from_proto);
2454
2455 let result = session
2456 .db()
2457 .await
2458 .move_channel(channel_id, to, session.user_id)
2459 .await?;
2460
2461 notify_channel_moved(result, session).await?;
2462
2463 response.send(Ack {})?;
2464 Ok(())
2465}
2466
2467async fn notify_channel_moved(result: Option<MoveChannelResult>, session: Session) -> Result<()> {
2468 let Some(MoveChannelResult {
2469 participants_to_remove,
2470 participants_to_update,
2471 moved_channels,
2472 }) = result
2473 else {
2474 return Ok(());
2475 };
2476 let moved_channels: Vec<u64> = moved_channels.iter().map(|id| id.to_proto()).collect();
2477
2478 let connection_pool = session.connection_pool().await;
2479 for (user_id, channels) in participants_to_update {
2480 let mut update = build_channels_update(channels, vec![]);
2481 update.delete_channels = moved_channels.clone();
2482 for connection_id in connection_pool.user_connection_ids(user_id) {
2483 session.peer.send(connection_id, update.clone())?;
2484 }
2485 }
2486
2487 for user_id in participants_to_remove {
2488 let update = proto::UpdateChannels {
2489 delete_channels: moved_channels.clone(),
2490 ..Default::default()
2491 };
2492 for connection_id in connection_pool.user_connection_ids(user_id) {
2493 session.peer.send(connection_id, update.clone())?;
2494 }
2495 }
2496 Ok(())
2497}
2498
2499async fn get_channel_members(
2500 request: proto::GetChannelMembers,
2501 response: Response<proto::GetChannelMembers>,
2502 session: Session,
2503) -> Result<()> {
2504 let db = session.db().await;
2505 let channel_id = ChannelId::from_proto(request.channel_id);
2506 let members = db
2507 .get_channel_participant_details(channel_id, session.user_id)
2508 .await?;
2509 response.send(proto::GetChannelMembersResponse { members })?;
2510 Ok(())
2511}
2512
2513async fn respond_to_channel_invite(
2514 request: proto::RespondToChannelInvite,
2515 response: Response<proto::RespondToChannelInvite>,
2516 session: Session,
2517) -> Result<()> {
2518 let db = session.db().await;
2519 let channel_id = ChannelId::from_proto(request.channel_id);
2520 let RespondToChannelInvite {
2521 membership_update,
2522 notifications,
2523 } = db
2524 .respond_to_channel_invite(channel_id, session.user_id, request.accept)
2525 .await?;
2526
2527 let connection_pool = session.connection_pool().await;
2528 if let Some(membership_update) = membership_update {
2529 notify_membership_updated(
2530 &connection_pool,
2531 membership_update,
2532 session.user_id,
2533 &session.peer,
2534 );
2535 } else {
2536 let update = proto::UpdateChannels {
2537 remove_channel_invitations: vec![channel_id.to_proto()],
2538 ..Default::default()
2539 };
2540
2541 for connection_id in connection_pool.user_connection_ids(session.user_id) {
2542 session.peer.send(connection_id, update.clone())?;
2543 }
2544 };
2545
2546 send_notifications(&*connection_pool, &session.peer, notifications);
2547
2548 response.send(proto::Ack {})?;
2549
2550 Ok(())
2551}
2552
2553async fn join_channel(
2554 request: proto::JoinChannel,
2555 response: Response<proto::JoinChannel>,
2556 session: Session,
2557) -> Result<()> {
2558 let channel_id = ChannelId::from_proto(request.channel_id);
2559 join_channel_internal(channel_id, Box::new(response), session).await
2560}
2561
2562trait JoinChannelInternalResponse {
2563 fn send(self, result: proto::JoinRoomResponse) -> Result<()>;
2564}
2565impl JoinChannelInternalResponse for Response<proto::JoinChannel> {
2566 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
2567 Response::<proto::JoinChannel>::send(self, result)
2568 }
2569}
2570impl JoinChannelInternalResponse for Response<proto::JoinRoom> {
2571 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
2572 Response::<proto::JoinRoom>::send(self, result)
2573 }
2574}
2575
2576async fn join_channel_internal(
2577 channel_id: ChannelId,
2578 response: Box<impl JoinChannelInternalResponse>,
2579 session: Session,
2580) -> Result<()> {
2581 let joined_room = {
2582 leave_room_for_session(&session).await?;
2583 let db = session.db().await;
2584
2585 let (joined_room, membership_updated, role) = db
2586 .join_channel(
2587 channel_id,
2588 session.user_id,
2589 session.connection_id,
2590 RELEASE_CHANNEL_NAME.as_str(),
2591 )
2592 .await?;
2593
2594 let live_kit_connection_info = session.live_kit_client.as_ref().and_then(|live_kit| {
2595 let (can_publish, token) = if role == ChannelRole::Guest {
2596 (
2597 false,
2598 live_kit
2599 .guest_token(
2600 &joined_room.room.live_kit_room,
2601 &session.user_id.to_string(),
2602 )
2603 .trace_err()?,
2604 )
2605 } else {
2606 (
2607 true,
2608 live_kit
2609 .room_token(
2610 &joined_room.room.live_kit_room,
2611 &session.user_id.to_string(),
2612 )
2613 .trace_err()?,
2614 )
2615 };
2616
2617 Some(LiveKitConnectionInfo {
2618 server_url: live_kit.url().into(),
2619 token,
2620 can_publish,
2621 })
2622 });
2623
2624 response.send(proto::JoinRoomResponse {
2625 room: Some(joined_room.room.clone()),
2626 channel_id: joined_room.channel_id.map(|id| id.to_proto()),
2627 live_kit_connection_info,
2628 })?;
2629
2630 let connection_pool = session.connection_pool().await;
2631 if let Some(membership_updated) = membership_updated {
2632 notify_membership_updated(
2633 &connection_pool,
2634 membership_updated,
2635 session.user_id,
2636 &session.peer,
2637 );
2638 }
2639
2640 room_updated(&joined_room.room, &session.peer);
2641
2642 joined_room
2643 };
2644
2645 channel_updated(
2646 channel_id,
2647 &joined_room.room,
2648 &joined_room.channel_members,
2649 &session.peer,
2650 &*session.connection_pool().await,
2651 );
2652
2653 update_user_contacts(session.user_id, &session).await?;
2654 Ok(())
2655}
2656
2657async fn join_channel_buffer(
2658 request: proto::JoinChannelBuffer,
2659 response: Response<proto::JoinChannelBuffer>,
2660 session: Session,
2661) -> Result<()> {
2662 let db = session.db().await;
2663 let channel_id = ChannelId::from_proto(request.channel_id);
2664
2665 let open_response = db
2666 .join_channel_buffer(channel_id, session.user_id, session.connection_id)
2667 .await?;
2668
2669 let collaborators = open_response.collaborators.clone();
2670 response.send(open_response)?;
2671
2672 let update = UpdateChannelBufferCollaborators {
2673 channel_id: channel_id.to_proto(),
2674 collaborators: collaborators.clone(),
2675 };
2676 channel_buffer_updated(
2677 session.connection_id,
2678 collaborators
2679 .iter()
2680 .filter_map(|collaborator| Some(collaborator.peer_id?.into())),
2681 &update,
2682 &session.peer,
2683 );
2684
2685 Ok(())
2686}
2687
2688async fn update_channel_buffer(
2689 request: proto::UpdateChannelBuffer,
2690 session: Session,
2691) -> Result<()> {
2692 let db = session.db().await;
2693 let channel_id = ChannelId::from_proto(request.channel_id);
2694
2695 let (collaborators, non_collaborators, epoch, version) = db
2696 .update_channel_buffer(channel_id, session.user_id, &request.operations)
2697 .await?;
2698
2699 channel_buffer_updated(
2700 session.connection_id,
2701 collaborators,
2702 &proto::UpdateChannelBuffer {
2703 channel_id: channel_id.to_proto(),
2704 operations: request.operations,
2705 },
2706 &session.peer,
2707 );
2708
2709 let pool = &*session.connection_pool().await;
2710
2711 broadcast(
2712 None,
2713 non_collaborators
2714 .iter()
2715 .flat_map(|user_id| pool.user_connection_ids(*user_id)),
2716 |peer_id| {
2717 session.peer.send(
2718 peer_id.into(),
2719 proto::UpdateChannels {
2720 unseen_channel_buffer_changes: vec![proto::UnseenChannelBufferChange {
2721 channel_id: channel_id.to_proto(),
2722 epoch: epoch as u64,
2723 version: version.clone(),
2724 }],
2725 ..Default::default()
2726 },
2727 )
2728 },
2729 );
2730
2731 Ok(())
2732}
2733
2734async fn rejoin_channel_buffers(
2735 request: proto::RejoinChannelBuffers,
2736 response: Response<proto::RejoinChannelBuffers>,
2737 session: Session,
2738) -> Result<()> {
2739 let db = session.db().await;
2740 let buffers = db
2741 .rejoin_channel_buffers(&request.buffers, session.user_id, session.connection_id)
2742 .await?;
2743
2744 for rejoined_buffer in &buffers {
2745 let collaborators_to_notify = rejoined_buffer
2746 .buffer
2747 .collaborators
2748 .iter()
2749 .filter_map(|c| Some(c.peer_id?.into()));
2750 channel_buffer_updated(
2751 session.connection_id,
2752 collaborators_to_notify,
2753 &proto::UpdateChannelBufferCollaborators {
2754 channel_id: rejoined_buffer.buffer.channel_id,
2755 collaborators: rejoined_buffer.buffer.collaborators.clone(),
2756 },
2757 &session.peer,
2758 );
2759 }
2760
2761 response.send(proto::RejoinChannelBuffersResponse {
2762 buffers: buffers.into_iter().map(|b| b.buffer).collect(),
2763 })?;
2764
2765 Ok(())
2766}
2767
2768async fn leave_channel_buffer(
2769 request: proto::LeaveChannelBuffer,
2770 response: Response<proto::LeaveChannelBuffer>,
2771 session: Session,
2772) -> Result<()> {
2773 let db = session.db().await;
2774 let channel_id = ChannelId::from_proto(request.channel_id);
2775
2776 let left_buffer = db
2777 .leave_channel_buffer(channel_id, session.connection_id)
2778 .await?;
2779
2780 response.send(Ack {})?;
2781
2782 channel_buffer_updated(
2783 session.connection_id,
2784 left_buffer.connections,
2785 &proto::UpdateChannelBufferCollaborators {
2786 channel_id: channel_id.to_proto(),
2787 collaborators: left_buffer.collaborators,
2788 },
2789 &session.peer,
2790 );
2791
2792 Ok(())
2793}
2794
2795fn channel_buffer_updated<T: EnvelopedMessage>(
2796 sender_id: ConnectionId,
2797 collaborators: impl IntoIterator<Item = ConnectionId>,
2798 message: &T,
2799 peer: &Peer,
2800) {
2801 broadcast(Some(sender_id), collaborators.into_iter(), |peer_id| {
2802 peer.send(peer_id.into(), message.clone())
2803 });
2804}
2805
2806fn send_notifications(
2807 connection_pool: &ConnectionPool,
2808 peer: &Peer,
2809 notifications: db::NotificationBatch,
2810) {
2811 for (user_id, notification) in notifications {
2812 for connection_id in connection_pool.user_connection_ids(user_id) {
2813 if let Err(error) = peer.send(
2814 connection_id,
2815 proto::AddNotification {
2816 notification: Some(notification.clone()),
2817 },
2818 ) {
2819 tracing::error!(
2820 "failed to send notification to {:?} {}",
2821 connection_id,
2822 error
2823 );
2824 }
2825 }
2826 }
2827}
2828
2829async fn send_channel_message(
2830 request: proto::SendChannelMessage,
2831 response: Response<proto::SendChannelMessage>,
2832 session: Session,
2833) -> Result<()> {
2834 // Validate the message body.
2835 let body = request.body.trim().to_string();
2836 if body.len() > MAX_MESSAGE_LEN {
2837 return Err(anyhow!("message is too long"))?;
2838 }
2839 if body.is_empty() {
2840 return Err(anyhow!("message can't be blank"))?;
2841 }
2842
2843 // TODO: adjust mentions if body is trimmed
2844
2845 let timestamp = OffsetDateTime::now_utc();
2846 let nonce = request
2847 .nonce
2848 .ok_or_else(|| anyhow!("nonce can't be blank"))?;
2849
2850 let channel_id = ChannelId::from_proto(request.channel_id);
2851 let CreatedChannelMessage {
2852 message_id,
2853 participant_connection_ids,
2854 channel_members,
2855 notifications,
2856 } = session
2857 .db()
2858 .await
2859 .create_channel_message(
2860 channel_id,
2861 session.user_id,
2862 &body,
2863 &request.mentions,
2864 timestamp,
2865 nonce.clone().into(),
2866 )
2867 .await?;
2868 let message = proto::ChannelMessage {
2869 sender_id: session.user_id.to_proto(),
2870 id: message_id.to_proto(),
2871 body,
2872 mentions: request.mentions,
2873 timestamp: timestamp.unix_timestamp() as u64,
2874 nonce: Some(nonce),
2875 };
2876 broadcast(
2877 Some(session.connection_id),
2878 participant_connection_ids,
2879 |connection| {
2880 session.peer.send(
2881 connection,
2882 proto::ChannelMessageSent {
2883 channel_id: channel_id.to_proto(),
2884 message: Some(message.clone()),
2885 },
2886 )
2887 },
2888 );
2889 response.send(proto::SendChannelMessageResponse {
2890 message: Some(message),
2891 })?;
2892
2893 let pool = &*session.connection_pool().await;
2894 broadcast(
2895 None,
2896 channel_members
2897 .iter()
2898 .flat_map(|user_id| pool.user_connection_ids(*user_id)),
2899 |peer_id| {
2900 session.peer.send(
2901 peer_id.into(),
2902 proto::UpdateChannels {
2903 unseen_channel_messages: vec![proto::UnseenChannelMessage {
2904 channel_id: channel_id.to_proto(),
2905 message_id: message_id.to_proto(),
2906 }],
2907 ..Default::default()
2908 },
2909 )
2910 },
2911 );
2912 send_notifications(pool, &session.peer, notifications);
2913
2914 Ok(())
2915}
2916
2917async fn remove_channel_message(
2918 request: proto::RemoveChannelMessage,
2919 response: Response<proto::RemoveChannelMessage>,
2920 session: Session,
2921) -> Result<()> {
2922 let channel_id = ChannelId::from_proto(request.channel_id);
2923 let message_id = MessageId::from_proto(request.message_id);
2924 let connection_ids = session
2925 .db()
2926 .await
2927 .remove_channel_message(channel_id, message_id, session.user_id)
2928 .await?;
2929 broadcast(Some(session.connection_id), connection_ids, |connection| {
2930 session.peer.send(connection, request.clone())
2931 });
2932 response.send(proto::Ack {})?;
2933 Ok(())
2934}
2935
2936async fn acknowledge_channel_message(
2937 request: proto::AckChannelMessage,
2938 session: Session,
2939) -> Result<()> {
2940 let channel_id = ChannelId::from_proto(request.channel_id);
2941 let message_id = MessageId::from_proto(request.message_id);
2942 let notifications = session
2943 .db()
2944 .await
2945 .observe_channel_message(channel_id, session.user_id, message_id)
2946 .await?;
2947 send_notifications(
2948 &*session.connection_pool().await,
2949 &session.peer,
2950 notifications,
2951 );
2952 Ok(())
2953}
2954
2955async fn acknowledge_buffer_version(
2956 request: proto::AckBufferOperation,
2957 session: Session,
2958) -> Result<()> {
2959 let buffer_id = BufferId::from_proto(request.buffer_id);
2960 session
2961 .db()
2962 .await
2963 .observe_buffer_version(
2964 buffer_id,
2965 session.user_id,
2966 request.epoch as i32,
2967 &request.version,
2968 )
2969 .await?;
2970 Ok(())
2971}
2972
2973async fn join_channel_chat(
2974 request: proto::JoinChannelChat,
2975 response: Response<proto::JoinChannelChat>,
2976 session: Session,
2977) -> Result<()> {
2978 let channel_id = ChannelId::from_proto(request.channel_id);
2979
2980 let db = session.db().await;
2981 db.join_channel_chat(channel_id, session.connection_id, session.user_id)
2982 .await?;
2983 let messages = db
2984 .get_channel_messages(channel_id, session.user_id, MESSAGE_COUNT_PER_PAGE, None)
2985 .await?;
2986 response.send(proto::JoinChannelChatResponse {
2987 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
2988 messages,
2989 })?;
2990 Ok(())
2991}
2992
2993async fn leave_channel_chat(request: proto::LeaveChannelChat, session: Session) -> Result<()> {
2994 let channel_id = ChannelId::from_proto(request.channel_id);
2995 session
2996 .db()
2997 .await
2998 .leave_channel_chat(channel_id, session.connection_id, session.user_id)
2999 .await?;
3000 Ok(())
3001}
3002
3003async fn get_channel_messages(
3004 request: proto::GetChannelMessages,
3005 response: Response<proto::GetChannelMessages>,
3006 session: Session,
3007) -> Result<()> {
3008 let channel_id = ChannelId::from_proto(request.channel_id);
3009 let messages = session
3010 .db()
3011 .await
3012 .get_channel_messages(
3013 channel_id,
3014 session.user_id,
3015 MESSAGE_COUNT_PER_PAGE,
3016 Some(MessageId::from_proto(request.before_message_id)),
3017 )
3018 .await?;
3019 response.send(proto::GetChannelMessagesResponse {
3020 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
3021 messages,
3022 })?;
3023 Ok(())
3024}
3025
3026async fn get_channel_messages_by_id(
3027 request: proto::GetChannelMessagesById,
3028 response: Response<proto::GetChannelMessagesById>,
3029 session: Session,
3030) -> Result<()> {
3031 let message_ids = request
3032 .message_ids
3033 .iter()
3034 .map(|id| MessageId::from_proto(*id))
3035 .collect::<Vec<_>>();
3036 let messages = session
3037 .db()
3038 .await
3039 .get_channel_messages_by_id(session.user_id, &message_ids)
3040 .await?;
3041 response.send(proto::GetChannelMessagesResponse {
3042 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
3043 messages,
3044 })?;
3045 Ok(())
3046}
3047
3048async fn get_notifications(
3049 request: proto::GetNotifications,
3050 response: Response<proto::GetNotifications>,
3051 session: Session,
3052) -> Result<()> {
3053 let notifications = session
3054 .db()
3055 .await
3056 .get_notifications(
3057 session.user_id,
3058 NOTIFICATION_COUNT_PER_PAGE,
3059 request
3060 .before_id
3061 .map(|id| db::NotificationId::from_proto(id)),
3062 )
3063 .await?;
3064 response.send(proto::GetNotificationsResponse {
3065 done: notifications.len() < NOTIFICATION_COUNT_PER_PAGE,
3066 notifications,
3067 })?;
3068 Ok(())
3069}
3070
3071async fn mark_notification_as_read(
3072 request: proto::MarkNotificationRead,
3073 response: Response<proto::MarkNotificationRead>,
3074 session: Session,
3075) -> Result<()> {
3076 let database = &session.db().await;
3077 let notifications = database
3078 .mark_notification_as_read_by_id(
3079 session.user_id,
3080 NotificationId::from_proto(request.notification_id),
3081 )
3082 .await?;
3083 send_notifications(
3084 &*session.connection_pool().await,
3085 &session.peer,
3086 notifications,
3087 );
3088 response.send(proto::Ack {})?;
3089 Ok(())
3090}
3091
3092async fn get_private_user_info(
3093 _request: proto::GetPrivateUserInfo,
3094 response: Response<proto::GetPrivateUserInfo>,
3095 session: Session,
3096) -> Result<()> {
3097 let db = session.db().await;
3098
3099 let metrics_id = db.get_user_metrics_id(session.user_id).await?;
3100 let user = db
3101 .get_user_by_id(session.user_id)
3102 .await?
3103 .ok_or_else(|| anyhow!("user not found"))?;
3104 let flags = db.get_user_flags(session.user_id).await?;
3105
3106 response.send(proto::GetPrivateUserInfoResponse {
3107 metrics_id,
3108 staff: user.admin,
3109 flags,
3110 })?;
3111 Ok(())
3112}
3113
3114fn to_axum_message(message: TungsteniteMessage) -> AxumMessage {
3115 match message {
3116 TungsteniteMessage::Text(payload) => AxumMessage::Text(payload),
3117 TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload),
3118 TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload),
3119 TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload),
3120 TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
3121 code: frame.code.into(),
3122 reason: frame.reason,
3123 })),
3124 }
3125}
3126
3127fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
3128 match message {
3129 AxumMessage::Text(payload) => TungsteniteMessage::Text(payload),
3130 AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload),
3131 AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload),
3132 AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload),
3133 AxumMessage::Close(frame) => {
3134 TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
3135 code: frame.code.into(),
3136 reason: frame.reason,
3137 }))
3138 }
3139 }
3140}
3141
3142fn notify_membership_updated(
3143 connection_pool: &ConnectionPool,
3144 result: MembershipUpdated,
3145 user_id: UserId,
3146 peer: &Peer,
3147) {
3148 let mut update = build_channels_update(result.new_channels, vec![]);
3149 update.delete_channels = result
3150 .removed_channels
3151 .into_iter()
3152 .map(|id| id.to_proto())
3153 .collect();
3154 update.remove_channel_invitations = vec![result.channel_id.to_proto()];
3155
3156 for connection_id in connection_pool.user_connection_ids(user_id) {
3157 peer.send(connection_id, update.clone()).trace_err();
3158 }
3159}
3160
3161fn build_channels_update(
3162 channels: ChannelsForUser,
3163 channel_invites: Vec<db::Channel>,
3164) -> proto::UpdateChannels {
3165 let mut update = proto::UpdateChannels::default();
3166
3167 for channel in channels.channels {
3168 update.channels.push(channel.to_proto());
3169 }
3170
3171 update.unseen_channel_buffer_changes = channels.unseen_buffer_changes;
3172 update.unseen_channel_messages = channels.channel_messages;
3173
3174 for (channel_id, participants) in channels.channel_participants {
3175 update
3176 .channel_participants
3177 .push(proto::ChannelParticipants {
3178 channel_id: channel_id.to_proto(),
3179 participant_user_ids: participants.into_iter().map(|id| id.to_proto()).collect(),
3180 });
3181 }
3182
3183 for channel in channel_invites {
3184 update.channel_invitations.push(channel.to_proto());
3185 }
3186
3187 update
3188}
3189
3190fn build_initial_contacts_update(
3191 contacts: Vec<db::Contact>,
3192 pool: &ConnectionPool,
3193) -> proto::UpdateContacts {
3194 let mut update = proto::UpdateContacts::default();
3195
3196 for contact in contacts {
3197 match contact {
3198 db::Contact::Accepted { user_id, busy } => {
3199 update.contacts.push(contact_for_user(user_id, busy, &pool));
3200 }
3201 db::Contact::Outgoing { user_id } => update.outgoing_requests.push(user_id.to_proto()),
3202 db::Contact::Incoming { user_id } => {
3203 update
3204 .incoming_requests
3205 .push(proto::IncomingContactRequest {
3206 requester_id: user_id.to_proto(),
3207 })
3208 }
3209 }
3210 }
3211
3212 update
3213}
3214
3215fn contact_for_user(user_id: UserId, busy: bool, pool: &ConnectionPool) -> proto::Contact {
3216 proto::Contact {
3217 user_id: user_id.to_proto(),
3218 online: pool.is_user_online(user_id),
3219 busy,
3220 }
3221}
3222
3223fn room_updated(room: &proto::Room, peer: &Peer) {
3224 broadcast(
3225 None,
3226 room.participants
3227 .iter()
3228 .filter_map(|participant| Some(participant.peer_id?.into())),
3229 |peer_id| {
3230 peer.send(
3231 peer_id.into(),
3232 proto::RoomUpdated {
3233 room: Some(room.clone()),
3234 },
3235 )
3236 },
3237 );
3238}
3239
3240fn channel_updated(
3241 channel_id: ChannelId,
3242 room: &proto::Room,
3243 channel_members: &[UserId],
3244 peer: &Peer,
3245 pool: &ConnectionPool,
3246) {
3247 let participants = room
3248 .participants
3249 .iter()
3250 .map(|p| p.user_id)
3251 .collect::<Vec<_>>();
3252
3253 broadcast(
3254 None,
3255 channel_members
3256 .iter()
3257 .flat_map(|user_id| pool.user_connection_ids(*user_id)),
3258 |peer_id| {
3259 peer.send(
3260 peer_id.into(),
3261 proto::UpdateChannels {
3262 channel_participants: vec![proto::ChannelParticipants {
3263 channel_id: channel_id.to_proto(),
3264 participant_user_ids: participants.clone(),
3265 }],
3266 ..Default::default()
3267 },
3268 )
3269 },
3270 );
3271}
3272
3273async fn update_user_contacts(user_id: UserId, session: &Session) -> Result<()> {
3274 let db = session.db().await;
3275
3276 let contacts = db.get_contacts(user_id).await?;
3277 let busy = db.is_user_busy(user_id).await?;
3278
3279 let pool = session.connection_pool().await;
3280 let updated_contact = contact_for_user(user_id, busy, &pool);
3281 for contact in contacts {
3282 if let db::Contact::Accepted {
3283 user_id: contact_user_id,
3284 ..
3285 } = contact
3286 {
3287 for contact_conn_id in pool.user_connection_ids(contact_user_id) {
3288 session
3289 .peer
3290 .send(
3291 contact_conn_id,
3292 proto::UpdateContacts {
3293 contacts: vec![updated_contact.clone()],
3294 remove_contacts: Default::default(),
3295 incoming_requests: Default::default(),
3296 remove_incoming_requests: Default::default(),
3297 outgoing_requests: Default::default(),
3298 remove_outgoing_requests: Default::default(),
3299 },
3300 )
3301 .trace_err();
3302 }
3303 }
3304 }
3305 Ok(())
3306}
3307
3308async fn leave_room_for_session(session: &Session) -> Result<()> {
3309 let mut contacts_to_update = HashSet::default();
3310
3311 let room_id;
3312 let canceled_calls_to_user_ids;
3313 let live_kit_room;
3314 let delete_live_kit_room;
3315 let room;
3316 let channel_members;
3317 let channel_id;
3318
3319 if let Some(mut left_room) = session.db().await.leave_room(session.connection_id).await? {
3320 contacts_to_update.insert(session.user_id);
3321
3322 for project in left_room.left_projects.values() {
3323 project_left(project, session);
3324 }
3325
3326 room_id = RoomId::from_proto(left_room.room.id);
3327 canceled_calls_to_user_ids = mem::take(&mut left_room.canceled_calls_to_user_ids);
3328 live_kit_room = mem::take(&mut left_room.room.live_kit_room);
3329 delete_live_kit_room = left_room.deleted;
3330 room = mem::take(&mut left_room.room);
3331 channel_members = mem::take(&mut left_room.channel_members);
3332 channel_id = left_room.channel_id;
3333
3334 room_updated(&room, &session.peer);
3335 } else {
3336 return Ok(());
3337 }
3338
3339 if let Some(channel_id) = channel_id {
3340 channel_updated(
3341 channel_id,
3342 &room,
3343 &channel_members,
3344 &session.peer,
3345 &*session.connection_pool().await,
3346 );
3347 }
3348
3349 {
3350 let pool = session.connection_pool().await;
3351 for canceled_user_id in canceled_calls_to_user_ids {
3352 for connection_id in pool.user_connection_ids(canceled_user_id) {
3353 session
3354 .peer
3355 .send(
3356 connection_id,
3357 proto::CallCanceled {
3358 room_id: room_id.to_proto(),
3359 },
3360 )
3361 .trace_err();
3362 }
3363 contacts_to_update.insert(canceled_user_id);
3364 }
3365 }
3366
3367 for contact_user_id in contacts_to_update {
3368 update_user_contacts(contact_user_id, &session).await?;
3369 }
3370
3371 if let Some(live_kit) = session.live_kit_client.as_ref() {
3372 live_kit
3373 .remove_participant(live_kit_room.clone(), session.user_id.to_string())
3374 .await
3375 .trace_err();
3376
3377 if delete_live_kit_room {
3378 live_kit.delete_room(live_kit_room).await.trace_err();
3379 }
3380 }
3381
3382 Ok(())
3383}
3384
3385async fn leave_channel_buffers_for_session(session: &Session) -> Result<()> {
3386 let left_channel_buffers = session
3387 .db()
3388 .await
3389 .leave_channel_buffers(session.connection_id)
3390 .await?;
3391
3392 for left_buffer in left_channel_buffers {
3393 channel_buffer_updated(
3394 session.connection_id,
3395 left_buffer.connections,
3396 &proto::UpdateChannelBufferCollaborators {
3397 channel_id: left_buffer.channel_id.to_proto(),
3398 collaborators: left_buffer.collaborators,
3399 },
3400 &session.peer,
3401 );
3402 }
3403
3404 Ok(())
3405}
3406
3407fn project_left(project: &db::LeftProject, session: &Session) {
3408 for connection_id in &project.connection_ids {
3409 if project.host_user_id == session.user_id {
3410 session
3411 .peer
3412 .send(
3413 *connection_id,
3414 proto::UnshareProject {
3415 project_id: project.id.to_proto(),
3416 },
3417 )
3418 .trace_err();
3419 } else {
3420 session
3421 .peer
3422 .send(
3423 *connection_id,
3424 proto::RemoveProjectCollaborator {
3425 project_id: project.id.to_proto(),
3426 peer_id: Some(session.connection_id.into()),
3427 },
3428 )
3429 .trace_err();
3430 }
3431 }
3432}
3433
3434pub trait ResultExt {
3435 type Ok;
3436
3437 fn trace_err(self) -> Option<Self::Ok>;
3438}
3439
3440impl<T, E> ResultExt for Result<T, E>
3441where
3442 E: std::fmt::Debug,
3443{
3444 type Ok = T;
3445
3446 fn trace_err(self) -> Option<T> {
3447 match self {
3448 Ok(value) => Some(value),
3449 Err(error) => {
3450 tracing::error!("{:?}", error);
3451 None
3452 }
3453 }
3454 }
3455}