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