1mod connection_pool;
2
3use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};
4use crate::{
5 AppState, Error, Result, auth,
6 db::{
7 self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, Database,
8 InviteMemberResult, MembershipUpdated, NotificationId, ProjectId, RejoinedProject,
9 RemoveChannelMemberResult, RespondToChannelInvite, RoomId, ServerId, User, UserId,
10 },
11 executor::Executor,
12};
13use anyhow::{Context as _, anyhow, bail};
14use async_tungstenite::tungstenite::{
15 Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame,
16};
17use axum::headers::UserAgent;
18use axum::{
19 Extension, Router, TypedHeader,
20 body::Body,
21 extract::{
22 ConnectInfo, WebSocketUpgrade,
23 ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage},
24 },
25 headers::{Header, HeaderName},
26 http::StatusCode,
27 middleware,
28 response::IntoResponse,
29 routing::get,
30};
31use collections::{HashMap, HashSet};
32pub use connection_pool::{ConnectionPool, ZedVersion};
33use core::fmt::{self, Debug, Formatter};
34use futures::TryFutureExt as _;
35use reqwest_client::ReqwestClient;
36use rpc::proto::split_repository_update;
37use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi};
38use tracing::Span;
39
40use futures::{
41 FutureExt, SinkExt, StreamExt, TryStreamExt, channel::oneshot, future::BoxFuture,
42 stream::FuturesUnordered,
43};
44use prometheus::{IntGauge, register_int_gauge};
45use rpc::{
46 Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope,
47 proto::{
48 self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo,
49 RequestMessage, ShareProject, UpdateChannelBufferCollaborators,
50 },
51};
52use semantic_version::SemanticVersion;
53use serde::{Serialize, Serializer};
54use std::{
55 any::TypeId,
56 future::Future,
57 marker::PhantomData,
58 mem,
59 net::SocketAddr,
60 ops::{Deref, DerefMut},
61 rc::Rc,
62 sync::{
63 Arc, OnceLock,
64 atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
65 },
66 time::{Duration, Instant},
67};
68use tokio::sync::{Semaphore, watch};
69use tower::ServiceBuilder;
70use tracing::{
71 Instrument,
72 field::{self},
73 info_span, instrument,
74};
75
76pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30);
77
78// kubernetes gives terminated pods 10s to shutdown gracefully. After they're gone, we can clean up old resources.
79pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
80
81const NOTIFICATION_COUNT_PER_PAGE: usize = 50;
82const MAX_CONCURRENT_CONNECTIONS: usize = 512;
83
84static CONCURRENT_CONNECTIONS: AtomicUsize = AtomicUsize::new(0);
85
86const TOTAL_DURATION_MS: &str = "total_duration_ms";
87const PROCESSING_DURATION_MS: &str = "processing_duration_ms";
88const QUEUE_DURATION_MS: &str = "queue_duration_ms";
89const HOST_WAITING_MS: &str = "host_waiting_ms";
90
91type MessageHandler =
92 Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session, Span) -> BoxFuture<'static, ()>>;
93
94pub struct ConnectionGuard;
95
96impl ConnectionGuard {
97 pub fn try_acquire() -> Result<Self, ()> {
98 let current_connections = CONCURRENT_CONNECTIONS.fetch_add(1, SeqCst);
99 if current_connections >= MAX_CONCURRENT_CONNECTIONS {
100 CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst);
101 tracing::error!(
102 "too many concurrent connections: {}",
103 current_connections + 1
104 );
105 return Err(());
106 }
107 Ok(ConnectionGuard)
108 }
109}
110
111impl Drop for ConnectionGuard {
112 fn drop(&mut self) {
113 CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst);
114 }
115}
116
117struct Response<R> {
118 peer: Arc<Peer>,
119 receipt: Receipt<R>,
120 responded: Arc<AtomicBool>,
121}
122
123impl<R: RequestMessage> Response<R> {
124 fn send(self, payload: R::Response) -> Result<()> {
125 self.responded.store(true, SeqCst);
126 self.peer.respond(self.receipt, payload)?;
127 Ok(())
128 }
129}
130
131#[derive(Clone, Debug)]
132pub enum Principal {
133 User(User),
134 Impersonated { user: User, admin: User },
135}
136
137impl Principal {
138 fn update_span(&self, span: &tracing::Span) {
139 match &self {
140 Principal::User(user) => {
141 span.record("user_id", user.id.0);
142 span.record("login", &user.github_login);
143 }
144 Principal::Impersonated { user, admin } => {
145 span.record("user_id", user.id.0);
146 span.record("login", &user.github_login);
147 span.record("impersonator", &admin.github_login);
148 }
149 }
150 }
151}
152
153#[derive(Clone)]
154struct MessageContext {
155 session: Session,
156 span: tracing::Span,
157}
158
159impl Deref for MessageContext {
160 type Target = Session;
161
162 fn deref(&self) -> &Self::Target {
163 &self.session
164 }
165}
166
167impl MessageContext {
168 pub fn forward_request<T: RequestMessage>(
169 &self,
170 receiver_id: ConnectionId,
171 request: T,
172 ) -> impl Future<Output = anyhow::Result<T::Response>> {
173 let request_start_time = Instant::now();
174 let span = self.span.clone();
175 tracing::info!("start forwarding request");
176 self.peer
177 .forward_request(self.connection_id, receiver_id, request)
178 .inspect(move |_| {
179 span.record(
180 HOST_WAITING_MS,
181 request_start_time.elapsed().as_micros() as f64 / 1000.0,
182 );
183 })
184 .inspect_err(|_| tracing::error!("error forwarding request"))
185 .inspect_ok(|_| tracing::info!("finished forwarding request"))
186 }
187}
188
189#[derive(Clone)]
190struct Session {
191 principal: Principal,
192 connection_id: ConnectionId,
193 db: Arc<tokio::sync::Mutex<DbHandle>>,
194 peer: Arc<Peer>,
195 connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
196 app_state: Arc<AppState>,
197 supermaven_client: Option<Arc<SupermavenAdminApi>>,
198 /// The GeoIP country code for the user.
199 #[allow(unused)]
200 geoip_country_code: Option<String>,
201 #[allow(unused)]
202 system_id: Option<String>,
203 _executor: Executor,
204}
205
206impl Session {
207 async fn db(&self) -> tokio::sync::MutexGuard<'_, DbHandle> {
208 #[cfg(test)]
209 tokio::task::yield_now().await;
210 let guard = self.db.lock().await;
211 #[cfg(test)]
212 tokio::task::yield_now().await;
213 guard
214 }
215
216 async fn connection_pool(&self) -> ConnectionPoolGuard<'_> {
217 #[cfg(test)]
218 tokio::task::yield_now().await;
219 let guard = self.connection_pool.lock();
220 ConnectionPoolGuard {
221 guard,
222 _not_send: PhantomData,
223 }
224 }
225
226 fn is_staff(&self) -> bool {
227 match &self.principal {
228 Principal::User(user) => user.admin,
229 Principal::Impersonated { .. } => true,
230 }
231 }
232
233 fn user_id(&self) -> UserId {
234 match &self.principal {
235 Principal::User(user) => user.id,
236 Principal::Impersonated { user, .. } => user.id,
237 }
238 }
239
240 pub fn email(&self) -> Option<String> {
241 match &self.principal {
242 Principal::User(user) => user.email_address.clone(),
243 Principal::Impersonated { user, .. } => user.email_address.clone(),
244 }
245 }
246}
247
248impl Debug for Session {
249 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
250 let mut result = f.debug_struct("Session");
251 match &self.principal {
252 Principal::User(user) => {
253 result.field("user", &user.github_login);
254 }
255 Principal::Impersonated { user, admin } => {
256 result.field("user", &user.github_login);
257 result.field("impersonator", &admin.github_login);
258 }
259 }
260 result.field("connection_id", &self.connection_id).finish()
261 }
262}
263
264struct DbHandle(Arc<Database>);
265
266impl Deref for DbHandle {
267 type Target = Database;
268
269 fn deref(&self) -> &Self::Target {
270 self.0.as_ref()
271 }
272}
273
274pub struct Server {
275 id: parking_lot::Mutex<ServerId>,
276 peer: Arc<Peer>,
277 pub(crate) connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
278 app_state: Arc<AppState>,
279 handlers: HashMap<TypeId, MessageHandler>,
280 teardown: watch::Sender<bool>,
281}
282
283pub(crate) struct ConnectionPoolGuard<'a> {
284 guard: parking_lot::MutexGuard<'a, ConnectionPool>,
285 _not_send: PhantomData<Rc<()>>,
286}
287
288#[derive(Serialize)]
289pub struct ServerSnapshot<'a> {
290 peer: &'a Peer,
291 #[serde(serialize_with = "serialize_deref")]
292 connection_pool: ConnectionPoolGuard<'a>,
293}
294
295pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
296where
297 S: Serializer,
298 T: Deref<Target = U>,
299 U: Serialize,
300{
301 Serialize::serialize(value.deref(), serializer)
302}
303
304impl Server {
305 pub fn new(id: ServerId, app_state: Arc<AppState>) -> Arc<Self> {
306 let mut server = Self {
307 id: parking_lot::Mutex::new(id),
308 peer: Peer::new(id.0 as u32),
309 app_state,
310 connection_pool: Default::default(),
311 handlers: Default::default(),
312 teardown: watch::channel(false).0,
313 };
314
315 server
316 .add_request_handler(ping)
317 .add_request_handler(create_room)
318 .add_request_handler(join_room)
319 .add_request_handler(rejoin_room)
320 .add_request_handler(leave_room)
321 .add_request_handler(set_room_participant_role)
322 .add_request_handler(call)
323 .add_request_handler(cancel_call)
324 .add_message_handler(decline_call)
325 .add_request_handler(update_participant_location)
326 .add_request_handler(share_project)
327 .add_message_handler(unshare_project)
328 .add_request_handler(join_project)
329 .add_message_handler(leave_project)
330 .add_request_handler(update_project)
331 .add_request_handler(update_worktree)
332 .add_request_handler(update_repository)
333 .add_request_handler(remove_repository)
334 .add_message_handler(start_language_server)
335 .add_message_handler(update_language_server)
336 .add_message_handler(update_diagnostic_summary)
337 .add_message_handler(update_worktree_settings)
338 .add_request_handler(forward_read_only_project_request::<proto::FindSearchCandidates>)
339 .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>)
340 .add_request_handler(forward_read_only_project_request::<proto::GetDocumentSymbols>)
341 .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>)
342 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)
343 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)
344 .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)
345 .add_request_handler(forward_read_only_project_request::<proto::InlayHints>)
346 .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>)
347 .add_request_handler(forward_read_only_project_request::<proto::GetColorPresentation>)
348 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)
349 .add_request_handler(forward_read_only_project_request::<proto::GitGetBranches>)
350 .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>)
351 .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>)
352 .add_request_handler(forward_read_only_project_request::<proto::LspExtExpandMacro>)
353 .add_request_handler(forward_read_only_project_request::<proto::LspExtOpenDocs>)
354 .add_request_handler(forward_mutating_project_request::<proto::LspExtRunnables>)
355 .add_request_handler(
356 forward_read_only_project_request::<proto::LspExtSwitchSourceHeader>,
357 )
358 .add_request_handler(forward_read_only_project_request::<proto::LspExtGoToParentModule>)
359 .add_request_handler(forward_read_only_project_request::<proto::LspExtCancelFlycheck>)
360 .add_request_handler(forward_read_only_project_request::<proto::LspExtRunFlycheck>)
361 .add_request_handler(forward_read_only_project_request::<proto::LspExtClearFlycheck>)
362 .add_request_handler(
363 forward_mutating_project_request::<proto::RegisterBufferWithLanguageServers>,
364 )
365 .add_request_handler(forward_mutating_project_request::<proto::UpdateGitBranch>)
366 .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>)
367 .add_request_handler(
368 forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>,
369 )
370 .add_request_handler(forward_mutating_project_request::<proto::OpenNewBuffer>)
371 .add_request_handler(
372 forward_mutating_project_request::<proto::ResolveCompletionDocumentation>,
373 )
374 .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>)
375 .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>)
376 .add_request_handler(forward_mutating_project_request::<proto::PerformRename>)
377 .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>)
378 .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeActionKind>)
379 .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>)
380 .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>)
381 .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>)
382 .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>)
383 .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>)
384 .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>)
385 .add_request_handler(
386 forward_mutating_project_request::<proto::ExpandAllForProjectEntry>,
387 )
388 .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>)
389 .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)
390 .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)
391 .add_request_handler(lsp_query)
392 .add_message_handler(broadcast_project_message_from_host::<proto::LspQueryResponse>)
393 .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>)
394 .add_request_handler(forward_mutating_project_request::<proto::StopLanguageServers>)
395 .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>)
396 .add_message_handler(create_buffer_for_peer)
397 .add_request_handler(update_buffer)
398 .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>)
399 .add_message_handler(broadcast_project_message_from_host::<proto::RefreshCodeLens>)
400 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>)
401 .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>)
402 .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>)
403 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBases>)
404 .add_message_handler(
405 broadcast_project_message_from_host::<proto::PullWorkspaceDiagnostics>,
406 )
407 .add_request_handler(get_users)
408 .add_request_handler(fuzzy_search_users)
409 .add_request_handler(request_contact)
410 .add_request_handler(remove_contact)
411 .add_request_handler(respond_to_contact_request)
412 .add_message_handler(subscribe_to_channels)
413 .add_request_handler(create_channel)
414 .add_request_handler(delete_channel)
415 .add_request_handler(invite_channel_member)
416 .add_request_handler(remove_channel_member)
417 .add_request_handler(set_channel_member_role)
418 .add_request_handler(set_channel_visibility)
419 .add_request_handler(rename_channel)
420 .add_request_handler(join_channel_buffer)
421 .add_request_handler(leave_channel_buffer)
422 .add_message_handler(update_channel_buffer)
423 .add_request_handler(rejoin_channel_buffers)
424 .add_request_handler(get_channel_members)
425 .add_request_handler(respond_to_channel_invite)
426 .add_request_handler(join_channel)
427 .add_request_handler(join_channel_chat)
428 .add_message_handler(leave_channel_chat)
429 .add_request_handler(send_channel_message)
430 .add_request_handler(remove_channel_message)
431 .add_request_handler(update_channel_message)
432 .add_request_handler(get_channel_messages)
433 .add_request_handler(get_channel_messages_by_id)
434 .add_request_handler(get_notifications)
435 .add_request_handler(mark_notification_as_read)
436 .add_request_handler(move_channel)
437 .add_request_handler(reorder_channel)
438 .add_request_handler(follow)
439 .add_message_handler(unfollow)
440 .add_message_handler(update_followers)
441 .add_message_handler(acknowledge_channel_message)
442 .add_message_handler(acknowledge_buffer_version)
443 .add_request_handler(get_supermaven_api_key)
444 .add_request_handler(forward_mutating_project_request::<proto::OpenContext>)
445 .add_request_handler(forward_mutating_project_request::<proto::CreateContext>)
446 .add_request_handler(forward_mutating_project_request::<proto::SynchronizeContexts>)
447 .add_request_handler(forward_mutating_project_request::<proto::Stage>)
448 .add_request_handler(forward_mutating_project_request::<proto::Unstage>)
449 .add_request_handler(forward_mutating_project_request::<proto::Stash>)
450 .add_request_handler(forward_mutating_project_request::<proto::StashPop>)
451 .add_request_handler(forward_mutating_project_request::<proto::StashDrop>)
452 .add_request_handler(forward_mutating_project_request::<proto::Commit>)
453 .add_request_handler(forward_mutating_project_request::<proto::GitInit>)
454 .add_request_handler(forward_read_only_project_request::<proto::GetRemotes>)
455 .add_request_handler(forward_read_only_project_request::<proto::GitShow>)
456 .add_request_handler(forward_read_only_project_request::<proto::LoadCommitDiff>)
457 .add_request_handler(forward_read_only_project_request::<proto::GitReset>)
458 .add_request_handler(forward_read_only_project_request::<proto::GitCheckoutFiles>)
459 .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>)
460 .add_request_handler(forward_mutating_project_request::<proto::ToggleBreakpoint>)
461 .add_message_handler(broadcast_project_message_from_host::<proto::BreakpointsForFile>)
462 .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>)
463 .add_request_handler(forward_mutating_project_request::<proto::GitDiff>)
464 .add_request_handler(forward_mutating_project_request::<proto::GitCreateBranch>)
465 .add_request_handler(forward_mutating_project_request::<proto::GitChangeBranch>)
466 .add_request_handler(forward_mutating_project_request::<proto::CheckForPushedCommits>)
467 .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>)
468 .add_message_handler(update_context)
469 .add_request_handler(forward_mutating_project_request::<proto::ToggleLspLogs>)
470 .add_message_handler(broadcast_project_message_from_host::<proto::LanguageServerLog>);
471
472 Arc::new(server)
473 }
474
475 pub async fn start(&self) -> Result<()> {
476 let server_id = *self.id.lock();
477 let app_state = self.app_state.clone();
478 let peer = self.peer.clone();
479 let timeout = self.app_state.executor.sleep(CLEANUP_TIMEOUT);
480 let pool = self.connection_pool.clone();
481 let livekit_client = self.app_state.livekit_client.clone();
482
483 let span = info_span!("start server");
484 self.app_state.executor.spawn_detached(
485 async move {
486 tracing::info!("waiting for cleanup timeout");
487 timeout.await;
488 tracing::info!("cleanup timeout expired, retrieving stale rooms");
489
490 app_state
491 .db
492 .delete_stale_channel_chat_participants(
493 &app_state.config.zed_environment,
494 server_id,
495 )
496 .await
497 .trace_err();
498
499 if let Some((room_ids, channel_ids)) = app_state
500 .db
501 .stale_server_resource_ids(&app_state.config.zed_environment, server_id)
502 .await
503 .trace_err()
504 {
505 tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms");
506 tracing::info!(
507 stale_channel_buffer_count = channel_ids.len(),
508 "retrieved stale channel buffers"
509 );
510
511 for channel_id in channel_ids {
512 if let Some(refreshed_channel_buffer) = app_state
513 .db
514 .clear_stale_channel_buffer_collaborators(channel_id, server_id)
515 .await
516 .trace_err()
517 {
518 for connection_id in refreshed_channel_buffer.connection_ids {
519 peer.send(
520 connection_id,
521 proto::UpdateChannelBufferCollaborators {
522 channel_id: channel_id.to_proto(),
523 collaborators: refreshed_channel_buffer
524 .collaborators
525 .clone(),
526 },
527 )
528 .trace_err();
529 }
530 }
531 }
532
533 for room_id in room_ids {
534 let mut contacts_to_update = HashSet::default();
535 let mut canceled_calls_to_user_ids = Vec::new();
536 let mut livekit_room = String::new();
537 let mut delete_livekit_room = false;
538
539 if let Some(mut refreshed_room) = app_state
540 .db
541 .clear_stale_room_participants(room_id, server_id)
542 .await
543 .trace_err()
544 {
545 tracing::info!(
546 room_id = room_id.0,
547 new_participant_count = refreshed_room.room.participants.len(),
548 "refreshed room"
549 );
550 room_updated(&refreshed_room.room, &peer);
551 if let Some(channel) = refreshed_room.channel.as_ref() {
552 channel_updated(channel, &refreshed_room.room, &peer, &pool.lock());
553 }
554 contacts_to_update
555 .extend(refreshed_room.stale_participant_user_ids.iter().copied());
556 contacts_to_update
557 .extend(refreshed_room.canceled_calls_to_user_ids.iter().copied());
558 canceled_calls_to_user_ids =
559 mem::take(&mut refreshed_room.canceled_calls_to_user_ids);
560 livekit_room = mem::take(&mut refreshed_room.room.livekit_room);
561 delete_livekit_room = refreshed_room.room.participants.is_empty();
562 }
563
564 {
565 let pool = pool.lock();
566 for canceled_user_id in canceled_calls_to_user_ids {
567 for connection_id in pool.user_connection_ids(canceled_user_id) {
568 peer.send(
569 connection_id,
570 proto::CallCanceled {
571 room_id: room_id.to_proto(),
572 },
573 )
574 .trace_err();
575 }
576 }
577 }
578
579 for user_id in contacts_to_update {
580 let busy = app_state.db.is_user_busy(user_id).await.trace_err();
581 let contacts = app_state.db.get_contacts(user_id).await.trace_err();
582 if let Some((busy, contacts)) = busy.zip(contacts) {
583 let pool = pool.lock();
584 let updated_contact = contact_for_user(user_id, busy, &pool);
585 for contact in contacts {
586 if let db::Contact::Accepted {
587 user_id: contact_user_id,
588 ..
589 } = contact
590 {
591 for contact_conn_id in
592 pool.user_connection_ids(contact_user_id)
593 {
594 peer.send(
595 contact_conn_id,
596 proto::UpdateContacts {
597 contacts: vec![updated_contact.clone()],
598 remove_contacts: Default::default(),
599 incoming_requests: Default::default(),
600 remove_incoming_requests: Default::default(),
601 outgoing_requests: Default::default(),
602 remove_outgoing_requests: Default::default(),
603 },
604 )
605 .trace_err();
606 }
607 }
608 }
609 }
610 }
611
612 if let Some(live_kit) = livekit_client.as_ref()
613 && delete_livekit_room
614 {
615 live_kit.delete_room(livekit_room).await.trace_err();
616 }
617 }
618 }
619
620 app_state
621 .db
622 .delete_stale_channel_chat_participants(
623 &app_state.config.zed_environment,
624 server_id,
625 )
626 .await
627 .trace_err();
628
629 app_state
630 .db
631 .clear_old_worktree_entries(server_id)
632 .await
633 .trace_err();
634
635 app_state
636 .db
637 .delete_stale_servers(&app_state.config.zed_environment, server_id)
638 .await
639 .trace_err();
640 }
641 .instrument(span),
642 );
643 Ok(())
644 }
645
646 pub fn teardown(&self) {
647 self.peer.teardown();
648 self.connection_pool.lock().reset();
649 let _ = self.teardown.send(true);
650 }
651
652 #[cfg(test)]
653 pub fn reset(&self, id: ServerId) {
654 self.teardown();
655 *self.id.lock() = id;
656 self.peer.reset(id.0 as u32);
657 let _ = self.teardown.send(false);
658 }
659
660 #[cfg(test)]
661 pub fn id(&self) -> ServerId {
662 *self.id.lock()
663 }
664
665 fn add_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
666 where
667 F: 'static + Send + Sync + Fn(TypedEnvelope<M>, MessageContext) -> Fut,
668 Fut: 'static + Send + Future<Output = Result<()>>,
669 M: EnvelopedMessage,
670 {
671 let prev_handler = self.handlers.insert(
672 TypeId::of::<M>(),
673 Box::new(move |envelope, session, span| {
674 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
675 let received_at = envelope.received_at;
676 tracing::info!("message received");
677 let start_time = Instant::now();
678 let future = (handler)(
679 *envelope,
680 MessageContext {
681 session,
682 span: span.clone(),
683 },
684 );
685 async move {
686 let result = future.await;
687 let total_duration_ms = received_at.elapsed().as_micros() as f64 / 1000.0;
688 let processing_duration_ms = start_time.elapsed().as_micros() as f64 / 1000.0;
689 let queue_duration_ms = total_duration_ms - processing_duration_ms;
690 span.record(TOTAL_DURATION_MS, total_duration_ms);
691 span.record(PROCESSING_DURATION_MS, processing_duration_ms);
692 span.record(QUEUE_DURATION_MS, queue_duration_ms);
693 match result {
694 Err(error) => {
695 tracing::error!(?error, "error handling message")
696 }
697 Ok(()) => tracing::info!("finished handling message"),
698 }
699 }
700 .boxed()
701 }),
702 );
703 if prev_handler.is_some() {
704 panic!("registered a handler for the same message twice");
705 }
706 self
707 }
708
709 fn add_message_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
710 where
711 F: 'static + Send + Sync + Fn(M, MessageContext) -> Fut,
712 Fut: 'static + Send + Future<Output = Result<()>>,
713 M: EnvelopedMessage,
714 {
715 self.add_handler(move |envelope, session| handler(envelope.payload, session));
716 self
717 }
718
719 fn add_request_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
720 where
721 F: 'static + Send + Sync + Fn(M, Response<M>, MessageContext) -> Fut,
722 Fut: Send + Future<Output = Result<()>>,
723 M: RequestMessage,
724 {
725 let handler = Arc::new(handler);
726 self.add_handler(move |envelope, session| {
727 let receipt = envelope.receipt();
728 let handler = handler.clone();
729 async move {
730 let peer = session.peer.clone();
731 let responded = Arc::new(AtomicBool::default());
732 let response = Response {
733 peer: peer.clone(),
734 responded: responded.clone(),
735 receipt,
736 };
737 match (handler)(envelope.payload, response, session).await {
738 Ok(()) => {
739 if responded.load(std::sync::atomic::Ordering::SeqCst) {
740 Ok(())
741 } else {
742 Err(anyhow!("handler did not send a response"))?
743 }
744 }
745 Err(error) => {
746 let proto_err = match &error {
747 Error::Internal(err) => err.to_proto(),
748 _ => ErrorCode::Internal.message(format!("{error}")).to_proto(),
749 };
750 peer.respond_with_error(receipt, proto_err)?;
751 Err(error)
752 }
753 }
754 }
755 })
756 }
757
758 pub fn handle_connection(
759 self: &Arc<Self>,
760 connection: Connection,
761 address: String,
762 principal: Principal,
763 zed_version: ZedVersion,
764 release_channel: Option<String>,
765 user_agent: Option<String>,
766 geoip_country_code: Option<String>,
767 system_id: Option<String>,
768 send_connection_id: Option<oneshot::Sender<ConnectionId>>,
769 executor: Executor,
770 connection_guard: Option<ConnectionGuard>,
771 ) -> impl Future<Output = ()> + use<> {
772 let this = self.clone();
773 let span = info_span!("handle connection", %address,
774 connection_id=field::Empty,
775 user_id=field::Empty,
776 login=field::Empty,
777 impersonator=field::Empty,
778 user_agent=field::Empty,
779 geoip_country_code=field::Empty,
780 release_channel=field::Empty,
781 );
782 principal.update_span(&span);
783 if let Some(user_agent) = user_agent {
784 span.record("user_agent", user_agent);
785 }
786 if let Some(release_channel) = release_channel {
787 span.record("release_channel", release_channel);
788 }
789
790 if let Some(country_code) = geoip_country_code.as_ref() {
791 span.record("geoip_country_code", country_code);
792 }
793
794 let mut teardown = self.teardown.subscribe();
795 async move {
796 if *teardown.borrow() {
797 tracing::error!("server is tearing down");
798 return;
799 }
800
801 let (connection_id, handle_io, mut incoming_rx) =
802 this.peer.add_connection(connection, {
803 let executor = executor.clone();
804 move |duration| executor.sleep(duration)
805 });
806 tracing::Span::current().record("connection_id", format!("{}", connection_id));
807
808 tracing::info!("connection opened");
809
810 let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
811 let http_client = match ReqwestClient::user_agent(&user_agent) {
812 Ok(http_client) => Arc::new(http_client),
813 Err(error) => {
814 tracing::error!(?error, "failed to create HTTP client");
815 return;
816 }
817 };
818
819 let supermaven_client = this.app_state.config.supermaven_admin_api_key.clone().map(
820 |supermaven_admin_api_key| {
821 Arc::new(SupermavenAdminApi::new(
822 supermaven_admin_api_key.to_string(),
823 http_client.clone(),
824 ))
825 },
826 );
827
828 let session = Session {
829 principal: principal.clone(),
830 connection_id,
831 db: Arc::new(tokio::sync::Mutex::new(DbHandle(this.app_state.db.clone()))),
832 peer: this.peer.clone(),
833 connection_pool: this.connection_pool.clone(),
834 app_state: this.app_state.clone(),
835 geoip_country_code,
836 system_id,
837 _executor: executor.clone(),
838 supermaven_client,
839 };
840
841 if let Err(error) = this
842 .send_initial_client_update(
843 connection_id,
844 zed_version,
845 send_connection_id,
846 &session,
847 )
848 .await
849 {
850 tracing::error!(?error, "failed to send initial client update");
851 return;
852 }
853 drop(connection_guard);
854
855 let handle_io = handle_io.fuse();
856 futures::pin_mut!(handle_io);
857
858 // Handlers for foreground messages are pushed into the following `FuturesUnordered`.
859 // This prevents deadlocks when e.g., client A performs a request to client B and
860 // client B performs a request to client A. If both clients stop processing further
861 // messages until their respective request completes, they won't have a chance to
862 // respond to the other client's request and cause a deadlock.
863 //
864 // This arrangement ensures we will attempt to process earlier messages first, but fall
865 // back to processing messages arrived later in the spirit of making progress.
866 const MAX_CONCURRENT_HANDLERS: usize = 256;
867 let mut foreground_message_handlers = FuturesUnordered::new();
868 let concurrent_handlers = Arc::new(Semaphore::new(MAX_CONCURRENT_HANDLERS));
869 let get_concurrent_handlers = {
870 let concurrent_handlers = concurrent_handlers.clone();
871 move || MAX_CONCURRENT_HANDLERS - concurrent_handlers.available_permits()
872 };
873 loop {
874 let next_message = async {
875 let permit = concurrent_handlers.clone().acquire_owned().await.unwrap();
876 let message = incoming_rx.next().await;
877 // Cache the concurrent_handlers here, so that we know what the
878 // queue looks like as each handler starts
879 (permit, message, get_concurrent_handlers())
880 }
881 .fuse();
882 futures::pin_mut!(next_message);
883 futures::select_biased! {
884 _ = teardown.changed().fuse() => return,
885 result = handle_io => {
886 if let Err(error) = result {
887 tracing::error!(?error, "error handling I/O");
888 }
889 break;
890 }
891 _ = foreground_message_handlers.next() => {}
892 next_message = next_message => {
893 let (permit, message, concurrent_handlers) = next_message;
894 if let Some(message) = message {
895 let type_name = message.payload_type_name();
896 // note: we copy all the fields from the parent span so we can query them in the logs.
897 // (https://github.com/tokio-rs/tracing/issues/2670).
898 let span = tracing::info_span!("receive message",
899 %connection_id,
900 %address,
901 type_name,
902 concurrent_handlers,
903 user_id=field::Empty,
904 login=field::Empty,
905 impersonator=field::Empty,
906 lsp_query_request=field::Empty,
907 release_channel=field::Empty,
908 { TOTAL_DURATION_MS }=field::Empty,
909 { PROCESSING_DURATION_MS }=field::Empty,
910 { QUEUE_DURATION_MS }=field::Empty,
911 { HOST_WAITING_MS }=field::Empty
912 );
913 principal.update_span(&span);
914 let span_enter = span.enter();
915 if let Some(handler) = this.handlers.get(&message.payload_type_id()) {
916 let is_background = message.is_background();
917 let handle_message = (handler)(message, session.clone(), span.clone());
918 drop(span_enter);
919
920 let handle_message = async move {
921 handle_message.await;
922 drop(permit);
923 }.instrument(span);
924 if is_background {
925 executor.spawn_detached(handle_message);
926 } else {
927 foreground_message_handlers.push(handle_message);
928 }
929 } else {
930 tracing::error!("no message handler");
931 }
932 } else {
933 tracing::info!("connection closed");
934 break;
935 }
936 }
937 }
938 }
939
940 drop(foreground_message_handlers);
941 let concurrent_handlers = get_concurrent_handlers();
942 tracing::info!(concurrent_handlers, "signing out");
943 if let Err(error) = connection_lost(session, teardown, executor).await {
944 tracing::error!(?error, "error signing out");
945 }
946 }
947 .instrument(span)
948 }
949
950 async fn send_initial_client_update(
951 &self,
952 connection_id: ConnectionId,
953 zed_version: ZedVersion,
954 mut send_connection_id: Option<oneshot::Sender<ConnectionId>>,
955 session: &Session,
956 ) -> Result<()> {
957 self.peer.send(
958 connection_id,
959 proto::Hello {
960 peer_id: Some(connection_id.into()),
961 },
962 )?;
963 tracing::info!("sent hello message");
964 if let Some(send_connection_id) = send_connection_id.take() {
965 let _ = send_connection_id.send(connection_id);
966 }
967
968 match &session.principal {
969 Principal::User(user) | Principal::Impersonated { user, admin: _ } => {
970 if !user.connected_once {
971 self.peer.send(connection_id, proto::ShowContacts {})?;
972 self.app_state
973 .db
974 .set_user_connected_once(user.id, true)
975 .await?;
976 }
977
978 let contacts = self.app_state.db.get_contacts(user.id).await?;
979
980 {
981 let mut pool = self.connection_pool.lock();
982 pool.add_connection(connection_id, user.id, user.admin, zed_version);
983 self.peer.send(
984 connection_id,
985 build_initial_contacts_update(contacts, &pool),
986 )?;
987 }
988
989 if should_auto_subscribe_to_channels(zed_version) {
990 subscribe_user_to_channels(user.id, session).await?;
991 }
992
993 if let Some(incoming_call) =
994 self.app_state.db.incoming_call_for_user(user.id).await?
995 {
996 self.peer.send(connection_id, incoming_call)?;
997 }
998
999 update_user_contacts(user.id, session).await?;
1000 }
1001 }
1002
1003 Ok(())
1004 }
1005
1006 pub async fn invite_code_redeemed(
1007 self: &Arc<Self>,
1008 inviter_id: UserId,
1009 invitee_id: UserId,
1010 ) -> Result<()> {
1011 if let Some(user) = self.app_state.db.get_user_by_id(inviter_id).await?
1012 && let Some(code) = &user.invite_code
1013 {
1014 let pool = self.connection_pool.lock();
1015 let invitee_contact = contact_for_user(invitee_id, false, &pool);
1016 for connection_id in pool.user_connection_ids(inviter_id) {
1017 self.peer.send(
1018 connection_id,
1019 proto::UpdateContacts {
1020 contacts: vec![invitee_contact.clone()],
1021 ..Default::default()
1022 },
1023 )?;
1024 self.peer.send(
1025 connection_id,
1026 proto::UpdateInviteInfo {
1027 url: format!("{}{}", self.app_state.config.invite_link_prefix, &code),
1028 count: user.invite_count as u32,
1029 },
1030 )?;
1031 }
1032 }
1033 Ok(())
1034 }
1035
1036 pub async fn invite_count_updated(self: &Arc<Self>, user_id: UserId) -> Result<()> {
1037 if let Some(user) = self.app_state.db.get_user_by_id(user_id).await?
1038 && let Some(invite_code) = &user.invite_code
1039 {
1040 let pool = self.connection_pool.lock();
1041 for connection_id in pool.user_connection_ids(user_id) {
1042 self.peer.send(
1043 connection_id,
1044 proto::UpdateInviteInfo {
1045 url: format!(
1046 "{}{}",
1047 self.app_state.config.invite_link_prefix, invite_code
1048 ),
1049 count: user.invite_count as u32,
1050 },
1051 )?;
1052 }
1053 }
1054 Ok(())
1055 }
1056
1057 pub async fn snapshot(self: &Arc<Self>) -> ServerSnapshot<'_> {
1058 ServerSnapshot {
1059 connection_pool: ConnectionPoolGuard {
1060 guard: self.connection_pool.lock(),
1061 _not_send: PhantomData,
1062 },
1063 peer: &self.peer,
1064 }
1065 }
1066}
1067
1068impl Deref for ConnectionPoolGuard<'_> {
1069 type Target = ConnectionPool;
1070
1071 fn deref(&self) -> &Self::Target {
1072 &self.guard
1073 }
1074}
1075
1076impl DerefMut for ConnectionPoolGuard<'_> {
1077 fn deref_mut(&mut self) -> &mut Self::Target {
1078 &mut self.guard
1079 }
1080}
1081
1082impl Drop for ConnectionPoolGuard<'_> {
1083 fn drop(&mut self) {
1084 #[cfg(test)]
1085 self.check_invariants();
1086 }
1087}
1088
1089fn broadcast<F>(
1090 sender_id: Option<ConnectionId>,
1091 receiver_ids: impl IntoIterator<Item = ConnectionId>,
1092 mut f: F,
1093) where
1094 F: FnMut(ConnectionId) -> anyhow::Result<()>,
1095{
1096 for receiver_id in receiver_ids {
1097 if Some(receiver_id) != sender_id
1098 && let Err(error) = f(receiver_id)
1099 {
1100 tracing::error!("failed to send to {:?} {}", receiver_id, error);
1101 }
1102 }
1103}
1104
1105pub struct ProtocolVersion(u32);
1106
1107impl Header for ProtocolVersion {
1108 fn name() -> &'static HeaderName {
1109 static ZED_PROTOCOL_VERSION: OnceLock<HeaderName> = OnceLock::new();
1110 ZED_PROTOCOL_VERSION.get_or_init(|| HeaderName::from_static("x-zed-protocol-version"))
1111 }
1112
1113 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1114 where
1115 Self: Sized,
1116 I: Iterator<Item = &'i axum::http::HeaderValue>,
1117 {
1118 let version = values
1119 .next()
1120 .ok_or_else(axum::headers::Error::invalid)?
1121 .to_str()
1122 .map_err(|_| axum::headers::Error::invalid())?
1123 .parse()
1124 .map_err(|_| axum::headers::Error::invalid())?;
1125 Ok(Self(version))
1126 }
1127
1128 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1129 values.extend([self.0.to_string().parse().unwrap()]);
1130 }
1131}
1132
1133pub struct AppVersionHeader(SemanticVersion);
1134impl Header for AppVersionHeader {
1135 fn name() -> &'static HeaderName {
1136 static ZED_APP_VERSION: OnceLock<HeaderName> = OnceLock::new();
1137 ZED_APP_VERSION.get_or_init(|| HeaderName::from_static("x-zed-app-version"))
1138 }
1139
1140 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1141 where
1142 Self: Sized,
1143 I: Iterator<Item = &'i axum::http::HeaderValue>,
1144 {
1145 let version = values
1146 .next()
1147 .ok_or_else(axum::headers::Error::invalid)?
1148 .to_str()
1149 .map_err(|_| axum::headers::Error::invalid())?
1150 .parse()
1151 .map_err(|_| axum::headers::Error::invalid())?;
1152 Ok(Self(version))
1153 }
1154
1155 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1156 values.extend([self.0.to_string().parse().unwrap()]);
1157 }
1158}
1159
1160#[derive(Debug)]
1161pub struct ReleaseChannelHeader(String);
1162
1163impl Header for ReleaseChannelHeader {
1164 fn name() -> &'static HeaderName {
1165 static ZED_RELEASE_CHANNEL: OnceLock<HeaderName> = OnceLock::new();
1166 ZED_RELEASE_CHANNEL.get_or_init(|| HeaderName::from_static("x-zed-release-channel"))
1167 }
1168
1169 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1170 where
1171 Self: Sized,
1172 I: Iterator<Item = &'i axum::http::HeaderValue>,
1173 {
1174 Ok(Self(
1175 values
1176 .next()
1177 .ok_or_else(axum::headers::Error::invalid)?
1178 .to_str()
1179 .map_err(|_| axum::headers::Error::invalid())?
1180 .to_owned(),
1181 ))
1182 }
1183
1184 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1185 values.extend([self.0.parse().unwrap()]);
1186 }
1187}
1188
1189pub fn routes(server: Arc<Server>) -> Router<(), Body> {
1190 Router::new()
1191 .route("/rpc", get(handle_websocket_request))
1192 .layer(
1193 ServiceBuilder::new()
1194 .layer(Extension(server.app_state.clone()))
1195 .layer(middleware::from_fn(auth::validate_header)),
1196 )
1197 .route("/metrics", get(handle_metrics))
1198 .layer(Extension(server))
1199}
1200
1201pub async fn handle_websocket_request(
1202 TypedHeader(ProtocolVersion(protocol_version)): TypedHeader<ProtocolVersion>,
1203 app_version_header: Option<TypedHeader<AppVersionHeader>>,
1204 release_channel_header: Option<TypedHeader<ReleaseChannelHeader>>,
1205 ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
1206 Extension(server): Extension<Arc<Server>>,
1207 Extension(principal): Extension<Principal>,
1208 user_agent: Option<TypedHeader<UserAgent>>,
1209 country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
1210 system_id_header: Option<TypedHeader<SystemIdHeader>>,
1211 ws: WebSocketUpgrade,
1212) -> axum::response::Response {
1213 if protocol_version != rpc::PROTOCOL_VERSION {
1214 return (
1215 StatusCode::UPGRADE_REQUIRED,
1216 "client must be upgraded".to_string(),
1217 )
1218 .into_response();
1219 }
1220
1221 let Some(version) = app_version_header.map(|header| ZedVersion(header.0.0)) else {
1222 return (
1223 StatusCode::UPGRADE_REQUIRED,
1224 "no version header found".to_string(),
1225 )
1226 .into_response();
1227 };
1228
1229 let release_channel = release_channel_header.map(|header| header.0.0);
1230
1231 if !version.can_collaborate() {
1232 return (
1233 StatusCode::UPGRADE_REQUIRED,
1234 "client must be upgraded".to_string(),
1235 )
1236 .into_response();
1237 }
1238
1239 let socket_address = socket_address.to_string();
1240
1241 // Acquire connection guard before WebSocket upgrade
1242 let connection_guard = match ConnectionGuard::try_acquire() {
1243 Ok(guard) => guard,
1244 Err(()) => {
1245 return (
1246 StatusCode::SERVICE_UNAVAILABLE,
1247 "Too many concurrent connections",
1248 )
1249 .into_response();
1250 }
1251 };
1252
1253 ws.on_upgrade(move |socket| {
1254 let socket = socket
1255 .map_ok(to_tungstenite_message)
1256 .err_into()
1257 .with(|message| async move { to_axum_message(message) });
1258 let connection = Connection::new(Box::pin(socket));
1259 async move {
1260 server
1261 .handle_connection(
1262 connection,
1263 socket_address,
1264 principal,
1265 version,
1266 release_channel,
1267 user_agent.map(|header| header.to_string()),
1268 country_code_header.map(|header| header.to_string()),
1269 system_id_header.map(|header| header.to_string()),
1270 None,
1271 Executor::Production,
1272 Some(connection_guard),
1273 )
1274 .await;
1275 }
1276 })
1277}
1278
1279pub async fn handle_metrics(Extension(server): Extension<Arc<Server>>) -> Result<String> {
1280 static CONNECTIONS_METRIC: OnceLock<IntGauge> = OnceLock::new();
1281 let connections_metric = CONNECTIONS_METRIC
1282 .get_or_init(|| register_int_gauge!("connections", "number of connections").unwrap());
1283
1284 let connections = server
1285 .connection_pool
1286 .lock()
1287 .connections()
1288 .filter(|connection| !connection.admin)
1289 .count();
1290 connections_metric.set(connections as _);
1291
1292 static SHARED_PROJECTS_METRIC: OnceLock<IntGauge> = OnceLock::new();
1293 let shared_projects_metric = SHARED_PROJECTS_METRIC.get_or_init(|| {
1294 register_int_gauge!(
1295 "shared_projects",
1296 "number of open projects with one or more guests"
1297 )
1298 .unwrap()
1299 });
1300
1301 let shared_projects = server.app_state.db.project_count_excluding_admins().await?;
1302 shared_projects_metric.set(shared_projects as _);
1303
1304 let encoder = prometheus::TextEncoder::new();
1305 let metric_families = prometheus::gather();
1306 let encoded_metrics = encoder
1307 .encode_to_string(&metric_families)
1308 .map_err(|err| anyhow!("{err}"))?;
1309 Ok(encoded_metrics)
1310}
1311
1312#[instrument(err, skip(executor))]
1313async fn connection_lost(
1314 session: Session,
1315 mut teardown: watch::Receiver<bool>,
1316 executor: Executor,
1317) -> Result<()> {
1318 session.peer.disconnect(session.connection_id);
1319 session
1320 .connection_pool()
1321 .await
1322 .remove_connection(session.connection_id)?;
1323
1324 session
1325 .db()
1326 .await
1327 .connection_lost(session.connection_id)
1328 .await
1329 .trace_err();
1330
1331 futures::select_biased! {
1332 _ = executor.sleep(RECONNECT_TIMEOUT).fuse() => {
1333
1334 log::info!("connection lost, removing all resources for user:{}, connection:{:?}", session.user_id(), session.connection_id);
1335 leave_room_for_session(&session, session.connection_id).await.trace_err();
1336 leave_channel_buffers_for_session(&session)
1337 .await
1338 .trace_err();
1339
1340 if !session
1341 .connection_pool()
1342 .await
1343 .is_user_online(session.user_id())
1344 {
1345 let db = session.db().await;
1346 if let Some(room) = db.decline_call(None, session.user_id()).await.trace_err().flatten() {
1347 room_updated(&room, &session.peer);
1348 }
1349 }
1350
1351 update_user_contacts(session.user_id(), &session).await?;
1352 },
1353 _ = teardown.changed().fuse() => {}
1354 }
1355
1356 Ok(())
1357}
1358
1359/// Acknowledges a ping from a client, used to keep the connection alive.
1360async fn ping(
1361 _: proto::Ping,
1362 response: Response<proto::Ping>,
1363 _session: MessageContext,
1364) -> Result<()> {
1365 response.send(proto::Ack {})?;
1366 Ok(())
1367}
1368
1369/// Creates a new room for calling (outside of channels)
1370async fn create_room(
1371 _request: proto::CreateRoom,
1372 response: Response<proto::CreateRoom>,
1373 session: MessageContext,
1374) -> Result<()> {
1375 let livekit_room = nanoid::nanoid!(30);
1376
1377 let live_kit_connection_info = util::maybe!(async {
1378 let live_kit = session.app_state.livekit_client.as_ref();
1379 let live_kit = live_kit?;
1380 let user_id = session.user_id().to_string();
1381
1382 let token = live_kit.room_token(&livekit_room, &user_id).trace_err()?;
1383
1384 Some(proto::LiveKitConnectionInfo {
1385 server_url: live_kit.url().into(),
1386 token,
1387 can_publish: true,
1388 })
1389 })
1390 .await;
1391
1392 let room = session
1393 .db()
1394 .await
1395 .create_room(session.user_id(), session.connection_id, &livekit_room)
1396 .await?;
1397
1398 response.send(proto::CreateRoomResponse {
1399 room: Some(room.clone()),
1400 live_kit_connection_info,
1401 })?;
1402
1403 update_user_contacts(session.user_id(), &session).await?;
1404 Ok(())
1405}
1406
1407/// Join a room from an invitation. Equivalent to joining a channel if there is one.
1408async fn join_room(
1409 request: proto::JoinRoom,
1410 response: Response<proto::JoinRoom>,
1411 session: MessageContext,
1412) -> Result<()> {
1413 let room_id = RoomId::from_proto(request.id);
1414
1415 let channel_id = session.db().await.channel_id_for_room(room_id).await?;
1416
1417 if let Some(channel_id) = channel_id {
1418 return join_channel_internal(channel_id, Box::new(response), session).await;
1419 }
1420
1421 let joined_room = {
1422 let room = session
1423 .db()
1424 .await
1425 .join_room(room_id, session.user_id(), session.connection_id)
1426 .await?;
1427 room_updated(&room.room, &session.peer);
1428 room.into_inner()
1429 };
1430
1431 for connection_id in session
1432 .connection_pool()
1433 .await
1434 .user_connection_ids(session.user_id())
1435 {
1436 session
1437 .peer
1438 .send(
1439 connection_id,
1440 proto::CallCanceled {
1441 room_id: room_id.to_proto(),
1442 },
1443 )
1444 .trace_err();
1445 }
1446
1447 let live_kit_connection_info = if let Some(live_kit) = session.app_state.livekit_client.as_ref()
1448 {
1449 live_kit
1450 .room_token(
1451 &joined_room.room.livekit_room,
1452 &session.user_id().to_string(),
1453 )
1454 .trace_err()
1455 .map(|token| proto::LiveKitConnectionInfo {
1456 server_url: live_kit.url().into(),
1457 token,
1458 can_publish: true,
1459 })
1460 } else {
1461 None
1462 };
1463
1464 response.send(proto::JoinRoomResponse {
1465 room: Some(joined_room.room),
1466 channel_id: None,
1467 live_kit_connection_info,
1468 })?;
1469
1470 update_user_contacts(session.user_id(), &session).await?;
1471 Ok(())
1472}
1473
1474/// Rejoin room is used to reconnect to a room after connection errors.
1475async fn rejoin_room(
1476 request: proto::RejoinRoom,
1477 response: Response<proto::RejoinRoom>,
1478 session: MessageContext,
1479) -> Result<()> {
1480 let room;
1481 let channel;
1482 {
1483 let mut rejoined_room = session
1484 .db()
1485 .await
1486 .rejoin_room(request, session.user_id(), session.connection_id)
1487 .await?;
1488
1489 response.send(proto::RejoinRoomResponse {
1490 room: Some(rejoined_room.room.clone()),
1491 reshared_projects: rejoined_room
1492 .reshared_projects
1493 .iter()
1494 .map(|project| proto::ResharedProject {
1495 id: project.id.to_proto(),
1496 collaborators: project
1497 .collaborators
1498 .iter()
1499 .map(|collaborator| collaborator.to_proto())
1500 .collect(),
1501 })
1502 .collect(),
1503 rejoined_projects: rejoined_room
1504 .rejoined_projects
1505 .iter()
1506 .map(|rejoined_project| rejoined_project.to_proto())
1507 .collect(),
1508 })?;
1509 room_updated(&rejoined_room.room, &session.peer);
1510
1511 for project in &rejoined_room.reshared_projects {
1512 for collaborator in &project.collaborators {
1513 session
1514 .peer
1515 .send(
1516 collaborator.connection_id,
1517 proto::UpdateProjectCollaborator {
1518 project_id: project.id.to_proto(),
1519 old_peer_id: Some(project.old_connection_id.into()),
1520 new_peer_id: Some(session.connection_id.into()),
1521 },
1522 )
1523 .trace_err();
1524 }
1525
1526 broadcast(
1527 Some(session.connection_id),
1528 project
1529 .collaborators
1530 .iter()
1531 .map(|collaborator| collaborator.connection_id),
1532 |connection_id| {
1533 session.peer.forward_send(
1534 session.connection_id,
1535 connection_id,
1536 proto::UpdateProject {
1537 project_id: project.id.to_proto(),
1538 worktrees: project.worktrees.clone(),
1539 },
1540 )
1541 },
1542 );
1543 }
1544
1545 notify_rejoined_projects(&mut rejoined_room.rejoined_projects, &session)?;
1546
1547 let rejoined_room = rejoined_room.into_inner();
1548
1549 room = rejoined_room.room;
1550 channel = rejoined_room.channel;
1551 }
1552
1553 if let Some(channel) = channel {
1554 channel_updated(
1555 &channel,
1556 &room,
1557 &session.peer,
1558 &*session.connection_pool().await,
1559 );
1560 }
1561
1562 update_user_contacts(session.user_id(), &session).await?;
1563 Ok(())
1564}
1565
1566fn notify_rejoined_projects(
1567 rejoined_projects: &mut Vec<RejoinedProject>,
1568 session: &Session,
1569) -> Result<()> {
1570 for project in rejoined_projects.iter() {
1571 for collaborator in &project.collaborators {
1572 session
1573 .peer
1574 .send(
1575 collaborator.connection_id,
1576 proto::UpdateProjectCollaborator {
1577 project_id: project.id.to_proto(),
1578 old_peer_id: Some(project.old_connection_id.into()),
1579 new_peer_id: Some(session.connection_id.into()),
1580 },
1581 )
1582 .trace_err();
1583 }
1584 }
1585
1586 for project in rejoined_projects {
1587 for worktree in mem::take(&mut project.worktrees) {
1588 // Stream this worktree's entries.
1589 let message = proto::UpdateWorktree {
1590 project_id: project.id.to_proto(),
1591 worktree_id: worktree.id,
1592 abs_path: worktree.abs_path.clone(),
1593 root_name: worktree.root_name,
1594 updated_entries: worktree.updated_entries,
1595 removed_entries: worktree.removed_entries,
1596 scan_id: worktree.scan_id,
1597 is_last_update: worktree.completed_scan_id == worktree.scan_id,
1598 updated_repositories: worktree.updated_repositories,
1599 removed_repositories: worktree.removed_repositories,
1600 };
1601 for update in proto::split_worktree_update(message) {
1602 session.peer.send(session.connection_id, update)?;
1603 }
1604
1605 // Stream this worktree's diagnostics.
1606 let mut worktree_diagnostics = worktree.diagnostic_summaries.into_iter();
1607 if let Some(summary) = worktree_diagnostics.next() {
1608 let message = proto::UpdateDiagnosticSummary {
1609 project_id: project.id.to_proto(),
1610 worktree_id: worktree.id,
1611 summary: Some(summary),
1612 more_summaries: worktree_diagnostics.collect(),
1613 };
1614 session.peer.send(session.connection_id, message)?;
1615 }
1616
1617 for settings_file in worktree.settings_files {
1618 session.peer.send(
1619 session.connection_id,
1620 proto::UpdateWorktreeSettings {
1621 project_id: project.id.to_proto(),
1622 worktree_id: worktree.id,
1623 path: settings_file.path,
1624 content: Some(settings_file.content),
1625 kind: Some(settings_file.kind.to_proto().into()),
1626 },
1627 )?;
1628 }
1629 }
1630
1631 for repository in mem::take(&mut project.updated_repositories) {
1632 for update in split_repository_update(repository) {
1633 session.peer.send(session.connection_id, update)?;
1634 }
1635 }
1636
1637 for id in mem::take(&mut project.removed_repositories) {
1638 session.peer.send(
1639 session.connection_id,
1640 proto::RemoveRepository {
1641 project_id: project.id.to_proto(),
1642 id,
1643 },
1644 )?;
1645 }
1646 }
1647
1648 Ok(())
1649}
1650
1651/// leave room disconnects from the room.
1652async fn leave_room(
1653 _: proto::LeaveRoom,
1654 response: Response<proto::LeaveRoom>,
1655 session: MessageContext,
1656) -> Result<()> {
1657 leave_room_for_session(&session, session.connection_id).await?;
1658 response.send(proto::Ack {})?;
1659 Ok(())
1660}
1661
1662/// Updates the permissions of someone else in the room.
1663async fn set_room_participant_role(
1664 request: proto::SetRoomParticipantRole,
1665 response: Response<proto::SetRoomParticipantRole>,
1666 session: MessageContext,
1667) -> Result<()> {
1668 let user_id = UserId::from_proto(request.user_id);
1669 let role = ChannelRole::from(request.role());
1670
1671 let (livekit_room, can_publish) = {
1672 let room = session
1673 .db()
1674 .await
1675 .set_room_participant_role(
1676 session.user_id(),
1677 RoomId::from_proto(request.room_id),
1678 user_id,
1679 role,
1680 )
1681 .await?;
1682
1683 let livekit_room = room.livekit_room.clone();
1684 let can_publish = ChannelRole::from(request.role()).can_use_microphone();
1685 room_updated(&room, &session.peer);
1686 (livekit_room, can_publish)
1687 };
1688
1689 if let Some(live_kit) = session.app_state.livekit_client.as_ref() {
1690 live_kit
1691 .update_participant(
1692 livekit_room.clone(),
1693 request.user_id.to_string(),
1694 livekit_api::proto::ParticipantPermission {
1695 can_subscribe: true,
1696 can_publish,
1697 can_publish_data: can_publish,
1698 hidden: false,
1699 recorder: false,
1700 },
1701 )
1702 .await
1703 .trace_err();
1704 }
1705
1706 response.send(proto::Ack {})?;
1707 Ok(())
1708}
1709
1710/// Call someone else into the current room
1711async fn call(
1712 request: proto::Call,
1713 response: Response<proto::Call>,
1714 session: MessageContext,
1715) -> Result<()> {
1716 let room_id = RoomId::from_proto(request.room_id);
1717 let calling_user_id = session.user_id();
1718 let calling_connection_id = session.connection_id;
1719 let called_user_id = UserId::from_proto(request.called_user_id);
1720 let initial_project_id = request.initial_project_id.map(ProjectId::from_proto);
1721 if !session
1722 .db()
1723 .await
1724 .has_contact(calling_user_id, called_user_id)
1725 .await?
1726 {
1727 return Err(anyhow!("cannot call a user who isn't a contact"))?;
1728 }
1729
1730 let incoming_call = {
1731 let (room, incoming_call) = &mut *session
1732 .db()
1733 .await
1734 .call(
1735 room_id,
1736 calling_user_id,
1737 calling_connection_id,
1738 called_user_id,
1739 initial_project_id,
1740 )
1741 .await?;
1742 room_updated(room, &session.peer);
1743 mem::take(incoming_call)
1744 };
1745 update_user_contacts(called_user_id, &session).await?;
1746
1747 let mut calls = session
1748 .connection_pool()
1749 .await
1750 .user_connection_ids(called_user_id)
1751 .map(|connection_id| session.peer.request(connection_id, incoming_call.clone()))
1752 .collect::<FuturesUnordered<_>>();
1753
1754 while let Some(call_response) = calls.next().await {
1755 match call_response.as_ref() {
1756 Ok(_) => {
1757 response.send(proto::Ack {})?;
1758 return Ok(());
1759 }
1760 Err(_) => {
1761 call_response.trace_err();
1762 }
1763 }
1764 }
1765
1766 {
1767 let room = session
1768 .db()
1769 .await
1770 .call_failed(room_id, called_user_id)
1771 .await?;
1772 room_updated(&room, &session.peer);
1773 }
1774 update_user_contacts(called_user_id, &session).await?;
1775
1776 Err(anyhow!("failed to ring user"))?
1777}
1778
1779/// Cancel an outgoing call.
1780async fn cancel_call(
1781 request: proto::CancelCall,
1782 response: Response<proto::CancelCall>,
1783 session: MessageContext,
1784) -> Result<()> {
1785 let called_user_id = UserId::from_proto(request.called_user_id);
1786 let room_id = RoomId::from_proto(request.room_id);
1787 {
1788 let room = session
1789 .db()
1790 .await
1791 .cancel_call(room_id, session.connection_id, called_user_id)
1792 .await?;
1793 room_updated(&room, &session.peer);
1794 }
1795
1796 for connection_id in session
1797 .connection_pool()
1798 .await
1799 .user_connection_ids(called_user_id)
1800 {
1801 session
1802 .peer
1803 .send(
1804 connection_id,
1805 proto::CallCanceled {
1806 room_id: room_id.to_proto(),
1807 },
1808 )
1809 .trace_err();
1810 }
1811 response.send(proto::Ack {})?;
1812
1813 update_user_contacts(called_user_id, &session).await?;
1814 Ok(())
1815}
1816
1817/// Decline an incoming call.
1818async fn decline_call(message: proto::DeclineCall, session: MessageContext) -> Result<()> {
1819 let room_id = RoomId::from_proto(message.room_id);
1820 {
1821 let room = session
1822 .db()
1823 .await
1824 .decline_call(Some(room_id), session.user_id())
1825 .await?
1826 .context("declining call")?;
1827 room_updated(&room, &session.peer);
1828 }
1829
1830 for connection_id in session
1831 .connection_pool()
1832 .await
1833 .user_connection_ids(session.user_id())
1834 {
1835 session
1836 .peer
1837 .send(
1838 connection_id,
1839 proto::CallCanceled {
1840 room_id: room_id.to_proto(),
1841 },
1842 )
1843 .trace_err();
1844 }
1845 update_user_contacts(session.user_id(), &session).await?;
1846 Ok(())
1847}
1848
1849/// Updates other participants in the room with your current location.
1850async fn update_participant_location(
1851 request: proto::UpdateParticipantLocation,
1852 response: Response<proto::UpdateParticipantLocation>,
1853 session: MessageContext,
1854) -> Result<()> {
1855 let room_id = RoomId::from_proto(request.room_id);
1856 let location = request.location.context("invalid location")?;
1857
1858 let db = session.db().await;
1859 let room = db
1860 .update_room_participant_location(room_id, session.connection_id, location)
1861 .await?;
1862
1863 room_updated(&room, &session.peer);
1864 response.send(proto::Ack {})?;
1865 Ok(())
1866}
1867
1868/// Share a project into the room.
1869async fn share_project(
1870 request: proto::ShareProject,
1871 response: Response<proto::ShareProject>,
1872 session: MessageContext,
1873) -> Result<()> {
1874 let (project_id, room) = &*session
1875 .db()
1876 .await
1877 .share_project(
1878 RoomId::from_proto(request.room_id),
1879 session.connection_id,
1880 &request.worktrees,
1881 request.is_ssh_project,
1882 )
1883 .await?;
1884 response.send(proto::ShareProjectResponse {
1885 project_id: project_id.to_proto(),
1886 })?;
1887 room_updated(room, &session.peer);
1888
1889 Ok(())
1890}
1891
1892/// Unshare a project from the room.
1893async fn unshare_project(message: proto::UnshareProject, session: MessageContext) -> Result<()> {
1894 let project_id = ProjectId::from_proto(message.project_id);
1895 unshare_project_internal(project_id, session.connection_id, &session).await
1896}
1897
1898async fn unshare_project_internal(
1899 project_id: ProjectId,
1900 connection_id: ConnectionId,
1901 session: &Session,
1902) -> Result<()> {
1903 let delete = {
1904 let room_guard = session
1905 .db()
1906 .await
1907 .unshare_project(project_id, connection_id)
1908 .await?;
1909
1910 let (delete, room, guest_connection_ids) = &*room_guard;
1911
1912 let message = proto::UnshareProject {
1913 project_id: project_id.to_proto(),
1914 };
1915
1916 broadcast(
1917 Some(connection_id),
1918 guest_connection_ids.iter().copied(),
1919 |conn_id| session.peer.send(conn_id, message.clone()),
1920 );
1921 if let Some(room) = room {
1922 room_updated(room, &session.peer);
1923 }
1924
1925 *delete
1926 };
1927
1928 if delete {
1929 let db = session.db().await;
1930 db.delete_project(project_id).await?;
1931 }
1932
1933 Ok(())
1934}
1935
1936/// Join someone elses shared project.
1937async fn join_project(
1938 request: proto::JoinProject,
1939 response: Response<proto::JoinProject>,
1940 session: MessageContext,
1941) -> Result<()> {
1942 let project_id = ProjectId::from_proto(request.project_id);
1943
1944 tracing::info!(%project_id, "join project");
1945
1946 let db = session.db().await;
1947 let (project, replica_id) = &mut *db
1948 .join_project(
1949 project_id,
1950 session.connection_id,
1951 session.user_id(),
1952 request.committer_name.clone(),
1953 request.committer_email.clone(),
1954 )
1955 .await?;
1956 drop(db);
1957 tracing::info!(%project_id, "join remote project");
1958 let collaborators = project
1959 .collaborators
1960 .iter()
1961 .filter(|collaborator| collaborator.connection_id != session.connection_id)
1962 .map(|collaborator| collaborator.to_proto())
1963 .collect::<Vec<_>>();
1964 let project_id = project.id;
1965 let guest_user_id = session.user_id();
1966
1967 let worktrees = project
1968 .worktrees
1969 .iter()
1970 .map(|(id, worktree)| proto::WorktreeMetadata {
1971 id: *id,
1972 root_name: worktree.root_name.clone(),
1973 visible: worktree.visible,
1974 abs_path: worktree.abs_path.clone(),
1975 })
1976 .collect::<Vec<_>>();
1977
1978 let add_project_collaborator = proto::AddProjectCollaborator {
1979 project_id: project_id.to_proto(),
1980 collaborator: Some(proto::Collaborator {
1981 peer_id: Some(session.connection_id.into()),
1982 replica_id: replica_id.0 as u32,
1983 user_id: guest_user_id.to_proto(),
1984 is_host: false,
1985 committer_name: request.committer_name.clone(),
1986 committer_email: request.committer_email.clone(),
1987 }),
1988 };
1989
1990 for collaborator in &collaborators {
1991 session
1992 .peer
1993 .send(
1994 collaborator.peer_id.unwrap().into(),
1995 add_project_collaborator.clone(),
1996 )
1997 .trace_err();
1998 }
1999
2000 // First, we send the metadata associated with each worktree.
2001 let (language_servers, language_server_capabilities) = project
2002 .language_servers
2003 .clone()
2004 .into_iter()
2005 .map(|server| (server.server, server.capabilities))
2006 .unzip();
2007 response.send(proto::JoinProjectResponse {
2008 project_id: project.id.0 as u64,
2009 worktrees,
2010 replica_id: replica_id.0 as u32,
2011 collaborators,
2012 language_servers,
2013 language_server_capabilities,
2014 role: project.role.into(),
2015 })?;
2016
2017 for (worktree_id, worktree) in mem::take(&mut project.worktrees) {
2018 // Stream this worktree's entries.
2019 let message = proto::UpdateWorktree {
2020 project_id: project_id.to_proto(),
2021 worktree_id,
2022 abs_path: worktree.abs_path.clone(),
2023 root_name: worktree.root_name,
2024 updated_entries: worktree.entries,
2025 removed_entries: Default::default(),
2026 scan_id: worktree.scan_id,
2027 is_last_update: worktree.scan_id == worktree.completed_scan_id,
2028 updated_repositories: worktree.legacy_repository_entries.into_values().collect(),
2029 removed_repositories: Default::default(),
2030 };
2031 for update in proto::split_worktree_update(message) {
2032 session.peer.send(session.connection_id, update.clone())?;
2033 }
2034
2035 // Stream this worktree's diagnostics.
2036 let mut worktree_diagnostics = worktree.diagnostic_summaries.into_iter();
2037 if let Some(summary) = worktree_diagnostics.next() {
2038 let message = proto::UpdateDiagnosticSummary {
2039 project_id: project.id.to_proto(),
2040 worktree_id: worktree.id,
2041 summary: Some(summary),
2042 more_summaries: worktree_diagnostics.collect(),
2043 };
2044 session.peer.send(session.connection_id, message)?;
2045 }
2046
2047 for settings_file in worktree.settings_files {
2048 session.peer.send(
2049 session.connection_id,
2050 proto::UpdateWorktreeSettings {
2051 project_id: project_id.to_proto(),
2052 worktree_id: worktree.id,
2053 path: settings_file.path,
2054 content: Some(settings_file.content),
2055 kind: Some(settings_file.kind.to_proto() as i32),
2056 },
2057 )?;
2058 }
2059 }
2060
2061 for repository in mem::take(&mut project.repositories) {
2062 for update in split_repository_update(repository) {
2063 session.peer.send(session.connection_id, update)?;
2064 }
2065 }
2066
2067 for language_server in &project.language_servers {
2068 session.peer.send(
2069 session.connection_id,
2070 proto::UpdateLanguageServer {
2071 project_id: project_id.to_proto(),
2072 server_name: Some(language_server.server.name.clone()),
2073 language_server_id: language_server.server.id,
2074 variant: Some(
2075 proto::update_language_server::Variant::DiskBasedDiagnosticsUpdated(
2076 proto::LspDiskBasedDiagnosticsUpdated {},
2077 ),
2078 ),
2079 },
2080 )?;
2081 }
2082
2083 Ok(())
2084}
2085
2086/// Leave someone elses shared project.
2087async fn leave_project(request: proto::LeaveProject, session: MessageContext) -> Result<()> {
2088 let sender_id = session.connection_id;
2089 let project_id = ProjectId::from_proto(request.project_id);
2090 let db = session.db().await;
2091
2092 let (room, project) = &*db.leave_project(project_id, sender_id).await?;
2093 tracing::info!(
2094 %project_id,
2095 "leave project"
2096 );
2097
2098 project_left(project, &session);
2099 if let Some(room) = room {
2100 room_updated(room, &session.peer);
2101 }
2102
2103 Ok(())
2104}
2105
2106/// Updates other participants with changes to the project
2107async fn update_project(
2108 request: proto::UpdateProject,
2109 response: Response<proto::UpdateProject>,
2110 session: MessageContext,
2111) -> Result<()> {
2112 let project_id = ProjectId::from_proto(request.project_id);
2113 let (room, guest_connection_ids) = &*session
2114 .db()
2115 .await
2116 .update_project(project_id, session.connection_id, &request.worktrees)
2117 .await?;
2118 broadcast(
2119 Some(session.connection_id),
2120 guest_connection_ids.iter().copied(),
2121 |connection_id| {
2122 session
2123 .peer
2124 .forward_send(session.connection_id, connection_id, request.clone())
2125 },
2126 );
2127 if let Some(room) = room {
2128 room_updated(room, &session.peer);
2129 }
2130 response.send(proto::Ack {})?;
2131
2132 Ok(())
2133}
2134
2135/// Updates other participants with changes to the worktree
2136async fn update_worktree(
2137 request: proto::UpdateWorktree,
2138 response: Response<proto::UpdateWorktree>,
2139 session: MessageContext,
2140) -> Result<()> {
2141 let guest_connection_ids = session
2142 .db()
2143 .await
2144 .update_worktree(&request, session.connection_id)
2145 .await?;
2146
2147 broadcast(
2148 Some(session.connection_id),
2149 guest_connection_ids.iter().copied(),
2150 |connection_id| {
2151 session
2152 .peer
2153 .forward_send(session.connection_id, connection_id, request.clone())
2154 },
2155 );
2156 response.send(proto::Ack {})?;
2157 Ok(())
2158}
2159
2160async fn update_repository(
2161 request: proto::UpdateRepository,
2162 response: Response<proto::UpdateRepository>,
2163 session: MessageContext,
2164) -> Result<()> {
2165 let guest_connection_ids = session
2166 .db()
2167 .await
2168 .update_repository(&request, session.connection_id)
2169 .await?;
2170
2171 broadcast(
2172 Some(session.connection_id),
2173 guest_connection_ids.iter().copied(),
2174 |connection_id| {
2175 session
2176 .peer
2177 .forward_send(session.connection_id, connection_id, request.clone())
2178 },
2179 );
2180 response.send(proto::Ack {})?;
2181 Ok(())
2182}
2183
2184async fn remove_repository(
2185 request: proto::RemoveRepository,
2186 response: Response<proto::RemoveRepository>,
2187 session: MessageContext,
2188) -> Result<()> {
2189 let guest_connection_ids = session
2190 .db()
2191 .await
2192 .remove_repository(&request, session.connection_id)
2193 .await?;
2194
2195 broadcast(
2196 Some(session.connection_id),
2197 guest_connection_ids.iter().copied(),
2198 |connection_id| {
2199 session
2200 .peer
2201 .forward_send(session.connection_id, connection_id, request.clone())
2202 },
2203 );
2204 response.send(proto::Ack {})?;
2205 Ok(())
2206}
2207
2208/// Updates other participants with changes to the diagnostics
2209async fn update_diagnostic_summary(
2210 message: proto::UpdateDiagnosticSummary,
2211 session: MessageContext,
2212) -> Result<()> {
2213 let guest_connection_ids = session
2214 .db()
2215 .await
2216 .update_diagnostic_summary(&message, session.connection_id)
2217 .await?;
2218
2219 broadcast(
2220 Some(session.connection_id),
2221 guest_connection_ids.iter().copied(),
2222 |connection_id| {
2223 session
2224 .peer
2225 .forward_send(session.connection_id, connection_id, message.clone())
2226 },
2227 );
2228
2229 Ok(())
2230}
2231
2232/// Updates other participants with changes to the worktree settings
2233async fn update_worktree_settings(
2234 message: proto::UpdateWorktreeSettings,
2235 session: MessageContext,
2236) -> Result<()> {
2237 let guest_connection_ids = session
2238 .db()
2239 .await
2240 .update_worktree_settings(&message, session.connection_id)
2241 .await?;
2242
2243 broadcast(
2244 Some(session.connection_id),
2245 guest_connection_ids.iter().copied(),
2246 |connection_id| {
2247 session
2248 .peer
2249 .forward_send(session.connection_id, connection_id, message.clone())
2250 },
2251 );
2252
2253 Ok(())
2254}
2255
2256/// Notify other participants that a language server has started.
2257async fn start_language_server(
2258 request: proto::StartLanguageServer,
2259 session: MessageContext,
2260) -> Result<()> {
2261 let guest_connection_ids = session
2262 .db()
2263 .await
2264 .start_language_server(&request, session.connection_id)
2265 .await?;
2266
2267 broadcast(
2268 Some(session.connection_id),
2269 guest_connection_ids.iter().copied(),
2270 |connection_id| {
2271 session
2272 .peer
2273 .forward_send(session.connection_id, connection_id, request.clone())
2274 },
2275 );
2276 Ok(())
2277}
2278
2279/// Notify other participants that a language server has changed.
2280async fn update_language_server(
2281 request: proto::UpdateLanguageServer,
2282 session: MessageContext,
2283) -> Result<()> {
2284 let project_id = ProjectId::from_proto(request.project_id);
2285 let db = session.db().await;
2286
2287 if let Some(proto::update_language_server::Variant::MetadataUpdated(update)) = &request.variant
2288 && let Some(capabilities) = update.capabilities.clone()
2289 {
2290 db.update_server_capabilities(project_id, request.language_server_id, capabilities)
2291 .await?;
2292 }
2293
2294 let project_connection_ids = db
2295 .project_connection_ids(project_id, session.connection_id, true)
2296 .await?;
2297 broadcast(
2298 Some(session.connection_id),
2299 project_connection_ids.iter().copied(),
2300 |connection_id| {
2301 session
2302 .peer
2303 .forward_send(session.connection_id, connection_id, request.clone())
2304 },
2305 );
2306 Ok(())
2307}
2308
2309/// forward a project request to the host. These requests should be read only
2310/// as guests are allowed to send them.
2311async fn forward_read_only_project_request<T>(
2312 request: T,
2313 response: Response<T>,
2314 session: MessageContext,
2315) -> Result<()>
2316where
2317 T: EntityMessage + RequestMessage,
2318{
2319 let project_id = ProjectId::from_proto(request.remote_entity_id());
2320 let host_connection_id = session
2321 .db()
2322 .await
2323 .host_for_read_only_project_request(project_id, session.connection_id)
2324 .await?;
2325 let payload = session.forward_request(host_connection_id, request).await?;
2326 response.send(payload)?;
2327 Ok(())
2328}
2329
2330/// forward a project request to the host. These requests are disallowed
2331/// for guests.
2332async fn forward_mutating_project_request<T>(
2333 request: T,
2334 response: Response<T>,
2335 session: MessageContext,
2336) -> Result<()>
2337where
2338 T: EntityMessage + RequestMessage,
2339{
2340 let project_id = ProjectId::from_proto(request.remote_entity_id());
2341
2342 let host_connection_id = session
2343 .db()
2344 .await
2345 .host_for_mutating_project_request(project_id, session.connection_id)
2346 .await?;
2347 let payload = session.forward_request(host_connection_id, request).await?;
2348 response.send(payload)?;
2349 Ok(())
2350}
2351
2352async fn lsp_query(
2353 request: proto::LspQuery,
2354 response: Response<proto::LspQuery>,
2355 session: MessageContext,
2356) -> Result<()> {
2357 let (name, should_write) = request.query_name_and_write_permissions();
2358 tracing::Span::current().record("lsp_query_request", name);
2359 tracing::info!("lsp_query message received");
2360 if should_write {
2361 forward_mutating_project_request(request, response, session).await
2362 } else {
2363 forward_read_only_project_request(request, response, session).await
2364 }
2365}
2366
2367/// Notify other participants that a new buffer has been created
2368async fn create_buffer_for_peer(
2369 request: proto::CreateBufferForPeer,
2370 session: MessageContext,
2371) -> Result<()> {
2372 session
2373 .db()
2374 .await
2375 .check_user_is_project_host(
2376 ProjectId::from_proto(request.project_id),
2377 session.connection_id,
2378 )
2379 .await?;
2380 let peer_id = request.peer_id.context("invalid peer id")?;
2381 session
2382 .peer
2383 .forward_send(session.connection_id, peer_id.into(), request)?;
2384 Ok(())
2385}
2386
2387/// Notify other participants that a buffer has been updated. This is
2388/// allowed for guests as long as the update is limited to selections.
2389async fn update_buffer(
2390 request: proto::UpdateBuffer,
2391 response: Response<proto::UpdateBuffer>,
2392 session: MessageContext,
2393) -> Result<()> {
2394 let project_id = ProjectId::from_proto(request.project_id);
2395 let mut capability = Capability::ReadOnly;
2396
2397 for op in request.operations.iter() {
2398 match op.variant {
2399 None | Some(proto::operation::Variant::UpdateSelections(_)) => {}
2400 Some(_) => capability = Capability::ReadWrite,
2401 }
2402 }
2403
2404 let host = {
2405 let guard = session
2406 .db()
2407 .await
2408 .connections_for_buffer_update(project_id, session.connection_id, capability)
2409 .await?;
2410
2411 let (host, guests) = &*guard;
2412
2413 broadcast(
2414 Some(session.connection_id),
2415 guests.clone(),
2416 |connection_id| {
2417 session
2418 .peer
2419 .forward_send(session.connection_id, connection_id, request.clone())
2420 },
2421 );
2422
2423 *host
2424 };
2425
2426 if host != session.connection_id {
2427 session.forward_request(host, request.clone()).await?;
2428 }
2429
2430 response.send(proto::Ack {})?;
2431 Ok(())
2432}
2433
2434async fn update_context(message: proto::UpdateContext, session: MessageContext) -> Result<()> {
2435 let project_id = ProjectId::from_proto(message.project_id);
2436
2437 let operation = message.operation.as_ref().context("invalid operation")?;
2438 let capability = match operation.variant.as_ref() {
2439 Some(proto::context_operation::Variant::BufferOperation(buffer_op)) => {
2440 if let Some(buffer_op) = buffer_op.operation.as_ref() {
2441 match buffer_op.variant {
2442 None | Some(proto::operation::Variant::UpdateSelections(_)) => {
2443 Capability::ReadOnly
2444 }
2445 _ => Capability::ReadWrite,
2446 }
2447 } else {
2448 Capability::ReadWrite
2449 }
2450 }
2451 Some(_) => Capability::ReadWrite,
2452 None => Capability::ReadOnly,
2453 };
2454
2455 let guard = session
2456 .db()
2457 .await
2458 .connections_for_buffer_update(project_id, session.connection_id, capability)
2459 .await?;
2460
2461 let (host, guests) = &*guard;
2462
2463 broadcast(
2464 Some(session.connection_id),
2465 guests.iter().chain([host]).copied(),
2466 |connection_id| {
2467 session
2468 .peer
2469 .forward_send(session.connection_id, connection_id, message.clone())
2470 },
2471 );
2472
2473 Ok(())
2474}
2475
2476/// Notify other participants that a project has been updated.
2477async fn broadcast_project_message_from_host<T: EntityMessage<Entity = ShareProject>>(
2478 request: T,
2479 session: MessageContext,
2480) -> Result<()> {
2481 let project_id = ProjectId::from_proto(request.remote_entity_id());
2482 let project_connection_ids = session
2483 .db()
2484 .await
2485 .project_connection_ids(project_id, session.connection_id, false)
2486 .await?;
2487
2488 broadcast(
2489 Some(session.connection_id),
2490 project_connection_ids.iter().copied(),
2491 |connection_id| {
2492 session
2493 .peer
2494 .forward_send(session.connection_id, connection_id, request.clone())
2495 },
2496 );
2497 Ok(())
2498}
2499
2500/// Start following another user in a call.
2501async fn follow(
2502 request: proto::Follow,
2503 response: Response<proto::Follow>,
2504 session: MessageContext,
2505) -> Result<()> {
2506 let room_id = RoomId::from_proto(request.room_id);
2507 let project_id = request.project_id.map(ProjectId::from_proto);
2508 let leader_id = request.leader_id.context("invalid leader id")?.into();
2509 let follower_id = session.connection_id;
2510
2511 session
2512 .db()
2513 .await
2514 .check_room_participants(room_id, leader_id, session.connection_id)
2515 .await?;
2516
2517 let response_payload = session.forward_request(leader_id, request).await?;
2518 response.send(response_payload)?;
2519
2520 if let Some(project_id) = project_id {
2521 let room = session
2522 .db()
2523 .await
2524 .follow(room_id, project_id, leader_id, follower_id)
2525 .await?;
2526 room_updated(&room, &session.peer);
2527 }
2528
2529 Ok(())
2530}
2531
2532/// Stop following another user in a call.
2533async fn unfollow(request: proto::Unfollow, session: MessageContext) -> Result<()> {
2534 let room_id = RoomId::from_proto(request.room_id);
2535 let project_id = request.project_id.map(ProjectId::from_proto);
2536 let leader_id = request.leader_id.context("invalid leader id")?.into();
2537 let follower_id = session.connection_id;
2538
2539 session
2540 .db()
2541 .await
2542 .check_room_participants(room_id, leader_id, session.connection_id)
2543 .await?;
2544
2545 session
2546 .peer
2547 .forward_send(session.connection_id, leader_id, request)?;
2548
2549 if let Some(project_id) = project_id {
2550 let room = session
2551 .db()
2552 .await
2553 .unfollow(room_id, project_id, leader_id, follower_id)
2554 .await?;
2555 room_updated(&room, &session.peer);
2556 }
2557
2558 Ok(())
2559}
2560
2561/// Notify everyone following you of your current location.
2562async fn update_followers(request: proto::UpdateFollowers, session: MessageContext) -> Result<()> {
2563 let room_id = RoomId::from_proto(request.room_id);
2564 let database = session.db.lock().await;
2565
2566 let connection_ids = if let Some(project_id) = request.project_id {
2567 let project_id = ProjectId::from_proto(project_id);
2568 database
2569 .project_connection_ids(project_id, session.connection_id, true)
2570 .await?
2571 } else {
2572 database
2573 .room_connection_ids(room_id, session.connection_id)
2574 .await?
2575 };
2576
2577 // For now, don't send view update messages back to that view's current leader.
2578 let peer_id_to_omit = request.variant.as_ref().and_then(|variant| match variant {
2579 proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
2580 _ => None,
2581 });
2582
2583 for connection_id in connection_ids.iter().cloned() {
2584 if Some(connection_id.into()) != peer_id_to_omit && connection_id != session.connection_id {
2585 session
2586 .peer
2587 .forward_send(session.connection_id, connection_id, request.clone())?;
2588 }
2589 }
2590 Ok(())
2591}
2592
2593/// Get public data about users.
2594async fn get_users(
2595 request: proto::GetUsers,
2596 response: Response<proto::GetUsers>,
2597 session: MessageContext,
2598) -> Result<()> {
2599 let user_ids = request
2600 .user_ids
2601 .into_iter()
2602 .map(UserId::from_proto)
2603 .collect();
2604 let users = session
2605 .db()
2606 .await
2607 .get_users_by_ids(user_ids)
2608 .await?
2609 .into_iter()
2610 .map(|user| proto::User {
2611 id: user.id.to_proto(),
2612 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
2613 github_login: user.github_login,
2614 name: user.name,
2615 })
2616 .collect();
2617 response.send(proto::UsersResponse { users })?;
2618 Ok(())
2619}
2620
2621/// Search for users (to invite) buy Github login
2622async fn fuzzy_search_users(
2623 request: proto::FuzzySearchUsers,
2624 response: Response<proto::FuzzySearchUsers>,
2625 session: MessageContext,
2626) -> Result<()> {
2627 let query = request.query;
2628 let users = match query.len() {
2629 0 => vec![],
2630 1 | 2 => session
2631 .db()
2632 .await
2633 .get_user_by_github_login(&query)
2634 .await?
2635 .into_iter()
2636 .collect(),
2637 _ => session.db().await.fuzzy_search_users(&query, 10).await?,
2638 };
2639 let users = users
2640 .into_iter()
2641 .filter(|user| user.id != session.user_id())
2642 .map(|user| proto::User {
2643 id: user.id.to_proto(),
2644 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
2645 github_login: user.github_login,
2646 name: user.name,
2647 })
2648 .collect();
2649 response.send(proto::UsersResponse { users })?;
2650 Ok(())
2651}
2652
2653/// Send a contact request to another user.
2654async fn request_contact(
2655 request: proto::RequestContact,
2656 response: Response<proto::RequestContact>,
2657 session: MessageContext,
2658) -> Result<()> {
2659 let requester_id = session.user_id();
2660 let responder_id = UserId::from_proto(request.responder_id);
2661 if requester_id == responder_id {
2662 return Err(anyhow!("cannot add yourself as a contact"))?;
2663 }
2664
2665 let notifications = session
2666 .db()
2667 .await
2668 .send_contact_request(requester_id, responder_id)
2669 .await?;
2670
2671 // Update outgoing contact requests of requester
2672 let mut update = proto::UpdateContacts::default();
2673 update.outgoing_requests.push(responder_id.to_proto());
2674 for connection_id in session
2675 .connection_pool()
2676 .await
2677 .user_connection_ids(requester_id)
2678 {
2679 session.peer.send(connection_id, update.clone())?;
2680 }
2681
2682 // Update incoming contact requests of responder
2683 let mut update = proto::UpdateContacts::default();
2684 update
2685 .incoming_requests
2686 .push(proto::IncomingContactRequest {
2687 requester_id: requester_id.to_proto(),
2688 });
2689 let connection_pool = session.connection_pool().await;
2690 for connection_id in connection_pool.user_connection_ids(responder_id) {
2691 session.peer.send(connection_id, update.clone())?;
2692 }
2693
2694 send_notifications(&connection_pool, &session.peer, notifications);
2695
2696 response.send(proto::Ack {})?;
2697 Ok(())
2698}
2699
2700/// Accept or decline a contact request
2701async fn respond_to_contact_request(
2702 request: proto::RespondToContactRequest,
2703 response: Response<proto::RespondToContactRequest>,
2704 session: MessageContext,
2705) -> Result<()> {
2706 let responder_id = session.user_id();
2707 let requester_id = UserId::from_proto(request.requester_id);
2708 let db = session.db().await;
2709 if request.response == proto::ContactRequestResponse::Dismiss as i32 {
2710 db.dismiss_contact_notification(responder_id, requester_id)
2711 .await?;
2712 } else {
2713 let accept = request.response == proto::ContactRequestResponse::Accept as i32;
2714
2715 let notifications = db
2716 .respond_to_contact_request(responder_id, requester_id, accept)
2717 .await?;
2718 let requester_busy = db.is_user_busy(requester_id).await?;
2719 let responder_busy = db.is_user_busy(responder_id).await?;
2720
2721 let pool = session.connection_pool().await;
2722 // Update responder with new contact
2723 let mut update = proto::UpdateContacts::default();
2724 if accept {
2725 update
2726 .contacts
2727 .push(contact_for_user(requester_id, requester_busy, &pool));
2728 }
2729 update
2730 .remove_incoming_requests
2731 .push(requester_id.to_proto());
2732 for connection_id in pool.user_connection_ids(responder_id) {
2733 session.peer.send(connection_id, update.clone())?;
2734 }
2735
2736 // Update requester with new contact
2737 let mut update = proto::UpdateContacts::default();
2738 if accept {
2739 update
2740 .contacts
2741 .push(contact_for_user(responder_id, responder_busy, &pool));
2742 }
2743 update
2744 .remove_outgoing_requests
2745 .push(responder_id.to_proto());
2746
2747 for connection_id in pool.user_connection_ids(requester_id) {
2748 session.peer.send(connection_id, update.clone())?;
2749 }
2750
2751 send_notifications(&pool, &session.peer, notifications);
2752 }
2753
2754 response.send(proto::Ack {})?;
2755 Ok(())
2756}
2757
2758/// Remove a contact.
2759async fn remove_contact(
2760 request: proto::RemoveContact,
2761 response: Response<proto::RemoveContact>,
2762 session: MessageContext,
2763) -> Result<()> {
2764 let requester_id = session.user_id();
2765 let responder_id = UserId::from_proto(request.user_id);
2766 let db = session.db().await;
2767 let (contact_accepted, deleted_notification_id) =
2768 db.remove_contact(requester_id, responder_id).await?;
2769
2770 let pool = session.connection_pool().await;
2771 // Update outgoing contact requests of requester
2772 let mut update = proto::UpdateContacts::default();
2773 if contact_accepted {
2774 update.remove_contacts.push(responder_id.to_proto());
2775 } else {
2776 update
2777 .remove_outgoing_requests
2778 .push(responder_id.to_proto());
2779 }
2780 for connection_id in pool.user_connection_ids(requester_id) {
2781 session.peer.send(connection_id, update.clone())?;
2782 }
2783
2784 // Update incoming contact requests of responder
2785 let mut update = proto::UpdateContacts::default();
2786 if contact_accepted {
2787 update.remove_contacts.push(requester_id.to_proto());
2788 } else {
2789 update
2790 .remove_incoming_requests
2791 .push(requester_id.to_proto());
2792 }
2793 for connection_id in pool.user_connection_ids(responder_id) {
2794 session.peer.send(connection_id, update.clone())?;
2795 if let Some(notification_id) = deleted_notification_id {
2796 session.peer.send(
2797 connection_id,
2798 proto::DeleteNotification {
2799 notification_id: notification_id.to_proto(),
2800 },
2801 )?;
2802 }
2803 }
2804
2805 response.send(proto::Ack {})?;
2806 Ok(())
2807}
2808
2809fn should_auto_subscribe_to_channels(version: ZedVersion) -> bool {
2810 version.0.minor() < 139
2811}
2812
2813async fn subscribe_to_channels(
2814 _: proto::SubscribeToChannels,
2815 session: MessageContext,
2816) -> Result<()> {
2817 subscribe_user_to_channels(session.user_id(), &session).await?;
2818 Ok(())
2819}
2820
2821async fn subscribe_user_to_channels(user_id: UserId, session: &Session) -> Result<(), Error> {
2822 let channels_for_user = session.db().await.get_channels_for_user(user_id).await?;
2823 let mut pool = session.connection_pool().await;
2824 for membership in &channels_for_user.channel_memberships {
2825 pool.subscribe_to_channel(user_id, membership.channel_id, membership.role)
2826 }
2827 session.peer.send(
2828 session.connection_id,
2829 build_update_user_channels(&channels_for_user),
2830 )?;
2831 session.peer.send(
2832 session.connection_id,
2833 build_channels_update(channels_for_user),
2834 )?;
2835 Ok(())
2836}
2837
2838/// Creates a new channel.
2839async fn create_channel(
2840 request: proto::CreateChannel,
2841 response: Response<proto::CreateChannel>,
2842 session: MessageContext,
2843) -> Result<()> {
2844 let db = session.db().await;
2845
2846 let parent_id = request.parent_id.map(ChannelId::from_proto);
2847 let (channel, membership) = db
2848 .create_channel(&request.name, parent_id, session.user_id())
2849 .await?;
2850
2851 let root_id = channel.root_id();
2852 let channel = Channel::from_model(channel);
2853
2854 response.send(proto::CreateChannelResponse {
2855 channel: Some(channel.to_proto()),
2856 parent_id: request.parent_id,
2857 })?;
2858
2859 let mut connection_pool = session.connection_pool().await;
2860 if let Some(membership) = membership {
2861 connection_pool.subscribe_to_channel(
2862 membership.user_id,
2863 membership.channel_id,
2864 membership.role,
2865 );
2866 let update = proto::UpdateUserChannels {
2867 channel_memberships: vec![proto::ChannelMembership {
2868 channel_id: membership.channel_id.to_proto(),
2869 role: membership.role.into(),
2870 }],
2871 ..Default::default()
2872 };
2873 for connection_id in connection_pool.user_connection_ids(membership.user_id) {
2874 session.peer.send(connection_id, update.clone())?;
2875 }
2876 }
2877
2878 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
2879 if !role.can_see_channel(channel.visibility) {
2880 continue;
2881 }
2882
2883 let update = proto::UpdateChannels {
2884 channels: vec![channel.to_proto()],
2885 ..Default::default()
2886 };
2887 session.peer.send(connection_id, update.clone())?;
2888 }
2889
2890 Ok(())
2891}
2892
2893/// Delete a channel
2894async fn delete_channel(
2895 request: proto::DeleteChannel,
2896 response: Response<proto::DeleteChannel>,
2897 session: MessageContext,
2898) -> Result<()> {
2899 let db = session.db().await;
2900
2901 let channel_id = request.channel_id;
2902 let (root_channel, removed_channels) = db
2903 .delete_channel(ChannelId::from_proto(channel_id), session.user_id())
2904 .await?;
2905 response.send(proto::Ack {})?;
2906
2907 // Notify members of removed channels
2908 let mut update = proto::UpdateChannels::default();
2909 update
2910 .delete_channels
2911 .extend(removed_channels.into_iter().map(|id| id.to_proto()));
2912
2913 let connection_pool = session.connection_pool().await;
2914 for (connection_id, _) in connection_pool.channel_connection_ids(root_channel) {
2915 session.peer.send(connection_id, update.clone())?;
2916 }
2917
2918 Ok(())
2919}
2920
2921/// Invite someone to join a channel.
2922async fn invite_channel_member(
2923 request: proto::InviteChannelMember,
2924 response: Response<proto::InviteChannelMember>,
2925 session: MessageContext,
2926) -> Result<()> {
2927 let db = session.db().await;
2928 let channel_id = ChannelId::from_proto(request.channel_id);
2929 let invitee_id = UserId::from_proto(request.user_id);
2930 let InviteMemberResult {
2931 channel,
2932 notifications,
2933 } = db
2934 .invite_channel_member(
2935 channel_id,
2936 invitee_id,
2937 session.user_id(),
2938 request.role().into(),
2939 )
2940 .await?;
2941
2942 let update = proto::UpdateChannels {
2943 channel_invitations: vec![channel.to_proto()],
2944 ..Default::default()
2945 };
2946
2947 let connection_pool = session.connection_pool().await;
2948 for connection_id in connection_pool.user_connection_ids(invitee_id) {
2949 session.peer.send(connection_id, update.clone())?;
2950 }
2951
2952 send_notifications(&connection_pool, &session.peer, notifications);
2953
2954 response.send(proto::Ack {})?;
2955 Ok(())
2956}
2957
2958/// remove someone from a channel
2959async fn remove_channel_member(
2960 request: proto::RemoveChannelMember,
2961 response: Response<proto::RemoveChannelMember>,
2962 session: MessageContext,
2963) -> Result<()> {
2964 let db = session.db().await;
2965 let channel_id = ChannelId::from_proto(request.channel_id);
2966 let member_id = UserId::from_proto(request.user_id);
2967
2968 let RemoveChannelMemberResult {
2969 membership_update,
2970 notification_id,
2971 } = db
2972 .remove_channel_member(channel_id, member_id, session.user_id())
2973 .await?;
2974
2975 let mut connection_pool = session.connection_pool().await;
2976 notify_membership_updated(
2977 &mut connection_pool,
2978 membership_update,
2979 member_id,
2980 &session.peer,
2981 );
2982 for connection_id in connection_pool.user_connection_ids(member_id) {
2983 if let Some(notification_id) = notification_id {
2984 session
2985 .peer
2986 .send(
2987 connection_id,
2988 proto::DeleteNotification {
2989 notification_id: notification_id.to_proto(),
2990 },
2991 )
2992 .trace_err();
2993 }
2994 }
2995
2996 response.send(proto::Ack {})?;
2997 Ok(())
2998}
2999
3000/// Toggle the channel between public and private.
3001/// Care is taken to maintain the invariant that public channels only descend from public channels,
3002/// (though members-only channels can appear at any point in the hierarchy).
3003async fn set_channel_visibility(
3004 request: proto::SetChannelVisibility,
3005 response: Response<proto::SetChannelVisibility>,
3006 session: MessageContext,
3007) -> Result<()> {
3008 let db = session.db().await;
3009 let channel_id = ChannelId::from_proto(request.channel_id);
3010 let visibility = request.visibility().into();
3011
3012 let channel_model = db
3013 .set_channel_visibility(channel_id, visibility, session.user_id())
3014 .await?;
3015 let root_id = channel_model.root_id();
3016 let channel = Channel::from_model(channel_model);
3017
3018 let mut connection_pool = session.connection_pool().await;
3019 for (user_id, role) in connection_pool
3020 .channel_user_ids(root_id)
3021 .collect::<Vec<_>>()
3022 .into_iter()
3023 {
3024 let update = if role.can_see_channel(channel.visibility) {
3025 connection_pool.subscribe_to_channel(user_id, channel_id, role);
3026 proto::UpdateChannels {
3027 channels: vec![channel.to_proto()],
3028 ..Default::default()
3029 }
3030 } else {
3031 connection_pool.unsubscribe_from_channel(&user_id, &channel_id);
3032 proto::UpdateChannels {
3033 delete_channels: vec![channel.id.to_proto()],
3034 ..Default::default()
3035 }
3036 };
3037
3038 for connection_id in connection_pool.user_connection_ids(user_id) {
3039 session.peer.send(connection_id, update.clone())?;
3040 }
3041 }
3042
3043 response.send(proto::Ack {})?;
3044 Ok(())
3045}
3046
3047/// Alter the role for a user in the channel.
3048async fn set_channel_member_role(
3049 request: proto::SetChannelMemberRole,
3050 response: Response<proto::SetChannelMemberRole>,
3051 session: MessageContext,
3052) -> Result<()> {
3053 let db = session.db().await;
3054 let channel_id = ChannelId::from_proto(request.channel_id);
3055 let member_id = UserId::from_proto(request.user_id);
3056 let result = db
3057 .set_channel_member_role(
3058 channel_id,
3059 session.user_id(),
3060 member_id,
3061 request.role().into(),
3062 )
3063 .await?;
3064
3065 match result {
3066 db::SetMemberRoleResult::MembershipUpdated(membership_update) => {
3067 let mut connection_pool = session.connection_pool().await;
3068 notify_membership_updated(
3069 &mut connection_pool,
3070 membership_update,
3071 member_id,
3072 &session.peer,
3073 )
3074 }
3075 db::SetMemberRoleResult::InviteUpdated(channel) => {
3076 let update = proto::UpdateChannels {
3077 channel_invitations: vec![channel.to_proto()],
3078 ..Default::default()
3079 };
3080
3081 for connection_id in session
3082 .connection_pool()
3083 .await
3084 .user_connection_ids(member_id)
3085 {
3086 session.peer.send(connection_id, update.clone())?;
3087 }
3088 }
3089 }
3090
3091 response.send(proto::Ack {})?;
3092 Ok(())
3093}
3094
3095/// Change the name of a channel
3096async fn rename_channel(
3097 request: proto::RenameChannel,
3098 response: Response<proto::RenameChannel>,
3099 session: MessageContext,
3100) -> Result<()> {
3101 let db = session.db().await;
3102 let channel_id = ChannelId::from_proto(request.channel_id);
3103 let channel_model = db
3104 .rename_channel(channel_id, session.user_id(), &request.name)
3105 .await?;
3106 let root_id = channel_model.root_id();
3107 let channel = Channel::from_model(channel_model);
3108
3109 response.send(proto::RenameChannelResponse {
3110 channel: Some(channel.to_proto()),
3111 })?;
3112
3113 let connection_pool = session.connection_pool().await;
3114 let update = proto::UpdateChannels {
3115 channels: vec![channel.to_proto()],
3116 ..Default::default()
3117 };
3118 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
3119 if role.can_see_channel(channel.visibility) {
3120 session.peer.send(connection_id, update.clone())?;
3121 }
3122 }
3123
3124 Ok(())
3125}
3126
3127/// Move a channel to a new parent.
3128async fn move_channel(
3129 request: proto::MoveChannel,
3130 response: Response<proto::MoveChannel>,
3131 session: MessageContext,
3132) -> Result<()> {
3133 let channel_id = ChannelId::from_proto(request.channel_id);
3134 let to = ChannelId::from_proto(request.to);
3135
3136 let (root_id, channels) = session
3137 .db()
3138 .await
3139 .move_channel(channel_id, to, session.user_id())
3140 .await?;
3141
3142 let connection_pool = session.connection_pool().await;
3143 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
3144 let channels = channels
3145 .iter()
3146 .filter_map(|channel| {
3147 if role.can_see_channel(channel.visibility) {
3148 Some(channel.to_proto())
3149 } else {
3150 None
3151 }
3152 })
3153 .collect::<Vec<_>>();
3154 if channels.is_empty() {
3155 continue;
3156 }
3157
3158 let update = proto::UpdateChannels {
3159 channels,
3160 ..Default::default()
3161 };
3162
3163 session.peer.send(connection_id, update.clone())?;
3164 }
3165
3166 response.send(Ack {})?;
3167 Ok(())
3168}
3169
3170async fn reorder_channel(
3171 request: proto::ReorderChannel,
3172 response: Response<proto::ReorderChannel>,
3173 session: MessageContext,
3174) -> Result<()> {
3175 let channel_id = ChannelId::from_proto(request.channel_id);
3176 let direction = request.direction();
3177
3178 let updated_channels = session
3179 .db()
3180 .await
3181 .reorder_channel(channel_id, direction, session.user_id())
3182 .await?;
3183
3184 if let Some(root_id) = updated_channels.first().map(|channel| channel.root_id()) {
3185 let connection_pool = session.connection_pool().await;
3186 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
3187 let channels = updated_channels
3188 .iter()
3189 .filter_map(|channel| {
3190 if role.can_see_channel(channel.visibility) {
3191 Some(channel.to_proto())
3192 } else {
3193 None
3194 }
3195 })
3196 .collect::<Vec<_>>();
3197
3198 if channels.is_empty() {
3199 continue;
3200 }
3201
3202 let update = proto::UpdateChannels {
3203 channels,
3204 ..Default::default()
3205 };
3206
3207 session.peer.send(connection_id, update.clone())?;
3208 }
3209 }
3210
3211 response.send(Ack {})?;
3212 Ok(())
3213}
3214
3215/// Get the list of channel members
3216async fn get_channel_members(
3217 request: proto::GetChannelMembers,
3218 response: Response<proto::GetChannelMembers>,
3219 session: MessageContext,
3220) -> Result<()> {
3221 let db = session.db().await;
3222 let channel_id = ChannelId::from_proto(request.channel_id);
3223 let limit = if request.limit == 0 {
3224 u16::MAX as u64
3225 } else {
3226 request.limit
3227 };
3228 let (members, users) = db
3229 .get_channel_participant_details(channel_id, &request.query, limit, session.user_id())
3230 .await?;
3231 response.send(proto::GetChannelMembersResponse { members, users })?;
3232 Ok(())
3233}
3234
3235/// Accept or decline a channel invitation.
3236async fn respond_to_channel_invite(
3237 request: proto::RespondToChannelInvite,
3238 response: Response<proto::RespondToChannelInvite>,
3239 session: MessageContext,
3240) -> Result<()> {
3241 let db = session.db().await;
3242 let channel_id = ChannelId::from_proto(request.channel_id);
3243 let RespondToChannelInvite {
3244 membership_update,
3245 notifications,
3246 } = db
3247 .respond_to_channel_invite(channel_id, session.user_id(), request.accept)
3248 .await?;
3249
3250 let mut connection_pool = session.connection_pool().await;
3251 if let Some(membership_update) = membership_update {
3252 notify_membership_updated(
3253 &mut connection_pool,
3254 membership_update,
3255 session.user_id(),
3256 &session.peer,
3257 );
3258 } else {
3259 let update = proto::UpdateChannels {
3260 remove_channel_invitations: vec![channel_id.to_proto()],
3261 ..Default::default()
3262 };
3263
3264 for connection_id in connection_pool.user_connection_ids(session.user_id()) {
3265 session.peer.send(connection_id, update.clone())?;
3266 }
3267 };
3268
3269 send_notifications(&connection_pool, &session.peer, notifications);
3270
3271 response.send(proto::Ack {})?;
3272
3273 Ok(())
3274}
3275
3276/// Join the channels' room
3277async fn join_channel(
3278 request: proto::JoinChannel,
3279 response: Response<proto::JoinChannel>,
3280 session: MessageContext,
3281) -> Result<()> {
3282 let channel_id = ChannelId::from_proto(request.channel_id);
3283 join_channel_internal(channel_id, Box::new(response), session).await
3284}
3285
3286trait JoinChannelInternalResponse {
3287 fn send(self, result: proto::JoinRoomResponse) -> Result<()>;
3288}
3289impl JoinChannelInternalResponse for Response<proto::JoinChannel> {
3290 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
3291 Response::<proto::JoinChannel>::send(self, result)
3292 }
3293}
3294impl JoinChannelInternalResponse for Response<proto::JoinRoom> {
3295 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
3296 Response::<proto::JoinRoom>::send(self, result)
3297 }
3298}
3299
3300async fn join_channel_internal(
3301 channel_id: ChannelId,
3302 response: Box<impl JoinChannelInternalResponse>,
3303 session: MessageContext,
3304) -> Result<()> {
3305 let joined_room = {
3306 let mut db = session.db().await;
3307 // If zed quits without leaving the room, and the user re-opens zed before the
3308 // RECONNECT_TIMEOUT, we need to make sure that we kick the user out of the previous
3309 // room they were in.
3310 if let Some(connection) = db.stale_room_connection(session.user_id()).await? {
3311 tracing::info!(
3312 stale_connection_id = %connection,
3313 "cleaning up stale connection",
3314 );
3315 drop(db);
3316 leave_room_for_session(&session, connection).await?;
3317 db = session.db().await;
3318 }
3319
3320 let (joined_room, membership_updated, role) = db
3321 .join_channel(channel_id, session.user_id(), session.connection_id)
3322 .await?;
3323
3324 let live_kit_connection_info =
3325 session
3326 .app_state
3327 .livekit_client
3328 .as_ref()
3329 .and_then(|live_kit| {
3330 let (can_publish, token) = if role == ChannelRole::Guest {
3331 (
3332 false,
3333 live_kit
3334 .guest_token(
3335 &joined_room.room.livekit_room,
3336 &session.user_id().to_string(),
3337 )
3338 .trace_err()?,
3339 )
3340 } else {
3341 (
3342 true,
3343 live_kit
3344 .room_token(
3345 &joined_room.room.livekit_room,
3346 &session.user_id().to_string(),
3347 )
3348 .trace_err()?,
3349 )
3350 };
3351
3352 Some(LiveKitConnectionInfo {
3353 server_url: live_kit.url().into(),
3354 token,
3355 can_publish,
3356 })
3357 });
3358
3359 response.send(proto::JoinRoomResponse {
3360 room: Some(joined_room.room.clone()),
3361 channel_id: joined_room
3362 .channel
3363 .as_ref()
3364 .map(|channel| channel.id.to_proto()),
3365 live_kit_connection_info,
3366 })?;
3367
3368 let mut connection_pool = session.connection_pool().await;
3369 if let Some(membership_updated) = membership_updated {
3370 notify_membership_updated(
3371 &mut connection_pool,
3372 membership_updated,
3373 session.user_id(),
3374 &session.peer,
3375 );
3376 }
3377
3378 room_updated(&joined_room.room, &session.peer);
3379
3380 joined_room
3381 };
3382
3383 channel_updated(
3384 &joined_room.channel.context("channel not returned")?,
3385 &joined_room.room,
3386 &session.peer,
3387 &*session.connection_pool().await,
3388 );
3389
3390 update_user_contacts(session.user_id(), &session).await?;
3391 Ok(())
3392}
3393
3394/// Start editing the channel notes
3395async fn join_channel_buffer(
3396 request: proto::JoinChannelBuffer,
3397 response: Response<proto::JoinChannelBuffer>,
3398 session: MessageContext,
3399) -> Result<()> {
3400 let db = session.db().await;
3401 let channel_id = ChannelId::from_proto(request.channel_id);
3402
3403 let open_response = db
3404 .join_channel_buffer(channel_id, session.user_id(), session.connection_id)
3405 .await?;
3406
3407 let collaborators = open_response.collaborators.clone();
3408 response.send(open_response)?;
3409
3410 let update = UpdateChannelBufferCollaborators {
3411 channel_id: channel_id.to_proto(),
3412 collaborators: collaborators.clone(),
3413 };
3414 channel_buffer_updated(
3415 session.connection_id,
3416 collaborators
3417 .iter()
3418 .filter_map(|collaborator| Some(collaborator.peer_id?.into())),
3419 &update,
3420 &session.peer,
3421 );
3422
3423 Ok(())
3424}
3425
3426/// Edit the channel notes
3427async fn update_channel_buffer(
3428 request: proto::UpdateChannelBuffer,
3429 session: MessageContext,
3430) -> Result<()> {
3431 let db = session.db().await;
3432 let channel_id = ChannelId::from_proto(request.channel_id);
3433
3434 let (collaborators, epoch, version) = db
3435 .update_channel_buffer(channel_id, session.user_id(), &request.operations)
3436 .await?;
3437
3438 channel_buffer_updated(
3439 session.connection_id,
3440 collaborators.clone(),
3441 &proto::UpdateChannelBuffer {
3442 channel_id: channel_id.to_proto(),
3443 operations: request.operations,
3444 },
3445 &session.peer,
3446 );
3447
3448 let pool = &*session.connection_pool().await;
3449
3450 let non_collaborators =
3451 pool.channel_connection_ids(channel_id)
3452 .filter_map(|(connection_id, _)| {
3453 if collaborators.contains(&connection_id) {
3454 None
3455 } else {
3456 Some(connection_id)
3457 }
3458 });
3459
3460 broadcast(None, non_collaborators, |peer_id| {
3461 session.peer.send(
3462 peer_id,
3463 proto::UpdateChannels {
3464 latest_channel_buffer_versions: vec![proto::ChannelBufferVersion {
3465 channel_id: channel_id.to_proto(),
3466 epoch: epoch as u64,
3467 version: version.clone(),
3468 }],
3469 ..Default::default()
3470 },
3471 )
3472 });
3473
3474 Ok(())
3475}
3476
3477/// Rejoin the channel notes after a connection blip
3478async fn rejoin_channel_buffers(
3479 request: proto::RejoinChannelBuffers,
3480 response: Response<proto::RejoinChannelBuffers>,
3481 session: MessageContext,
3482) -> Result<()> {
3483 let db = session.db().await;
3484 let buffers = db
3485 .rejoin_channel_buffers(&request.buffers, session.user_id(), session.connection_id)
3486 .await?;
3487
3488 for rejoined_buffer in &buffers {
3489 let collaborators_to_notify = rejoined_buffer
3490 .buffer
3491 .collaborators
3492 .iter()
3493 .filter_map(|c| Some(c.peer_id?.into()));
3494 channel_buffer_updated(
3495 session.connection_id,
3496 collaborators_to_notify,
3497 &proto::UpdateChannelBufferCollaborators {
3498 channel_id: rejoined_buffer.buffer.channel_id,
3499 collaborators: rejoined_buffer.buffer.collaborators.clone(),
3500 },
3501 &session.peer,
3502 );
3503 }
3504
3505 response.send(proto::RejoinChannelBuffersResponse {
3506 buffers: buffers.into_iter().map(|b| b.buffer).collect(),
3507 })?;
3508
3509 Ok(())
3510}
3511
3512/// Stop editing the channel notes
3513async fn leave_channel_buffer(
3514 request: proto::LeaveChannelBuffer,
3515 response: Response<proto::LeaveChannelBuffer>,
3516 session: MessageContext,
3517) -> Result<()> {
3518 let db = session.db().await;
3519 let channel_id = ChannelId::from_proto(request.channel_id);
3520
3521 let left_buffer = db
3522 .leave_channel_buffer(channel_id, session.connection_id)
3523 .await?;
3524
3525 response.send(Ack {})?;
3526
3527 channel_buffer_updated(
3528 session.connection_id,
3529 left_buffer.connections,
3530 &proto::UpdateChannelBufferCollaborators {
3531 channel_id: channel_id.to_proto(),
3532 collaborators: left_buffer.collaborators,
3533 },
3534 &session.peer,
3535 );
3536
3537 Ok(())
3538}
3539
3540fn channel_buffer_updated<T: EnvelopedMessage>(
3541 sender_id: ConnectionId,
3542 collaborators: impl IntoIterator<Item = ConnectionId>,
3543 message: &T,
3544 peer: &Peer,
3545) {
3546 broadcast(Some(sender_id), collaborators, |peer_id| {
3547 peer.send(peer_id, message.clone())
3548 });
3549}
3550
3551fn send_notifications(
3552 connection_pool: &ConnectionPool,
3553 peer: &Peer,
3554 notifications: db::NotificationBatch,
3555) {
3556 for (user_id, notification) in notifications {
3557 for connection_id in connection_pool.user_connection_ids(user_id) {
3558 if let Err(error) = peer.send(
3559 connection_id,
3560 proto::AddNotification {
3561 notification: Some(notification.clone()),
3562 },
3563 ) {
3564 tracing::error!(
3565 "failed to send notification to {:?} {}",
3566 connection_id,
3567 error
3568 );
3569 }
3570 }
3571 }
3572}
3573
3574/// Send a message to the channel
3575async fn send_channel_message(
3576 _request: proto::SendChannelMessage,
3577 _response: Response<proto::SendChannelMessage>,
3578 _session: MessageContext,
3579) -> Result<()> {
3580 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3581}
3582
3583/// Delete a channel message
3584async fn remove_channel_message(
3585 _request: proto::RemoveChannelMessage,
3586 _response: Response<proto::RemoveChannelMessage>,
3587 _session: MessageContext,
3588) -> Result<()> {
3589 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3590}
3591
3592async fn update_channel_message(
3593 _request: proto::UpdateChannelMessage,
3594 _response: Response<proto::UpdateChannelMessage>,
3595 _session: MessageContext,
3596) -> Result<()> {
3597 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3598}
3599
3600/// Mark a channel message as read
3601async fn acknowledge_channel_message(
3602 _request: proto::AckChannelMessage,
3603 _session: MessageContext,
3604) -> Result<()> {
3605 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3606}
3607
3608/// Mark a buffer version as synced
3609async fn acknowledge_buffer_version(
3610 request: proto::AckBufferOperation,
3611 session: MessageContext,
3612) -> Result<()> {
3613 let buffer_id = BufferId::from_proto(request.buffer_id);
3614 session
3615 .db()
3616 .await
3617 .observe_buffer_version(
3618 buffer_id,
3619 session.user_id(),
3620 request.epoch as i32,
3621 &request.version,
3622 )
3623 .await?;
3624 Ok(())
3625}
3626
3627/// Get a Supermaven API key for the user
3628async fn get_supermaven_api_key(
3629 _request: proto::GetSupermavenApiKey,
3630 response: Response<proto::GetSupermavenApiKey>,
3631 session: MessageContext,
3632) -> Result<()> {
3633 let user_id: String = session.user_id().to_string();
3634 if !session.is_staff() {
3635 return Err(anyhow!("supermaven not enabled for this account"))?;
3636 }
3637
3638 let email = session.email().context("user must have an email")?;
3639
3640 let supermaven_admin_api = session
3641 .supermaven_client
3642 .as_ref()
3643 .context("supermaven not configured")?;
3644
3645 let result = supermaven_admin_api
3646 .try_get_or_create_user(CreateExternalUserRequest { id: user_id, email })
3647 .await?;
3648
3649 response.send(proto::GetSupermavenApiKeyResponse {
3650 api_key: result.api_key,
3651 })?;
3652
3653 Ok(())
3654}
3655
3656/// Start receiving chat updates for a channel
3657async fn join_channel_chat(
3658 _request: proto::JoinChannelChat,
3659 _response: Response<proto::JoinChannelChat>,
3660 _session: MessageContext,
3661) -> Result<()> {
3662 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3663}
3664
3665/// Stop receiving chat updates for a channel
3666async fn leave_channel_chat(
3667 _request: proto::LeaveChannelChat,
3668 _session: MessageContext,
3669) -> Result<()> {
3670 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3671}
3672
3673/// Retrieve the chat history for a channel
3674async fn get_channel_messages(
3675 _request: proto::GetChannelMessages,
3676 _response: Response<proto::GetChannelMessages>,
3677 _session: MessageContext,
3678) -> Result<()> {
3679 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3680}
3681
3682/// Retrieve specific chat messages
3683async fn get_channel_messages_by_id(
3684 _request: proto::GetChannelMessagesById,
3685 _response: Response<proto::GetChannelMessagesById>,
3686 _session: MessageContext,
3687) -> Result<()> {
3688 Err(anyhow!("chat has been removed in the latest version of Zed").into())
3689}
3690
3691/// Retrieve the current users notifications
3692async fn get_notifications(
3693 request: proto::GetNotifications,
3694 response: Response<proto::GetNotifications>,
3695 session: MessageContext,
3696) -> Result<()> {
3697 let notifications = session
3698 .db()
3699 .await
3700 .get_notifications(
3701 session.user_id(),
3702 NOTIFICATION_COUNT_PER_PAGE,
3703 request.before_id.map(db::NotificationId::from_proto),
3704 )
3705 .await?;
3706 response.send(proto::GetNotificationsResponse {
3707 done: notifications.len() < NOTIFICATION_COUNT_PER_PAGE,
3708 notifications,
3709 })?;
3710 Ok(())
3711}
3712
3713/// Mark notifications as read
3714async fn mark_notification_as_read(
3715 request: proto::MarkNotificationRead,
3716 response: Response<proto::MarkNotificationRead>,
3717 session: MessageContext,
3718) -> Result<()> {
3719 let database = &session.db().await;
3720 let notifications = database
3721 .mark_notification_as_read_by_id(
3722 session.user_id(),
3723 NotificationId::from_proto(request.notification_id),
3724 )
3725 .await?;
3726 send_notifications(
3727 &*session.connection_pool().await,
3728 &session.peer,
3729 notifications,
3730 );
3731 response.send(proto::Ack {})?;
3732 Ok(())
3733}
3734
3735fn to_axum_message(message: TungsteniteMessage) -> anyhow::Result<AxumMessage> {
3736 let message = match message {
3737 TungsteniteMessage::Text(payload) => AxumMessage::Text(payload.as_str().to_string()),
3738 TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload.into()),
3739 TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload.into()),
3740 TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload.into()),
3741 TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
3742 code: frame.code.into(),
3743 reason: frame.reason.as_str().to_owned().into(),
3744 })),
3745 // We should never receive a frame while reading the message, according
3746 // to the `tungstenite` maintainers:
3747 //
3748 // > It cannot occur when you read messages from the WebSocket, but it
3749 // > can be used when you want to send the raw frames (e.g. you want to
3750 // > send the frames to the WebSocket without composing the full message first).
3751 // >
3752 // > — https://github.com/snapview/tungstenite-rs/issues/268
3753 TungsteniteMessage::Frame(_) => {
3754 bail!("received an unexpected frame while reading the message")
3755 }
3756 };
3757
3758 Ok(message)
3759}
3760
3761fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
3762 match message {
3763 AxumMessage::Text(payload) => TungsteniteMessage::Text(payload.into()),
3764 AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload.into()),
3765 AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload.into()),
3766 AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload.into()),
3767 AxumMessage::Close(frame) => {
3768 TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
3769 code: frame.code.into(),
3770 reason: frame.reason.as_ref().into(),
3771 }))
3772 }
3773 }
3774}
3775
3776fn notify_membership_updated(
3777 connection_pool: &mut ConnectionPool,
3778 result: MembershipUpdated,
3779 user_id: UserId,
3780 peer: &Peer,
3781) {
3782 for membership in &result.new_channels.channel_memberships {
3783 connection_pool.subscribe_to_channel(user_id, membership.channel_id, membership.role)
3784 }
3785 for channel_id in &result.removed_channels {
3786 connection_pool.unsubscribe_from_channel(&user_id, channel_id)
3787 }
3788
3789 let user_channels_update = proto::UpdateUserChannels {
3790 channel_memberships: result
3791 .new_channels
3792 .channel_memberships
3793 .iter()
3794 .map(|cm| proto::ChannelMembership {
3795 channel_id: cm.channel_id.to_proto(),
3796 role: cm.role.into(),
3797 })
3798 .collect(),
3799 ..Default::default()
3800 };
3801
3802 let mut update = build_channels_update(result.new_channels);
3803 update.delete_channels = result
3804 .removed_channels
3805 .into_iter()
3806 .map(|id| id.to_proto())
3807 .collect();
3808 update.remove_channel_invitations = vec![result.channel_id.to_proto()];
3809
3810 for connection_id in connection_pool.user_connection_ids(user_id) {
3811 peer.send(connection_id, user_channels_update.clone())
3812 .trace_err();
3813 peer.send(connection_id, update.clone()).trace_err();
3814 }
3815}
3816
3817fn build_update_user_channels(channels: &ChannelsForUser) -> proto::UpdateUserChannels {
3818 proto::UpdateUserChannels {
3819 channel_memberships: channels
3820 .channel_memberships
3821 .iter()
3822 .map(|m| proto::ChannelMembership {
3823 channel_id: m.channel_id.to_proto(),
3824 role: m.role.into(),
3825 })
3826 .collect(),
3827 observed_channel_buffer_version: channels.observed_buffer_versions.clone(),
3828 }
3829}
3830
3831fn build_channels_update(channels: ChannelsForUser) -> proto::UpdateChannels {
3832 let mut update = proto::UpdateChannels::default();
3833
3834 for channel in channels.channels {
3835 update.channels.push(channel.to_proto());
3836 }
3837
3838 update.latest_channel_buffer_versions = channels.latest_buffer_versions;
3839
3840 for (channel_id, participants) in channels.channel_participants {
3841 update
3842 .channel_participants
3843 .push(proto::ChannelParticipants {
3844 channel_id: channel_id.to_proto(),
3845 participant_user_ids: participants.into_iter().map(|id| id.to_proto()).collect(),
3846 });
3847 }
3848
3849 for channel in channels.invited_channels {
3850 update.channel_invitations.push(channel.to_proto());
3851 }
3852
3853 update
3854}
3855
3856fn build_initial_contacts_update(
3857 contacts: Vec<db::Contact>,
3858 pool: &ConnectionPool,
3859) -> proto::UpdateContacts {
3860 let mut update = proto::UpdateContacts::default();
3861
3862 for contact in contacts {
3863 match contact {
3864 db::Contact::Accepted { user_id, busy } => {
3865 update.contacts.push(contact_for_user(user_id, busy, pool));
3866 }
3867 db::Contact::Outgoing { user_id } => update.outgoing_requests.push(user_id.to_proto()),
3868 db::Contact::Incoming { user_id } => {
3869 update
3870 .incoming_requests
3871 .push(proto::IncomingContactRequest {
3872 requester_id: user_id.to_proto(),
3873 })
3874 }
3875 }
3876 }
3877
3878 update
3879}
3880
3881fn contact_for_user(user_id: UserId, busy: bool, pool: &ConnectionPool) -> proto::Contact {
3882 proto::Contact {
3883 user_id: user_id.to_proto(),
3884 online: pool.is_user_online(user_id),
3885 busy,
3886 }
3887}
3888
3889fn room_updated(room: &proto::Room, peer: &Peer) {
3890 broadcast(
3891 None,
3892 room.participants
3893 .iter()
3894 .filter_map(|participant| Some(participant.peer_id?.into())),
3895 |peer_id| {
3896 peer.send(
3897 peer_id,
3898 proto::RoomUpdated {
3899 room: Some(room.clone()),
3900 },
3901 )
3902 },
3903 );
3904}
3905
3906fn channel_updated(
3907 channel: &db::channel::Model,
3908 room: &proto::Room,
3909 peer: &Peer,
3910 pool: &ConnectionPool,
3911) {
3912 let participants = room
3913 .participants
3914 .iter()
3915 .map(|p| p.user_id)
3916 .collect::<Vec<_>>();
3917
3918 broadcast(
3919 None,
3920 pool.channel_connection_ids(channel.root_id())
3921 .filter_map(|(channel_id, role)| {
3922 role.can_see_channel(channel.visibility)
3923 .then_some(channel_id)
3924 }),
3925 |peer_id| {
3926 peer.send(
3927 peer_id,
3928 proto::UpdateChannels {
3929 channel_participants: vec![proto::ChannelParticipants {
3930 channel_id: channel.id.to_proto(),
3931 participant_user_ids: participants.clone(),
3932 }],
3933 ..Default::default()
3934 },
3935 )
3936 },
3937 );
3938}
3939
3940async fn update_user_contacts(user_id: UserId, session: &Session) -> Result<()> {
3941 let db = session.db().await;
3942
3943 let contacts = db.get_contacts(user_id).await?;
3944 let busy = db.is_user_busy(user_id).await?;
3945
3946 let pool = session.connection_pool().await;
3947 let updated_contact = contact_for_user(user_id, busy, &pool);
3948 for contact in contacts {
3949 if let db::Contact::Accepted {
3950 user_id: contact_user_id,
3951 ..
3952 } = contact
3953 {
3954 for contact_conn_id in pool.user_connection_ids(contact_user_id) {
3955 session
3956 .peer
3957 .send(
3958 contact_conn_id,
3959 proto::UpdateContacts {
3960 contacts: vec![updated_contact.clone()],
3961 remove_contacts: Default::default(),
3962 incoming_requests: Default::default(),
3963 remove_incoming_requests: Default::default(),
3964 outgoing_requests: Default::default(),
3965 remove_outgoing_requests: Default::default(),
3966 },
3967 )
3968 .trace_err();
3969 }
3970 }
3971 }
3972 Ok(())
3973}
3974
3975async fn leave_room_for_session(session: &Session, connection_id: ConnectionId) -> Result<()> {
3976 let mut contacts_to_update = HashSet::default();
3977
3978 let room_id;
3979 let canceled_calls_to_user_ids;
3980 let livekit_room;
3981 let delete_livekit_room;
3982 let room;
3983 let channel;
3984
3985 if let Some(mut left_room) = session.db().await.leave_room(connection_id).await? {
3986 contacts_to_update.insert(session.user_id());
3987
3988 for project in left_room.left_projects.values() {
3989 project_left(project, session);
3990 }
3991
3992 room_id = RoomId::from_proto(left_room.room.id);
3993 canceled_calls_to_user_ids = mem::take(&mut left_room.canceled_calls_to_user_ids);
3994 livekit_room = mem::take(&mut left_room.room.livekit_room);
3995 delete_livekit_room = left_room.deleted;
3996 room = mem::take(&mut left_room.room);
3997 channel = mem::take(&mut left_room.channel);
3998
3999 room_updated(&room, &session.peer);
4000 } else {
4001 return Ok(());
4002 }
4003
4004 if let Some(channel) = channel {
4005 channel_updated(
4006 &channel,
4007 &room,
4008 &session.peer,
4009 &*session.connection_pool().await,
4010 );
4011 }
4012
4013 {
4014 let pool = session.connection_pool().await;
4015 for canceled_user_id in canceled_calls_to_user_ids {
4016 for connection_id in pool.user_connection_ids(canceled_user_id) {
4017 session
4018 .peer
4019 .send(
4020 connection_id,
4021 proto::CallCanceled {
4022 room_id: room_id.to_proto(),
4023 },
4024 )
4025 .trace_err();
4026 }
4027 contacts_to_update.insert(canceled_user_id);
4028 }
4029 }
4030
4031 for contact_user_id in contacts_to_update {
4032 update_user_contacts(contact_user_id, session).await?;
4033 }
4034
4035 if let Some(live_kit) = session.app_state.livekit_client.as_ref() {
4036 live_kit
4037 .remove_participant(livekit_room.clone(), session.user_id().to_string())
4038 .await
4039 .trace_err();
4040
4041 if delete_livekit_room {
4042 live_kit.delete_room(livekit_room).await.trace_err();
4043 }
4044 }
4045
4046 Ok(())
4047}
4048
4049async fn leave_channel_buffers_for_session(session: &Session) -> Result<()> {
4050 let left_channel_buffers = session
4051 .db()
4052 .await
4053 .leave_channel_buffers(session.connection_id)
4054 .await?;
4055
4056 for left_buffer in left_channel_buffers {
4057 channel_buffer_updated(
4058 session.connection_id,
4059 left_buffer.connections,
4060 &proto::UpdateChannelBufferCollaborators {
4061 channel_id: left_buffer.channel_id.to_proto(),
4062 collaborators: left_buffer.collaborators,
4063 },
4064 &session.peer,
4065 );
4066 }
4067
4068 Ok(())
4069}
4070
4071fn project_left(project: &db::LeftProject, session: &Session) {
4072 for connection_id in &project.connection_ids {
4073 if project.should_unshare {
4074 session
4075 .peer
4076 .send(
4077 *connection_id,
4078 proto::UnshareProject {
4079 project_id: project.id.to_proto(),
4080 },
4081 )
4082 .trace_err();
4083 } else {
4084 session
4085 .peer
4086 .send(
4087 *connection_id,
4088 proto::RemoveProjectCollaborator {
4089 project_id: project.id.to_proto(),
4090 peer_id: Some(session.connection_id.into()),
4091 },
4092 )
4093 .trace_err();
4094 }
4095 }
4096}
4097
4098pub trait ResultExt {
4099 type Ok;
4100
4101 fn trace_err(self) -> Option<Self::Ok>;
4102}
4103
4104impl<T, E> ResultExt for Result<T, E>
4105where
4106 E: std::fmt::Debug,
4107{
4108 type Ok = T;
4109
4110 #[track_caller]
4111 fn trace_err(self) -> Option<T> {
4112 match self {
4113 Ok(value) => Some(value),
4114 Err(error) => {
4115 tracing::error!("{:?}", error);
4116 None
4117 }
4118 }
4119 }
4120}