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