1#[cfg(any(test, feature = "test-support"))]
2pub mod test;
3
4pub mod channel;
5pub mod http;
6pub mod telemetry;
7pub mod user;
8
9use anyhow::{anyhow, Context, Result};
10use async_recursion::async_recursion;
11use async_tungstenite::tungstenite::{
12 error::Error as WebsocketError,
13 http::{Request, StatusCode},
14};
15use db::Db;
16use futures::{future::LocalBoxFuture, FutureExt, SinkExt, StreamExt, TryStreamExt};
17use gpui::{
18 actions, serde_json::Value, AnyModelHandle, AnyViewHandle, AnyWeakModelHandle,
19 AnyWeakViewHandle, AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle,
20 MutableAppContext, Task, View, ViewContext, ViewHandle,
21};
22use http::HttpClient;
23use lazy_static::lazy_static;
24use parking_lot::RwLock;
25use postage::watch;
26use rand::prelude::*;
27use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, RequestMessage};
28use std::{
29 any::TypeId,
30 collections::HashMap,
31 convert::TryFrom,
32 fmt::Write as _,
33 future::Future,
34 path::PathBuf,
35 sync::{Arc, Weak},
36 time::{Duration, Instant},
37};
38use telemetry::Telemetry;
39use thiserror::Error;
40use url::Url;
41use util::{ResultExt, TryFutureExt};
42
43pub use channel::*;
44pub use rpc::*;
45pub use user::*;
46
47lazy_static! {
48 pub static ref ZED_SERVER_URL: String =
49 std::env::var("ZED_SERVER_URL").unwrap_or_else(|_| "https://zed.dev".to_string());
50 pub static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE")
51 .ok()
52 .and_then(|s| if s.is_empty() { None } else { Some(s) });
53}
54
55pub const ZED_SECRET_CLIENT_TOKEN: &str = "618033988749894";
56pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(100);
57pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(5);
58
59actions!(client, [Authenticate]);
60
61pub fn init(client: Arc<Client>, cx: &mut MutableAppContext) {
62 cx.add_global_action({
63 let client = client.clone();
64 move |_: &Authenticate, cx| {
65 let client = client.clone();
66 cx.spawn(
67 |cx| async move { client.authenticate_and_connect(true, &cx).log_err().await },
68 )
69 .detach();
70 }
71 });
72}
73
74pub struct Client {
75 id: usize,
76 peer: Arc<Peer>,
77 http: Arc<dyn HttpClient>,
78 telemetry: Arc<Telemetry>,
79 state: RwLock<ClientState>,
80
81 #[allow(clippy::type_complexity)]
82 #[cfg(any(test, feature = "test-support"))]
83 authenticate: RwLock<
84 Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,
85 >,
86
87 #[allow(clippy::type_complexity)]
88 #[cfg(any(test, feature = "test-support"))]
89 establish_connection: RwLock<
90 Option<
91 Box<
92 dyn 'static
93 + Send
94 + Sync
95 + Fn(
96 &Credentials,
97 &AsyncAppContext,
98 ) -> Task<Result<Connection, EstablishConnectionError>>,
99 >,
100 >,
101 >,
102}
103
104#[derive(Error, Debug)]
105pub enum EstablishConnectionError {
106 #[error("upgrade required")]
107 UpgradeRequired,
108 #[error("unauthorized")]
109 Unauthorized,
110 #[error("{0}")]
111 Other(#[from] anyhow::Error),
112 #[error("{0}")]
113 Http(#[from] http::Error),
114 #[error("{0}")]
115 Io(#[from] std::io::Error),
116 #[error("{0}")]
117 Websocket(#[from] async_tungstenite::tungstenite::http::Error),
118}
119
120impl From<WebsocketError> for EstablishConnectionError {
121 fn from(error: WebsocketError) -> Self {
122 if let WebsocketError::Http(response) = &error {
123 match response.status() {
124 StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,
125 StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,
126 _ => {}
127 }
128 }
129 EstablishConnectionError::Other(error.into())
130 }
131}
132
133impl EstablishConnectionError {
134 pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
135 Self::Other(error.into())
136 }
137}
138
139#[derive(Copy, Clone, Debug, Eq, PartialEq)]
140pub enum Status {
141 SignedOut,
142 UpgradeRequired,
143 Authenticating,
144 Connecting,
145 ConnectionError,
146 Connected { connection_id: ConnectionId },
147 ConnectionLost,
148 Reauthenticating,
149 Reconnecting,
150 ReconnectionError { next_reconnection: Instant },
151}
152
153impl Status {
154 pub fn is_connected(&self) -> bool {
155 matches!(self, Self::Connected { .. })
156 }
157}
158
159struct ClientState {
160 credentials: Option<Credentials>,
161 status: (watch::Sender<Status>, watch::Receiver<Status>),
162 entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
163 _reconnect_task: Option<Task<()>>,
164 reconnect_interval: Duration,
165 entities_by_type_and_remote_id: HashMap<(TypeId, u64), AnyWeakEntityHandle>,
166 models_by_message_type: HashMap<TypeId, AnyWeakModelHandle>,
167 entity_types_by_message_type: HashMap<TypeId, TypeId>,
168 #[allow(clippy::type_complexity)]
169 message_handlers: HashMap<
170 TypeId,
171 Arc<
172 dyn Send
173 + Sync
174 + Fn(
175 AnyEntityHandle,
176 Box<dyn AnyTypedEnvelope>,
177 &Arc<Client>,
178 AsyncAppContext,
179 ) -> LocalBoxFuture<'static, Result<()>>,
180 >,
181 >,
182}
183
184enum AnyWeakEntityHandle {
185 Model(AnyWeakModelHandle),
186 View(AnyWeakViewHandle),
187}
188
189enum AnyEntityHandle {
190 Model(AnyModelHandle),
191 View(AnyViewHandle),
192}
193
194#[derive(Clone, Debug)]
195pub struct Credentials {
196 pub user_id: u64,
197 pub access_token: String,
198}
199
200impl Default for ClientState {
201 fn default() -> Self {
202 Self {
203 credentials: None,
204 status: watch::channel_with(Status::SignedOut),
205 entity_id_extractors: Default::default(),
206 _reconnect_task: None,
207 reconnect_interval: Duration::from_secs(5),
208 models_by_message_type: Default::default(),
209 entities_by_type_and_remote_id: Default::default(),
210 entity_types_by_message_type: Default::default(),
211 message_handlers: Default::default(),
212 }
213 }
214}
215
216pub enum Subscription {
217 Entity {
218 client: Weak<Client>,
219 id: (TypeId, u64),
220 },
221 Message {
222 client: Weak<Client>,
223 id: TypeId,
224 },
225}
226
227impl Drop for Subscription {
228 fn drop(&mut self) {
229 match self {
230 Subscription::Entity { client, id } => {
231 if let Some(client) = client.upgrade() {
232 let mut state = client.state.write();
233 let _ = state.entities_by_type_and_remote_id.remove(id);
234 }
235 }
236 Subscription::Message { client, id } => {
237 if let Some(client) = client.upgrade() {
238 let mut state = client.state.write();
239 let _ = state.entity_types_by_message_type.remove(id);
240 let _ = state.message_handlers.remove(id);
241 }
242 }
243 }
244 }
245}
246
247impl Client {
248 pub fn new(http: Arc<dyn HttpClient>, cx: &AppContext) -> Arc<Self> {
249 Arc::new(Self {
250 id: 0,
251 peer: Peer::new(),
252 telemetry: Telemetry::new(http.clone(), cx),
253 http,
254 state: Default::default(),
255
256 #[cfg(any(test, feature = "test-support"))]
257 authenticate: Default::default(),
258 #[cfg(any(test, feature = "test-support"))]
259 establish_connection: Default::default(),
260 })
261 }
262
263 pub fn id(&self) -> usize {
264 self.id
265 }
266
267 pub fn http_client(&self) -> Arc<dyn HttpClient> {
268 self.http.clone()
269 }
270
271 #[cfg(any(test, feature = "test-support"))]
272 pub fn set_id(&mut self, id: usize) -> &Self {
273 self.id = id;
274 self
275 }
276
277 #[cfg(any(test, feature = "test-support"))]
278 pub fn tear_down(&self) {
279 let mut state = self.state.write();
280 state._reconnect_task.take();
281 state.message_handlers.clear();
282 state.models_by_message_type.clear();
283 state.entities_by_type_and_remote_id.clear();
284 state.entity_id_extractors.clear();
285 self.peer.reset();
286 }
287
288 #[cfg(any(test, feature = "test-support"))]
289 pub fn override_authenticate<F>(&self, authenticate: F) -> &Self
290 where
291 F: 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>,
292 {
293 *self.authenticate.write() = Some(Box::new(authenticate));
294 self
295 }
296
297 #[cfg(any(test, feature = "test-support"))]
298 pub fn override_establish_connection<F>(&self, connect: F) -> &Self
299 where
300 F: 'static
301 + Send
302 + Sync
303 + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,
304 {
305 *self.establish_connection.write() = Some(Box::new(connect));
306 self
307 }
308
309 pub fn user_id(&self) -> Option<u64> {
310 self.state
311 .read()
312 .credentials
313 .as_ref()
314 .map(|credentials| credentials.user_id)
315 }
316
317 pub fn status(&self) -> watch::Receiver<Status> {
318 self.state.read().status.1.clone()
319 }
320
321 fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncAppContext) {
322 log::info!("set status on client {}: {:?}", self.id, status);
323 let mut state = self.state.write();
324 *state.status.0.borrow_mut() = status;
325
326 match status {
327 Status::Connected { .. } => {
328 state._reconnect_task = None;
329 }
330 Status::ConnectionLost => {
331 let this = self.clone();
332 let reconnect_interval = state.reconnect_interval;
333 state._reconnect_task = Some(cx.spawn(|cx| async move {
334 let mut rng = StdRng::from_entropy();
335 let mut delay = INITIAL_RECONNECTION_DELAY;
336 while let Err(error) = this.authenticate_and_connect(true, &cx).await {
337 log::error!("failed to connect {}", error);
338 if matches!(*this.status().borrow(), Status::ConnectionError) {
339 this.set_status(
340 Status::ReconnectionError {
341 next_reconnection: Instant::now() + delay,
342 },
343 &cx,
344 );
345 cx.background().timer(delay).await;
346 delay = delay
347 .mul_f32(rng.gen_range(1.0..=2.0))
348 .min(reconnect_interval);
349 } else {
350 break;
351 }
352 }
353 }));
354 }
355 Status::SignedOut | Status::UpgradeRequired => {
356 self.telemetry.set_authenticated_user_info(None, false);
357 state._reconnect_task.take();
358 }
359 _ => {}
360 }
361 }
362
363 pub fn add_view_for_remote_entity<T: View>(
364 self: &Arc<Self>,
365 remote_id: u64,
366 cx: &mut ViewContext<T>,
367 ) -> Subscription {
368 let id = (TypeId::of::<T>(), remote_id);
369 self.state
370 .write()
371 .entities_by_type_and_remote_id
372 .insert(id, AnyWeakEntityHandle::View(cx.weak_handle().into()));
373 Subscription::Entity {
374 client: Arc::downgrade(self),
375 id,
376 }
377 }
378
379 pub fn add_model_for_remote_entity<T: Entity>(
380 self: &Arc<Self>,
381 remote_id: u64,
382 cx: &mut ModelContext<T>,
383 ) -> Subscription {
384 let id = (TypeId::of::<T>(), remote_id);
385 self.state
386 .write()
387 .entities_by_type_and_remote_id
388 .insert(id, AnyWeakEntityHandle::Model(cx.weak_handle().into()));
389 Subscription::Entity {
390 client: Arc::downgrade(self),
391 id,
392 }
393 }
394
395 pub fn add_message_handler<M, E, H, F>(
396 self: &Arc<Self>,
397 model: ModelHandle<E>,
398 handler: H,
399 ) -> Subscription
400 where
401 M: EnvelopedMessage,
402 E: Entity,
403 H: 'static
404 + Send
405 + Sync
406 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
407 F: 'static + Future<Output = Result<()>>,
408 {
409 let message_type_id = TypeId::of::<M>();
410
411 let mut state = self.state.write();
412 state
413 .models_by_message_type
414 .insert(message_type_id, model.downgrade().into());
415
416 let prev_handler = state.message_handlers.insert(
417 message_type_id,
418 Arc::new(move |handle, envelope, client, cx| {
419 let handle = if let AnyEntityHandle::Model(handle) = handle {
420 handle
421 } else {
422 unreachable!();
423 };
424 let model = handle.downcast::<E>().unwrap();
425 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
426 handler(model, *envelope, client.clone(), cx).boxed_local()
427 }),
428 );
429 if prev_handler.is_some() {
430 panic!("registered handler for the same message twice");
431 }
432
433 Subscription::Message {
434 client: Arc::downgrade(self),
435 id: message_type_id,
436 }
437 }
438
439 pub fn add_request_handler<M, E, H, F>(
440 self: &Arc<Self>,
441 model: ModelHandle<E>,
442 handler: H,
443 ) -> Subscription
444 where
445 M: RequestMessage,
446 E: Entity,
447 H: 'static
448 + Send
449 + Sync
450 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
451 F: 'static + Future<Output = Result<M::Response>>,
452 {
453 self.add_message_handler(model, move |handle, envelope, this, cx| {
454 Self::respond_to_request(
455 envelope.receipt(),
456 handler(handle, envelope, this.clone(), cx),
457 this,
458 )
459 })
460 }
461
462 pub fn add_view_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
463 where
464 M: EntityMessage,
465 E: View,
466 H: 'static
467 + Send
468 + Sync
469 + Fn(ViewHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
470 F: 'static + Future<Output = Result<()>>,
471 {
472 self.add_entity_message_handler::<M, E, _, _>(move |handle, message, client, cx| {
473 if let AnyEntityHandle::View(handle) = handle {
474 handler(handle.downcast::<E>().unwrap(), message, client, cx)
475 } else {
476 unreachable!();
477 }
478 })
479 }
480
481 pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
482 where
483 M: EntityMessage,
484 E: Entity,
485 H: 'static
486 + Send
487 + Sync
488 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
489 F: 'static + Future<Output = Result<()>>,
490 {
491 self.add_entity_message_handler::<M, E, _, _>(move |handle, message, client, cx| {
492 if let AnyEntityHandle::Model(handle) = handle {
493 handler(handle.downcast::<E>().unwrap(), message, client, cx)
494 } else {
495 unreachable!();
496 }
497 })
498 }
499
500 fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
501 where
502 M: EntityMessage,
503 E: Entity,
504 H: 'static
505 + Send
506 + Sync
507 + Fn(AnyEntityHandle, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
508 F: 'static + Future<Output = Result<()>>,
509 {
510 let model_type_id = TypeId::of::<E>();
511 let message_type_id = TypeId::of::<M>();
512
513 let mut state = self.state.write();
514 state
515 .entity_types_by_message_type
516 .insert(message_type_id, model_type_id);
517 state
518 .entity_id_extractors
519 .entry(message_type_id)
520 .or_insert_with(|| {
521 |envelope| {
522 envelope
523 .as_any()
524 .downcast_ref::<TypedEnvelope<M>>()
525 .unwrap()
526 .payload
527 .remote_entity_id()
528 }
529 });
530 let prev_handler = state.message_handlers.insert(
531 message_type_id,
532 Arc::new(move |handle, envelope, client, cx| {
533 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
534 handler(handle, *envelope, client.clone(), cx).boxed_local()
535 }),
536 );
537 if prev_handler.is_some() {
538 panic!("registered handler for the same message twice");
539 }
540 }
541
542 pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
543 where
544 M: EntityMessage + RequestMessage,
545 E: Entity,
546 H: 'static
547 + Send
548 + Sync
549 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
550 F: 'static + Future<Output = Result<M::Response>>,
551 {
552 self.add_model_message_handler(move |entity, envelope, client, cx| {
553 Self::respond_to_request::<M, _>(
554 envelope.receipt(),
555 handler(entity, envelope, client.clone(), cx),
556 client,
557 )
558 })
559 }
560
561 pub fn add_view_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
562 where
563 M: EntityMessage + RequestMessage,
564 E: View,
565 H: 'static
566 + Send
567 + Sync
568 + Fn(ViewHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
569 F: 'static + Future<Output = Result<M::Response>>,
570 {
571 self.add_view_message_handler(move |entity, envelope, client, cx| {
572 Self::respond_to_request::<M, _>(
573 envelope.receipt(),
574 handler(entity, envelope, client.clone(), cx),
575 client,
576 )
577 })
578 }
579
580 async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
581 receipt: Receipt<T>,
582 response: F,
583 client: Arc<Self>,
584 ) -> Result<()> {
585 match response.await {
586 Ok(response) => {
587 client.respond(receipt, response)?;
588 Ok(())
589 }
590 Err(error) => {
591 client.respond_with_error(
592 receipt,
593 proto::Error {
594 message: format!("{:?}", error),
595 },
596 )?;
597 Err(error)
598 }
599 }
600 }
601
602 pub fn has_keychain_credentials(&self, cx: &AsyncAppContext) -> bool {
603 read_credentials_from_keychain(cx).is_some()
604 }
605
606 #[async_recursion(?Send)]
607 pub async fn authenticate_and_connect(
608 self: &Arc<Self>,
609 try_keychain: bool,
610 cx: &AsyncAppContext,
611 ) -> anyhow::Result<()> {
612 let was_disconnected = match *self.status().borrow() {
613 Status::SignedOut => true,
614 Status::ConnectionError
615 | Status::ConnectionLost
616 | Status::Authenticating { .. }
617 | Status::Reauthenticating { .. }
618 | Status::ReconnectionError { .. } => false,
619 Status::Connected { .. } | Status::Connecting { .. } | Status::Reconnecting { .. } => {
620 return Ok(())
621 }
622 Status::UpgradeRequired => return Err(EstablishConnectionError::UpgradeRequired)?,
623 };
624
625 if was_disconnected {
626 self.set_status(Status::Authenticating, cx);
627 } else {
628 self.set_status(Status::Reauthenticating, cx)
629 }
630
631 let mut read_from_keychain = false;
632 let mut credentials = self.state.read().credentials.clone();
633 if credentials.is_none() && try_keychain {
634 credentials = read_credentials_from_keychain(cx);
635 read_from_keychain = credentials.is_some();
636 if read_from_keychain {
637 self.report_event("read credentials from keychain", Default::default());
638 }
639 }
640 if credentials.is_none() {
641 let mut status_rx = self.status();
642 let _ = status_rx.next().await;
643 futures::select_biased! {
644 authenticate = self.authenticate(cx).fuse() => {
645 match authenticate {
646 Ok(creds) => credentials = Some(creds),
647 Err(err) => {
648 self.set_status(Status::ConnectionError, cx);
649 return Err(err);
650 }
651 }
652 }
653 _ = status_rx.next().fuse() => {
654 return Err(anyhow!("authentication canceled"));
655 }
656 }
657 }
658 let credentials = credentials.unwrap();
659
660 if was_disconnected {
661 self.set_status(Status::Connecting, cx);
662 } else {
663 self.set_status(Status::Reconnecting, cx);
664 }
665
666 let mut timeout = cx.background().timer(CONNECTION_TIMEOUT).fuse();
667 futures::select_biased! {
668 connection = self.establish_connection(&credentials, cx).fuse() => {
669 match connection {
670 Ok(conn) => {
671 self.state.write().credentials = Some(credentials.clone());
672 if !read_from_keychain && IMPERSONATE_LOGIN.is_none() {
673 write_credentials_to_keychain(&credentials, cx).log_err();
674 }
675 self.set_connection(conn, cx).await;
676 Ok(())
677 }
678 Err(EstablishConnectionError::Unauthorized) => {
679 self.state.write().credentials.take();
680 if read_from_keychain {
681 cx.platform().delete_credentials(&ZED_SERVER_URL).log_err();
682 self.set_status(Status::SignedOut, cx);
683 self.authenticate_and_connect(false, cx).await
684 } else {
685 self.set_status(Status::ConnectionError, cx);
686 Err(EstablishConnectionError::Unauthorized)?
687 }
688 }
689 Err(EstablishConnectionError::UpgradeRequired) => {
690 self.set_status(Status::UpgradeRequired, cx);
691 Err(EstablishConnectionError::UpgradeRequired)?
692 }
693 Err(error) => {
694 self.set_status(Status::ConnectionError, cx);
695 Err(error)?
696 }
697 }
698 }
699 _ = timeout => {
700 self.set_status(Status::ConnectionError, cx);
701 Err(anyhow!("timed out trying to establish connection"))
702 }
703 }
704 }
705
706 async fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncAppContext) {
707 let executor = cx.background();
708 log::info!("add connection to peer");
709 let (connection_id, handle_io, mut incoming) = self
710 .peer
711 .add_connection(conn, move |duration| executor.timer(duration))
712 .await;
713 log::info!("set status to connected {}", connection_id);
714 self.set_status(Status::Connected { connection_id }, cx);
715 cx.foreground()
716 .spawn({
717 let cx = cx.clone();
718 let this = self.clone();
719 async move {
720 let mut message_id = 0_usize;
721 while let Some(message) = incoming.next().await {
722 let mut state = this.state.write();
723 message_id += 1;
724 let type_name = message.payload_type_name();
725 let payload_type_id = message.payload_type_id();
726 let sender_id = message.original_sender_id().map(|id| id.0);
727
728 let model = state
729 .models_by_message_type
730 .get(&payload_type_id)
731 .and_then(|model| model.upgrade(&cx))
732 .map(AnyEntityHandle::Model)
733 .or_else(|| {
734 let entity_type_id =
735 *state.entity_types_by_message_type.get(&payload_type_id)?;
736 let entity_id = state
737 .entity_id_extractors
738 .get(&message.payload_type_id())
739 .map(|extract_entity_id| {
740 (extract_entity_id)(message.as_ref())
741 })?;
742
743 let entity = state
744 .entities_by_type_and_remote_id
745 .get(&(entity_type_id, entity_id))?;
746 if let Some(entity) = entity.upgrade(&cx) {
747 Some(entity)
748 } else {
749 state
750 .entities_by_type_and_remote_id
751 .remove(&(entity_type_id, entity_id));
752 None
753 }
754 });
755
756 let model = if let Some(model) = model {
757 model
758 } else {
759 log::info!("unhandled message {}", type_name);
760 continue;
761 };
762
763 if let Some(handler) = state.message_handlers.get(&payload_type_id).cloned()
764 {
765 drop(state); // Avoid deadlocks if the handler interacts with rpc::Client
766 let future = handler(model, message, &this, cx.clone());
767
768 let client_id = this.id;
769 log::debug!(
770 "rpc message received. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
771 client_id,
772 message_id,
773 sender_id,
774 type_name
775 );
776 cx.foreground()
777 .spawn(async move {
778 match future.await {
779 Ok(()) => {
780 log::debug!(
781 "rpc message handled. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
782 client_id,
783 message_id,
784 sender_id,
785 type_name
786 );
787 }
788 Err(error) => {
789 log::error!(
790 "error handling message. client_id:{}, message_id:{}, sender_id:{:?}, type:{}, error:{:?}",
791 client_id,
792 message_id,
793 sender_id,
794 type_name,
795 error
796 );
797 }
798 }
799 })
800 .detach();
801 } else {
802 log::info!("unhandled message {}", type_name);
803 }
804
805 // Don't starve the main thread when receiving lots of messages at once.
806 smol::future::yield_now().await;
807 }
808 }
809 })
810 .detach();
811
812 let handle_io = cx.background().spawn(handle_io);
813 let this = self.clone();
814 let cx = cx.clone();
815 cx.foreground()
816 .spawn(async move {
817 match handle_io.await {
818 Ok(()) => {
819 if *this.status().borrow() == (Status::Connected { connection_id }) {
820 this.set_status(Status::SignedOut, &cx);
821 }
822 }
823 Err(err) => {
824 log::error!("connection error: {:?}", err);
825 this.set_status(Status::ConnectionLost, &cx);
826 }
827 }
828 })
829 .detach();
830 }
831
832 fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {
833 #[cfg(any(test, feature = "test-support"))]
834 if let Some(callback) = self.authenticate.read().as_ref() {
835 return callback(cx);
836 }
837
838 self.authenticate_with_browser(cx)
839 }
840
841 fn establish_connection(
842 self: &Arc<Self>,
843 credentials: &Credentials,
844 cx: &AsyncAppContext,
845 ) -> Task<Result<Connection, EstablishConnectionError>> {
846 #[cfg(any(test, feature = "test-support"))]
847 if let Some(callback) = self.establish_connection.read().as_ref() {
848 return callback(credentials, cx);
849 }
850
851 self.establish_websocket_connection(credentials, cx)
852 }
853
854 fn establish_websocket_connection(
855 self: &Arc<Self>,
856 credentials: &Credentials,
857 cx: &AsyncAppContext,
858 ) -> Task<Result<Connection, EstablishConnectionError>> {
859 let request = Request::builder()
860 .header(
861 "Authorization",
862 format!("{} {}", credentials.user_id, credentials.access_token),
863 )
864 .header("x-zed-protocol-version", rpc::PROTOCOL_VERSION);
865
866 let http = self.http.clone();
867 cx.background().spawn(async move {
868 let mut rpc_url = format!("{}/rpc", *ZED_SERVER_URL);
869 let rpc_response = http.get(&rpc_url, Default::default(), false).await?;
870 if rpc_response.status().is_redirection() {
871 rpc_url = rpc_response
872 .headers()
873 .get("Location")
874 .ok_or_else(|| anyhow!("missing location header in /rpc response"))?
875 .to_str()
876 .map_err(EstablishConnectionError::other)?
877 .to_string();
878 }
879 // Until we switch the zed.dev domain to point to the new Next.js app, there
880 // will be no redirect required, and the app will connect directly to
881 // wss://zed.dev/rpc.
882 else if rpc_response.status() != StatusCode::UPGRADE_REQUIRED {
883 Err(anyhow!(
884 "unexpected /rpc response status {}",
885 rpc_response.status()
886 ))?
887 }
888
889 let mut rpc_url = Url::parse(&rpc_url).context("invalid rpc url")?;
890 let rpc_host = rpc_url
891 .host_str()
892 .zip(rpc_url.port_or_known_default())
893 .ok_or_else(|| anyhow!("missing host in rpc url"))?;
894 let stream = smol::net::TcpStream::connect(rpc_host).await?;
895
896 log::info!("connected to rpc endpoint {}", rpc_url);
897
898 match rpc_url.scheme() {
899 "https" => {
900 rpc_url.set_scheme("wss").unwrap();
901 let request = request.uri(rpc_url.as_str()).body(())?;
902 let (stream, _) =
903 async_tungstenite::async_tls::client_async_tls(request, stream).await?;
904 Ok(Connection::new(
905 stream
906 .map_err(|error| anyhow!(error))
907 .sink_map_err(|error| anyhow!(error)),
908 ))
909 }
910 "http" => {
911 rpc_url.set_scheme("ws").unwrap();
912 let request = request.uri(rpc_url.as_str()).body(())?;
913 let (stream, _) = async_tungstenite::client_async(request, stream).await?;
914 Ok(Connection::new(
915 stream
916 .map_err(|error| anyhow!(error))
917 .sink_map_err(|error| anyhow!(error)),
918 ))
919 }
920 _ => Err(anyhow!("invalid rpc url: {}", rpc_url))?,
921 }
922 })
923 }
924
925 pub fn authenticate_with_browser(
926 self: &Arc<Self>,
927 cx: &AsyncAppContext,
928 ) -> Task<Result<Credentials>> {
929 let platform = cx.platform();
930 let executor = cx.background();
931 let telemetry = self.telemetry.clone();
932 executor.clone().spawn(async move {
933 // Generate a pair of asymmetric encryption keys. The public key will be used by the
934 // zed server to encrypt the user's access token, so that it can'be intercepted by
935 // any other app running on the user's device.
936 let (public_key, private_key) =
937 rpc::auth::keypair().expect("failed to generate keypair for auth");
938 let public_key_string =
939 String::try_from(public_key).expect("failed to serialize public key for auth");
940
941 // Start an HTTP server to receive the redirect from Zed's sign-in page.
942 let server = tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
943 let port = server.server_addr().port();
944
945 // Open the Zed sign-in page in the user's browser, with query parameters that indicate
946 // that the user is signing in from a Zed app running on the same device.
947 let mut url = format!(
948 "{}/native_app_signin?native_app_port={}&native_app_public_key={}",
949 *ZED_SERVER_URL, port, public_key_string
950 );
951
952 if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
953 log::info!("impersonating user @{}", impersonate_login);
954 write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
955 }
956
957 platform.open_url(&url);
958
959 // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
960 // access token from the query params.
961 //
962 // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
963 // custom URL scheme instead of this local HTTP server.
964 let (user_id, access_token) = executor
965 .spawn(async move {
966 for _ in 0..100 {
967 if let Some(req) = server.recv_timeout(Duration::from_secs(1))? {
968 let path = req.url();
969 let mut user_id = None;
970 let mut access_token = None;
971 let url = Url::parse(&format!("http://example.com{}", path))
972 .context("failed to parse login notification url")?;
973 for (key, value) in url.query_pairs() {
974 if key == "access_token" {
975 access_token = Some(value.to_string());
976 } else if key == "user_id" {
977 user_id = Some(value.to_string());
978 }
979 }
980
981 let post_auth_url =
982 format!("{}/native_app_signin_succeeded", *ZED_SERVER_URL);
983 req.respond(
984 tiny_http::Response::empty(302).with_header(
985 tiny_http::Header::from_bytes(
986 &b"Location"[..],
987 post_auth_url.as_bytes(),
988 )
989 .unwrap(),
990 ),
991 )
992 .context("failed to respond to login http request")?;
993 return Ok((
994 user_id.ok_or_else(|| anyhow!("missing user_id parameter"))?,
995 access_token
996 .ok_or_else(|| anyhow!("missing access_token parameter"))?,
997 ));
998 }
999 }
1000
1001 Err(anyhow!("didn't receive login redirect"))
1002 })
1003 .await?;
1004
1005 let access_token = private_key
1006 .decrypt_string(&access_token)
1007 .context("failed to decrypt access token")?;
1008 platform.activate(true);
1009
1010 telemetry.report_event("authenticate with browser", Default::default());
1011
1012 Ok(Credentials {
1013 user_id: user_id.parse()?,
1014 access_token,
1015 })
1016 })
1017 }
1018
1019 pub fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) -> Result<()> {
1020 let conn_id = self.connection_id()?;
1021 self.peer.disconnect(conn_id);
1022 self.set_status(Status::SignedOut, cx);
1023 Ok(())
1024 }
1025
1026 fn connection_id(&self) -> Result<ConnectionId> {
1027 if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1028 Ok(connection_id)
1029 } else {
1030 Err(anyhow!("not connected"))
1031 }
1032 }
1033
1034 pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1035 log::debug!("rpc send. client_id:{}, name:{}", self.id, T::NAME);
1036 self.peer.send(self.connection_id()?, message)
1037 }
1038
1039 pub fn request<T: RequestMessage>(
1040 &self,
1041 request: T,
1042 ) -> impl Future<Output = Result<T::Response>> {
1043 let client_id = self.id;
1044 log::debug!(
1045 "rpc request start. client_id:{}. name:{}",
1046 client_id,
1047 T::NAME
1048 );
1049 let response = self
1050 .connection_id()
1051 .map(|conn_id| self.peer.request(conn_id, request));
1052 async move {
1053 let response = response?.await;
1054 log::debug!(
1055 "rpc request finish. client_id:{}. name:{}",
1056 client_id,
1057 T::NAME
1058 );
1059 response
1060 }
1061 }
1062
1063 fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
1064 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1065 self.peer.respond(receipt, response)
1066 }
1067
1068 fn respond_with_error<T: RequestMessage>(
1069 &self,
1070 receipt: Receipt<T>,
1071 error: proto::Error,
1072 ) -> Result<()> {
1073 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1074 self.peer.respond_with_error(receipt, error)
1075 }
1076
1077 pub fn start_telemetry(&self, db: Arc<Db>) {
1078 self.telemetry.start(db);
1079 }
1080
1081 pub fn report_event(&self, kind: &str, properties: Value) {
1082 self.telemetry.report_event(kind, properties)
1083 }
1084
1085 pub fn telemetry_log_file_path(&self) -> Option<PathBuf> {
1086 self.telemetry.log_file_path()
1087 }
1088}
1089
1090impl AnyWeakEntityHandle {
1091 fn upgrade(&self, cx: &AsyncAppContext) -> Option<AnyEntityHandle> {
1092 match self {
1093 AnyWeakEntityHandle::Model(handle) => handle.upgrade(cx).map(AnyEntityHandle::Model),
1094 AnyWeakEntityHandle::View(handle) => handle.upgrade(cx).map(AnyEntityHandle::View),
1095 }
1096 }
1097}
1098
1099fn read_credentials_from_keychain(cx: &AsyncAppContext) -> Option<Credentials> {
1100 if IMPERSONATE_LOGIN.is_some() {
1101 return None;
1102 }
1103
1104 let (user_id, access_token) = cx
1105 .platform()
1106 .read_credentials(&ZED_SERVER_URL)
1107 .log_err()
1108 .flatten()?;
1109 Some(Credentials {
1110 user_id: user_id.parse().ok()?,
1111 access_token: String::from_utf8(access_token).ok()?,
1112 })
1113}
1114
1115fn write_credentials_to_keychain(credentials: &Credentials, cx: &AsyncAppContext) -> Result<()> {
1116 cx.platform().write_credentials(
1117 &ZED_SERVER_URL,
1118 &credentials.user_id.to_string(),
1119 credentials.access_token.as_bytes(),
1120 )
1121}
1122
1123const WORKTREE_URL_PREFIX: &str = "zed://worktrees/";
1124
1125pub fn encode_worktree_url(id: u64, access_token: &str) -> String {
1126 format!("{}{}/{}", WORKTREE_URL_PREFIX, id, access_token)
1127}
1128
1129pub fn decode_worktree_url(url: &str) -> Option<(u64, String)> {
1130 let path = url.trim().strip_prefix(WORKTREE_URL_PREFIX)?;
1131 let mut parts = path.split('/');
1132 let id = parts.next()?.parse::<u64>().ok()?;
1133 let access_token = parts.next()?;
1134 if access_token.is_empty() {
1135 return None;
1136 }
1137 Some((id, access_token.to_string()))
1138}
1139
1140#[cfg(test)]
1141mod tests {
1142 use super::*;
1143 use crate::test::{FakeHttpClient, FakeServer};
1144 use gpui::{executor::Deterministic, TestAppContext};
1145 use parking_lot::Mutex;
1146 use std::future;
1147
1148 #[gpui::test(iterations = 10)]
1149 async fn test_reconnection(cx: &mut TestAppContext) {
1150 cx.foreground().forbid_parking();
1151
1152 let user_id = 5;
1153 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1154 let server = FakeServer::for_client(user_id, &client, cx).await;
1155 let mut status = client.status();
1156 assert!(matches!(
1157 status.next().await,
1158 Some(Status::Connected { .. })
1159 ));
1160 assert_eq!(server.auth_count(), 1);
1161
1162 server.forbid_connections();
1163 server.disconnect();
1164 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1165
1166 server.allow_connections();
1167 cx.foreground().advance_clock(Duration::from_secs(10));
1168 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1169 assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1170
1171 server.forbid_connections();
1172 server.disconnect();
1173 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1174
1175 // Clear cached credentials after authentication fails
1176 server.roll_access_token();
1177 server.allow_connections();
1178 cx.foreground().advance_clock(Duration::from_secs(10));
1179 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1180 assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1181 }
1182
1183 #[gpui::test(iterations = 10)]
1184 async fn test_connection_timeout(deterministic: Arc<Deterministic>, cx: &mut TestAppContext) {
1185 deterministic.forbid_parking();
1186
1187 let user_id = 5;
1188 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1189 let mut status = client.status();
1190
1191 // Time out when client tries to connect.
1192 client.override_authenticate(move |cx| {
1193 cx.foreground().spawn(async move {
1194 Ok(Credentials {
1195 user_id,
1196 access_token: "token".into(),
1197 })
1198 })
1199 });
1200 client.override_establish_connection(|_, cx| {
1201 cx.foreground().spawn(async move {
1202 future::pending::<()>().await;
1203 unreachable!()
1204 })
1205 });
1206 let auth_and_connect = cx.spawn({
1207 let client = client.clone();
1208 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1209 });
1210 deterministic.run_until_parked();
1211 assert!(matches!(status.next().await, Some(Status::Connecting)));
1212
1213 deterministic.advance_clock(CONNECTION_TIMEOUT);
1214 assert!(matches!(
1215 status.next().await,
1216 Some(Status::ConnectionError { .. })
1217 ));
1218 auth_and_connect.await.unwrap_err();
1219
1220 // Allow the connection to be established.
1221 let server = FakeServer::for_client(user_id, &client, cx).await;
1222 assert!(matches!(
1223 status.next().await,
1224 Some(Status::Connected { .. })
1225 ));
1226
1227 // Disconnect client.
1228 server.forbid_connections();
1229 server.disconnect();
1230 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1231
1232 // Time out when re-establishing the connection.
1233 server.allow_connections();
1234 client.override_establish_connection(|_, cx| {
1235 cx.foreground().spawn(async move {
1236 future::pending::<()>().await;
1237 unreachable!()
1238 })
1239 });
1240 deterministic.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1241 assert!(matches!(
1242 status.next().await,
1243 Some(Status::Reconnecting { .. })
1244 ));
1245
1246 deterministic.advance_clock(CONNECTION_TIMEOUT);
1247 assert!(matches!(
1248 status.next().await,
1249 Some(Status::ReconnectionError { .. })
1250 ));
1251 }
1252
1253 #[gpui::test(iterations = 10)]
1254 async fn test_authenticating_more_than_once(
1255 cx: &mut TestAppContext,
1256 deterministic: Arc<Deterministic>,
1257 ) {
1258 cx.foreground().forbid_parking();
1259
1260 let auth_count = Arc::new(Mutex::new(0));
1261 let dropped_auth_count = Arc::new(Mutex::new(0));
1262 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1263 client.override_authenticate({
1264 let auth_count = auth_count.clone();
1265 let dropped_auth_count = dropped_auth_count.clone();
1266 move |cx| {
1267 let auth_count = auth_count.clone();
1268 let dropped_auth_count = dropped_auth_count.clone();
1269 cx.foreground().spawn(async move {
1270 *auth_count.lock() += 1;
1271 let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
1272 future::pending::<()>().await;
1273 unreachable!()
1274 })
1275 }
1276 });
1277
1278 let _authenticate = cx.spawn(|cx| {
1279 let client = client.clone();
1280 async move { client.authenticate_and_connect(false, &cx).await }
1281 });
1282 deterministic.run_until_parked();
1283 assert_eq!(*auth_count.lock(), 1);
1284 assert_eq!(*dropped_auth_count.lock(), 0);
1285
1286 let _authenticate = cx.spawn(|cx| {
1287 let client = client.clone();
1288 async move { client.authenticate_and_connect(false, &cx).await }
1289 });
1290 deterministic.run_until_parked();
1291 assert_eq!(*auth_count.lock(), 2);
1292 assert_eq!(*dropped_auth_count.lock(), 1);
1293 }
1294
1295 #[test]
1296 fn test_encode_and_decode_worktree_url() {
1297 let url = encode_worktree_url(5, "deadbeef");
1298 assert_eq!(decode_worktree_url(&url), Some((5, "deadbeef".to_string())));
1299 assert_eq!(
1300 decode_worktree_url(&format!("\n {}\t", url)),
1301 Some((5, "deadbeef".to_string()))
1302 );
1303 assert_eq!(decode_worktree_url("not://the-right-format"), None);
1304 }
1305
1306 #[gpui::test]
1307 async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
1308 cx.foreground().forbid_parking();
1309
1310 let user_id = 5;
1311 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1312 let server = FakeServer::for_client(user_id, &client, cx).await;
1313
1314 let (done_tx1, mut done_rx1) = smol::channel::unbounded();
1315 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1316 client.add_model_message_handler(
1317 move |model: ModelHandle<Model>, _: TypedEnvelope<proto::JoinProject>, _, cx| {
1318 match model.read_with(&cx, |model, _| model.id) {
1319 1 => done_tx1.try_send(()).unwrap(),
1320 2 => done_tx2.try_send(()).unwrap(),
1321 _ => unreachable!(),
1322 }
1323 async { Ok(()) }
1324 },
1325 );
1326 let model1 = cx.add_model(|_| Model {
1327 id: 1,
1328 subscription: None,
1329 });
1330 let model2 = cx.add_model(|_| Model {
1331 id: 2,
1332 subscription: None,
1333 });
1334 let model3 = cx.add_model(|_| Model {
1335 id: 3,
1336 subscription: None,
1337 });
1338
1339 let _subscription1 = model1.update(cx, |_, cx| client.add_model_for_remote_entity(1, cx));
1340 let _subscription2 = model2.update(cx, |_, cx| client.add_model_for_remote_entity(2, cx));
1341 // Ensure dropping a subscription for the same entity type still allows receiving of
1342 // messages for other entity IDs of the same type.
1343 let subscription3 = model3.update(cx, |_, cx| client.add_model_for_remote_entity(3, cx));
1344 drop(subscription3);
1345
1346 server.send(proto::JoinProject { project_id: 1 });
1347 server.send(proto::JoinProject { project_id: 2 });
1348 done_rx1.next().await.unwrap();
1349 done_rx2.next().await.unwrap();
1350 }
1351
1352 #[gpui::test]
1353 async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
1354 cx.foreground().forbid_parking();
1355
1356 let user_id = 5;
1357 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1358 let server = FakeServer::for_client(user_id, &client, cx).await;
1359
1360 let model = cx.add_model(|_| Model::default());
1361 let (done_tx1, _done_rx1) = smol::channel::unbounded();
1362 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1363 let subscription1 = client.add_message_handler(
1364 model.clone(),
1365 move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1366 done_tx1.try_send(()).unwrap();
1367 async { Ok(()) }
1368 },
1369 );
1370 drop(subscription1);
1371 let _subscription2 =
1372 client.add_message_handler(model, move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1373 done_tx2.try_send(()).unwrap();
1374 async { Ok(()) }
1375 });
1376 server.send(proto::Ping {});
1377 done_rx2.next().await.unwrap();
1378 }
1379
1380 #[gpui::test]
1381 async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
1382 cx.foreground().forbid_parking();
1383
1384 let user_id = 5;
1385 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1386 let server = FakeServer::for_client(user_id, &client, cx).await;
1387
1388 let model = cx.add_model(|_| Model::default());
1389 let (done_tx, mut done_rx) = smol::channel::unbounded();
1390 let subscription = client.add_message_handler(
1391 model.clone(),
1392 move |model, _: TypedEnvelope<proto::Ping>, _, mut cx| {
1393 model.update(&mut cx, |model, _| model.subscription.take());
1394 done_tx.try_send(()).unwrap();
1395 async { Ok(()) }
1396 },
1397 );
1398 model.update(cx, |model, _| {
1399 model.subscription = Some(subscription);
1400 });
1401 server.send(proto::Ping {});
1402 done_rx.next().await.unwrap();
1403 }
1404
1405 #[derive(Default)]
1406 struct Model {
1407 id: usize,
1408 subscription: Option<Subscription>,
1409 }
1410
1411 impl Entity for Model {
1412 type Event = ();
1413 }
1414}