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 futures::select_biased! {
667 connection = self.establish_connection(&credentials, cx).fuse() => {
668 match connection {
669 Ok(conn) => {
670 self.state.write().credentials = Some(credentials.clone());
671 if !read_from_keychain && IMPERSONATE_LOGIN.is_none() {
672 write_credentials_to_keychain(&credentials, cx).log_err();
673 }
674 self.set_connection(conn, cx);
675 Ok(())
676 }
677 Err(EstablishConnectionError::Unauthorized) => {
678 self.state.write().credentials.take();
679 if read_from_keychain {
680 cx.platform().delete_credentials(&ZED_SERVER_URL).log_err();
681 self.set_status(Status::SignedOut, cx);
682 self.authenticate_and_connect(false, cx).await
683 } else {
684 self.set_status(Status::ConnectionError, cx);
685 Err(EstablishConnectionError::Unauthorized)?
686 }
687 }
688 Err(EstablishConnectionError::UpgradeRequired) => {
689 self.set_status(Status::UpgradeRequired, cx);
690 Err(EstablishConnectionError::UpgradeRequired)?
691 }
692 Err(error) => {
693 self.set_status(Status::ConnectionError, cx);
694 Err(error)?
695 }
696 }
697 }
698 _ = cx.background().timer(CONNECTION_TIMEOUT).fuse() => {
699 self.set_status(Status::ConnectionError, cx);
700 Err(anyhow!("timed out trying to establish connection"))
701 }
702 }
703 }
704
705 fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncAppContext) {
706 let executor = cx.background();
707 log::info!("add connection to peer");
708 let (connection_id, handle_io, mut incoming) = self
709 .peer
710 .add_connection(conn, move |duration| executor.timer(duration));
711 log::info!("set status to connected {}", connection_id);
712 self.set_status(Status::Connected { connection_id }, cx);
713 cx.foreground()
714 .spawn({
715 let cx = cx.clone();
716 let this = self.clone();
717 async move {
718 let mut message_id = 0_usize;
719 while let Some(message) = incoming.next().await {
720 let mut state = this.state.write();
721 message_id += 1;
722 let type_name = message.payload_type_name();
723 let payload_type_id = message.payload_type_id();
724 let sender_id = message.original_sender_id().map(|id| id.0);
725
726 let model = state
727 .models_by_message_type
728 .get(&payload_type_id)
729 .and_then(|model| model.upgrade(&cx))
730 .map(AnyEntityHandle::Model)
731 .or_else(|| {
732 let entity_type_id =
733 *state.entity_types_by_message_type.get(&payload_type_id)?;
734 let entity_id = state
735 .entity_id_extractors
736 .get(&message.payload_type_id())
737 .map(|extract_entity_id| {
738 (extract_entity_id)(message.as_ref())
739 })?;
740
741 let entity = state
742 .entities_by_type_and_remote_id
743 .get(&(entity_type_id, entity_id))?;
744 if let Some(entity) = entity.upgrade(&cx) {
745 Some(entity)
746 } else {
747 state
748 .entities_by_type_and_remote_id
749 .remove(&(entity_type_id, entity_id));
750 None
751 }
752 });
753
754 let model = if let Some(model) = model {
755 model
756 } else {
757 log::info!("unhandled message {}", type_name);
758 continue;
759 };
760
761 if let Some(handler) = state.message_handlers.get(&payload_type_id).cloned()
762 {
763 drop(state); // Avoid deadlocks if the handler interacts with rpc::Client
764 let future = handler(model, message, &this, cx.clone());
765
766 let client_id = this.id;
767 log::debug!(
768 "rpc message received. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
769 client_id,
770 message_id,
771 sender_id,
772 type_name
773 );
774 cx.foreground()
775 .spawn(async move {
776 match future.await {
777 Ok(()) => {
778 log::debug!(
779 "rpc message handled. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
780 client_id,
781 message_id,
782 sender_id,
783 type_name
784 );
785 }
786 Err(error) => {
787 log::error!(
788 "error handling message. client_id:{}, message_id:{}, sender_id:{:?}, type:{}, error:{:?}",
789 client_id,
790 message_id,
791 sender_id,
792 type_name,
793 error
794 );
795 }
796 }
797 })
798 .detach();
799 } else {
800 log::info!("unhandled message {}", type_name);
801 }
802
803 // Don't starve the main thread when receiving lots of messages at once.
804 smol::future::yield_now().await;
805 }
806 }
807 })
808 .detach();
809
810 let handle_io = cx.background().spawn(handle_io);
811 let this = self.clone();
812 let cx = cx.clone();
813 cx.foreground()
814 .spawn(async move {
815 match handle_io.await {
816 Ok(()) => {
817 if *this.status().borrow() == (Status::Connected { connection_id }) {
818 this.set_status(Status::SignedOut, &cx);
819 }
820 }
821 Err(err) => {
822 log::error!("connection error: {:?}", err);
823 this.set_status(Status::ConnectionLost, &cx);
824 }
825 }
826 })
827 .detach();
828 }
829
830 fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {
831 #[cfg(any(test, feature = "test-support"))]
832 if let Some(callback) = self.authenticate.read().as_ref() {
833 return callback(cx);
834 }
835
836 self.authenticate_with_browser(cx)
837 }
838
839 fn establish_connection(
840 self: &Arc<Self>,
841 credentials: &Credentials,
842 cx: &AsyncAppContext,
843 ) -> Task<Result<Connection, EstablishConnectionError>> {
844 #[cfg(any(test, feature = "test-support"))]
845 if let Some(callback) = self.establish_connection.read().as_ref() {
846 return callback(credentials, cx);
847 }
848
849 self.establish_websocket_connection(credentials, cx)
850 }
851
852 fn establish_websocket_connection(
853 self: &Arc<Self>,
854 credentials: &Credentials,
855 cx: &AsyncAppContext,
856 ) -> Task<Result<Connection, EstablishConnectionError>> {
857 let request = Request::builder()
858 .header(
859 "Authorization",
860 format!("{} {}", credentials.user_id, credentials.access_token),
861 )
862 .header("x-zed-protocol-version", rpc::PROTOCOL_VERSION);
863
864 let http = self.http.clone();
865 cx.background().spawn(async move {
866 let mut rpc_url = format!("{}/rpc", *ZED_SERVER_URL);
867 let rpc_response = http.get(&rpc_url, Default::default(), false).await?;
868 if rpc_response.status().is_redirection() {
869 rpc_url = rpc_response
870 .headers()
871 .get("Location")
872 .ok_or_else(|| anyhow!("missing location header in /rpc response"))?
873 .to_str()
874 .map_err(EstablishConnectionError::other)?
875 .to_string();
876 }
877 // Until we switch the zed.dev domain to point to the new Next.js app, there
878 // will be no redirect required, and the app will connect directly to
879 // wss://zed.dev/rpc.
880 else if rpc_response.status() != StatusCode::UPGRADE_REQUIRED {
881 Err(anyhow!(
882 "unexpected /rpc response status {}",
883 rpc_response.status()
884 ))?
885 }
886
887 let mut rpc_url = Url::parse(&rpc_url).context("invalid rpc url")?;
888 let rpc_host = rpc_url
889 .host_str()
890 .zip(rpc_url.port_or_known_default())
891 .ok_or_else(|| anyhow!("missing host in rpc url"))?;
892 let stream = smol::net::TcpStream::connect(rpc_host).await?;
893
894 log::info!("connected to rpc endpoint {}", rpc_url);
895
896 match rpc_url.scheme() {
897 "https" => {
898 rpc_url.set_scheme("wss").unwrap();
899 let request = request.uri(rpc_url.as_str()).body(())?;
900 let (stream, _) =
901 async_tungstenite::async_tls::client_async_tls(request, stream).await?;
902 Ok(Connection::new(
903 stream
904 .map_err(|error| anyhow!(error))
905 .sink_map_err(|error| anyhow!(error)),
906 ))
907 }
908 "http" => {
909 rpc_url.set_scheme("ws").unwrap();
910 let request = request.uri(rpc_url.as_str()).body(())?;
911 let (stream, _) = async_tungstenite::client_async(request, stream).await?;
912 Ok(Connection::new(
913 stream
914 .map_err(|error| anyhow!(error))
915 .sink_map_err(|error| anyhow!(error)),
916 ))
917 }
918 _ => Err(anyhow!("invalid rpc url: {}", rpc_url))?,
919 }
920 })
921 }
922
923 pub fn authenticate_with_browser(
924 self: &Arc<Self>,
925 cx: &AsyncAppContext,
926 ) -> Task<Result<Credentials>> {
927 let platform = cx.platform();
928 let executor = cx.background();
929 let telemetry = self.telemetry.clone();
930 executor.clone().spawn(async move {
931 // Generate a pair of asymmetric encryption keys. The public key will be used by the
932 // zed server to encrypt the user's access token, so that it can'be intercepted by
933 // any other app running on the user's device.
934 let (public_key, private_key) =
935 rpc::auth::keypair().expect("failed to generate keypair for auth");
936 let public_key_string =
937 String::try_from(public_key).expect("failed to serialize public key for auth");
938
939 // Start an HTTP server to receive the redirect from Zed's sign-in page.
940 let server = tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
941 let port = server.server_addr().port();
942
943 // Open the Zed sign-in page in the user's browser, with query parameters that indicate
944 // that the user is signing in from a Zed app running on the same device.
945 let mut url = format!(
946 "{}/native_app_signin?native_app_port={}&native_app_public_key={}",
947 *ZED_SERVER_URL, port, public_key_string
948 );
949
950 if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
951 log::info!("impersonating user @{}", impersonate_login);
952 write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
953 }
954
955 platform.open_url(&url);
956
957 // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
958 // access token from the query params.
959 //
960 // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
961 // custom URL scheme instead of this local HTTP server.
962 let (user_id, access_token) = executor
963 .spawn(async move {
964 for _ in 0..100 {
965 if let Some(req) = server.recv_timeout(Duration::from_secs(1))? {
966 let path = req.url();
967 let mut user_id = None;
968 let mut access_token = None;
969 let url = Url::parse(&format!("http://example.com{}", path))
970 .context("failed to parse login notification url")?;
971 for (key, value) in url.query_pairs() {
972 if key == "access_token" {
973 access_token = Some(value.to_string());
974 } else if key == "user_id" {
975 user_id = Some(value.to_string());
976 }
977 }
978
979 let post_auth_url =
980 format!("{}/native_app_signin_succeeded", *ZED_SERVER_URL);
981 req.respond(
982 tiny_http::Response::empty(302).with_header(
983 tiny_http::Header::from_bytes(
984 &b"Location"[..],
985 post_auth_url.as_bytes(),
986 )
987 .unwrap(),
988 ),
989 )
990 .context("failed to respond to login http request")?;
991 return Ok((
992 user_id.ok_or_else(|| anyhow!("missing user_id parameter"))?,
993 access_token
994 .ok_or_else(|| anyhow!("missing access_token parameter"))?,
995 ));
996 }
997 }
998
999 Err(anyhow!("didn't receive login redirect"))
1000 })
1001 .await?;
1002
1003 let access_token = private_key
1004 .decrypt_string(&access_token)
1005 .context("failed to decrypt access token")?;
1006 platform.activate(true);
1007
1008 telemetry.report_event("authenticate with browser", Default::default());
1009
1010 Ok(Credentials {
1011 user_id: user_id.parse()?,
1012 access_token,
1013 })
1014 })
1015 }
1016
1017 pub fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) -> Result<()> {
1018 let conn_id = self.connection_id()?;
1019 self.peer.disconnect(conn_id);
1020 self.set_status(Status::SignedOut, cx);
1021 Ok(())
1022 }
1023
1024 fn connection_id(&self) -> Result<ConnectionId> {
1025 if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1026 Ok(connection_id)
1027 } else {
1028 Err(anyhow!("not connected"))
1029 }
1030 }
1031
1032 pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1033 log::debug!("rpc send. client_id:{}, name:{}", self.id, T::NAME);
1034 self.peer.send(self.connection_id()?, message)
1035 }
1036
1037 pub fn request<T: RequestMessage>(
1038 &self,
1039 request: T,
1040 ) -> impl Future<Output = Result<T::Response>> {
1041 let client_id = self.id;
1042 log::debug!(
1043 "rpc request start. client_id:{}. name:{}",
1044 client_id,
1045 T::NAME
1046 );
1047 let response = self
1048 .connection_id()
1049 .map(|conn_id| self.peer.request(conn_id, request));
1050 async move {
1051 let response = response?.await;
1052 log::debug!(
1053 "rpc request finish. client_id:{}. name:{}",
1054 client_id,
1055 T::NAME
1056 );
1057 response
1058 }
1059 }
1060
1061 fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
1062 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1063 self.peer.respond(receipt, response)
1064 }
1065
1066 fn respond_with_error<T: RequestMessage>(
1067 &self,
1068 receipt: Receipt<T>,
1069 error: proto::Error,
1070 ) -> Result<()> {
1071 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1072 self.peer.respond_with_error(receipt, error)
1073 }
1074
1075 pub fn start_telemetry(&self, db: Arc<Db>) {
1076 self.telemetry.start(db);
1077 }
1078
1079 pub fn report_event(&self, kind: &str, properties: Value) {
1080 self.telemetry.report_event(kind, properties)
1081 }
1082
1083 pub fn telemetry_log_file_path(&self) -> Option<PathBuf> {
1084 self.telemetry.log_file_path()
1085 }
1086}
1087
1088impl AnyWeakEntityHandle {
1089 fn upgrade(&self, cx: &AsyncAppContext) -> Option<AnyEntityHandle> {
1090 match self {
1091 AnyWeakEntityHandle::Model(handle) => handle.upgrade(cx).map(AnyEntityHandle::Model),
1092 AnyWeakEntityHandle::View(handle) => handle.upgrade(cx).map(AnyEntityHandle::View),
1093 }
1094 }
1095}
1096
1097fn read_credentials_from_keychain(cx: &AsyncAppContext) -> Option<Credentials> {
1098 if IMPERSONATE_LOGIN.is_some() {
1099 return None;
1100 }
1101
1102 let (user_id, access_token) = cx
1103 .platform()
1104 .read_credentials(&ZED_SERVER_URL)
1105 .log_err()
1106 .flatten()?;
1107 Some(Credentials {
1108 user_id: user_id.parse().ok()?,
1109 access_token: String::from_utf8(access_token).ok()?,
1110 })
1111}
1112
1113fn write_credentials_to_keychain(credentials: &Credentials, cx: &AsyncAppContext) -> Result<()> {
1114 cx.platform().write_credentials(
1115 &ZED_SERVER_URL,
1116 &credentials.user_id.to_string(),
1117 credentials.access_token.as_bytes(),
1118 )
1119}
1120
1121const WORKTREE_URL_PREFIX: &str = "zed://worktrees/";
1122
1123pub fn encode_worktree_url(id: u64, access_token: &str) -> String {
1124 format!("{}{}/{}", WORKTREE_URL_PREFIX, id, access_token)
1125}
1126
1127pub fn decode_worktree_url(url: &str) -> Option<(u64, String)> {
1128 let path = url.trim().strip_prefix(WORKTREE_URL_PREFIX)?;
1129 let mut parts = path.split('/');
1130 let id = parts.next()?.parse::<u64>().ok()?;
1131 let access_token = parts.next()?;
1132 if access_token.is_empty() {
1133 return None;
1134 }
1135 Some((id, access_token.to_string()))
1136}
1137
1138#[cfg(test)]
1139mod tests {
1140 use super::*;
1141 use crate::test::{FakeHttpClient, FakeServer};
1142 use gpui::{executor::Deterministic, TestAppContext};
1143 use parking_lot::Mutex;
1144 use std::future;
1145
1146 #[gpui::test(iterations = 10)]
1147 async fn test_reconnection(cx: &mut TestAppContext) {
1148 cx.foreground().forbid_parking();
1149
1150 let user_id = 5;
1151 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1152 let server = FakeServer::for_client(user_id, &client, cx).await;
1153 let mut status = client.status();
1154 assert!(matches!(
1155 status.next().await,
1156 Some(Status::Connected { .. })
1157 ));
1158 assert_eq!(server.auth_count(), 1);
1159
1160 server.forbid_connections();
1161 server.disconnect();
1162 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1163
1164 server.allow_connections();
1165 cx.foreground().advance_clock(Duration::from_secs(10));
1166 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1167 assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1168
1169 server.forbid_connections();
1170 server.disconnect();
1171 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1172
1173 // Clear cached credentials after authentication fails
1174 server.roll_access_token();
1175 server.allow_connections();
1176 cx.foreground().advance_clock(Duration::from_secs(10));
1177 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1178 assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1179 }
1180
1181 #[gpui::test(iterations = 10)]
1182 async fn test_connection_timeout(deterministic: Arc<Deterministic>, cx: &mut TestAppContext) {
1183 deterministic.forbid_parking();
1184
1185 let user_id = 5;
1186 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1187 let mut status = client.status();
1188
1189 // Time out when client tries to connect.
1190 client.override_authenticate(move |cx| {
1191 cx.foreground().spawn(async move {
1192 Ok(Credentials {
1193 user_id,
1194 access_token: "token".into(),
1195 })
1196 })
1197 });
1198 client.override_establish_connection(|_, cx| {
1199 cx.foreground().spawn(async move {
1200 future::pending::<()>().await;
1201 unreachable!()
1202 })
1203 });
1204 let auth_and_connect = cx.spawn({
1205 let client = client.clone();
1206 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1207 });
1208 deterministic.run_until_parked();
1209 assert!(matches!(status.next().await, Some(Status::Connecting)));
1210
1211 deterministic.advance_clock(CONNECTION_TIMEOUT);
1212 assert!(matches!(
1213 status.next().await,
1214 Some(Status::ConnectionError { .. })
1215 ));
1216 auth_and_connect.await.unwrap_err();
1217
1218 // Allow the connection to be established.
1219 let server = FakeServer::for_client(user_id, &client, cx).await;
1220 assert!(matches!(
1221 status.next().await,
1222 Some(Status::Connected { .. })
1223 ));
1224
1225 // Disconnect client.
1226 server.forbid_connections();
1227 server.disconnect();
1228 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1229
1230 // Time out when re-establishing the connection.
1231 server.allow_connections();
1232 client.override_establish_connection(|_, cx| {
1233 cx.foreground().spawn(async move {
1234 future::pending::<()>().await;
1235 unreachable!()
1236 })
1237 });
1238 deterministic.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1239 assert!(matches!(
1240 status.next().await,
1241 Some(Status::Reconnecting { .. })
1242 ));
1243
1244 deterministic.advance_clock(CONNECTION_TIMEOUT);
1245 assert!(matches!(
1246 status.next().await,
1247 Some(Status::ReconnectionError { .. })
1248 ));
1249 }
1250
1251 #[gpui::test(iterations = 10)]
1252 async fn test_authenticating_more_than_once(
1253 cx: &mut TestAppContext,
1254 deterministic: Arc<Deterministic>,
1255 ) {
1256 cx.foreground().forbid_parking();
1257
1258 let auth_count = Arc::new(Mutex::new(0));
1259 let dropped_auth_count = Arc::new(Mutex::new(0));
1260 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1261 client.override_authenticate({
1262 let auth_count = auth_count.clone();
1263 let dropped_auth_count = dropped_auth_count.clone();
1264 move |cx| {
1265 let auth_count = auth_count.clone();
1266 let dropped_auth_count = dropped_auth_count.clone();
1267 cx.foreground().spawn(async move {
1268 *auth_count.lock() += 1;
1269 let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
1270 future::pending::<()>().await;
1271 unreachable!()
1272 })
1273 }
1274 });
1275
1276 let _authenticate = cx.spawn(|cx| {
1277 let client = client.clone();
1278 async move { client.authenticate_and_connect(false, &cx).await }
1279 });
1280 deterministic.run_until_parked();
1281 assert_eq!(*auth_count.lock(), 1);
1282 assert_eq!(*dropped_auth_count.lock(), 0);
1283
1284 let _authenticate = cx.spawn(|cx| {
1285 let client = client.clone();
1286 async move { client.authenticate_and_connect(false, &cx).await }
1287 });
1288 deterministic.run_until_parked();
1289 assert_eq!(*auth_count.lock(), 2);
1290 assert_eq!(*dropped_auth_count.lock(), 1);
1291 }
1292
1293 #[test]
1294 fn test_encode_and_decode_worktree_url() {
1295 let url = encode_worktree_url(5, "deadbeef");
1296 assert_eq!(decode_worktree_url(&url), Some((5, "deadbeef".to_string())));
1297 assert_eq!(
1298 decode_worktree_url(&format!("\n {}\t", url)),
1299 Some((5, "deadbeef".to_string()))
1300 );
1301 assert_eq!(decode_worktree_url("not://the-right-format"), None);
1302 }
1303
1304 #[gpui::test]
1305 async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
1306 cx.foreground().forbid_parking();
1307
1308 let user_id = 5;
1309 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1310 let server = FakeServer::for_client(user_id, &client, cx).await;
1311
1312 let (done_tx1, mut done_rx1) = smol::channel::unbounded();
1313 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1314 client.add_model_message_handler(
1315 move |model: ModelHandle<Model>, _: TypedEnvelope<proto::JoinProject>, _, cx| {
1316 match model.read_with(&cx, |model, _| model.id) {
1317 1 => done_tx1.try_send(()).unwrap(),
1318 2 => done_tx2.try_send(()).unwrap(),
1319 _ => unreachable!(),
1320 }
1321 async { Ok(()) }
1322 },
1323 );
1324 let model1 = cx.add_model(|_| Model {
1325 id: 1,
1326 subscription: None,
1327 });
1328 let model2 = cx.add_model(|_| Model {
1329 id: 2,
1330 subscription: None,
1331 });
1332 let model3 = cx.add_model(|_| Model {
1333 id: 3,
1334 subscription: None,
1335 });
1336
1337 let _subscription1 = model1.update(cx, |_, cx| client.add_model_for_remote_entity(1, cx));
1338 let _subscription2 = model2.update(cx, |_, cx| client.add_model_for_remote_entity(2, cx));
1339 // Ensure dropping a subscription for the same entity type still allows receiving of
1340 // messages for other entity IDs of the same type.
1341 let subscription3 = model3.update(cx, |_, cx| client.add_model_for_remote_entity(3, cx));
1342 drop(subscription3);
1343
1344 server.send(proto::JoinProject { project_id: 1 });
1345 server.send(proto::JoinProject { project_id: 2 });
1346 done_rx1.next().await.unwrap();
1347 done_rx2.next().await.unwrap();
1348 }
1349
1350 #[gpui::test]
1351 async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
1352 cx.foreground().forbid_parking();
1353
1354 let user_id = 5;
1355 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1356 let server = FakeServer::for_client(user_id, &client, cx).await;
1357
1358 let model = cx.add_model(|_| Model::default());
1359 let (done_tx1, _done_rx1) = smol::channel::unbounded();
1360 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1361 let subscription1 = client.add_message_handler(
1362 model.clone(),
1363 move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1364 done_tx1.try_send(()).unwrap();
1365 async { Ok(()) }
1366 },
1367 );
1368 drop(subscription1);
1369 let _subscription2 =
1370 client.add_message_handler(model, move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1371 done_tx2.try_send(()).unwrap();
1372 async { Ok(()) }
1373 });
1374 server.send(proto::Ping {});
1375 done_rx2.next().await.unwrap();
1376 }
1377
1378 #[gpui::test]
1379 async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
1380 cx.foreground().forbid_parking();
1381
1382 let user_id = 5;
1383 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1384 let server = FakeServer::for_client(user_id, &client, cx).await;
1385
1386 let model = cx.add_model(|_| Model::default());
1387 let (done_tx, mut done_rx) = smol::channel::unbounded();
1388 let subscription = client.add_message_handler(
1389 model.clone(),
1390 move |model, _: TypedEnvelope<proto::Ping>, _, mut cx| {
1391 model.update(&mut cx, |model, _| model.subscription.take());
1392 done_tx.try_send(()).unwrap();
1393 async { Ok(()) }
1394 },
1395 );
1396 model.update(cx, |model, _| {
1397 model.subscription = Some(subscription);
1398 });
1399 server.send(proto::Ping {});
1400 done_rx.next().await.unwrap();
1401 }
1402
1403 #[derive(Default)]
1404 struct Model {
1405 id: usize,
1406 subscription: Option<Subscription>,
1407 }
1408
1409 impl Entity for Model {
1410 type Event = ();
1411 }
1412}