1#[cfg(any(test, feature = "test-support"))]
2pub mod test;
3
4pub mod telemetry;
5pub mod user;
6
7use anyhow::{anyhow, Context as _, Result};
8use async_recursion::async_recursion;
9use async_tungstenite::tungstenite::{
10 error::Error as WebsocketError,
11 http::{Request, StatusCode},
12};
13use clock::SystemClock;
14use collections::HashMap;
15use futures::{
16 channel::oneshot, future::LocalBoxFuture, AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt,
17 TryFutureExt as _, TryStreamExt,
18};
19use gpui::{
20 actions, AnyModel, AnyWeakModel, AppContext, AsyncAppContext, Global, Model, Task, WeakModel,
21};
22use http::{HttpClient, HttpClientWithUrl};
23use lazy_static::lazy_static;
24use parking_lot::RwLock;
25use postage::watch;
26use rand::prelude::*;
27use release_channel::{AppVersion, ReleaseChannel};
28use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, PeerId, RequestMessage};
29use schemars::JsonSchema;
30use serde::{Deserialize, Serialize};
31use settings::{Settings, SettingsSources};
32use std::fmt;
33use std::pin::Pin;
34use std::{
35 any::TypeId,
36 convert::TryFrom,
37 fmt::Write as _,
38 future::Future,
39 marker::PhantomData,
40 path::PathBuf,
41 sync::{
42 atomic::{AtomicU64, Ordering},
43 Arc, Weak,
44 },
45 time::{Duration, Instant},
46};
47use telemetry::Telemetry;
48use thiserror::Error;
49use url::Url;
50use util::{ResultExt, TryFutureExt};
51
52pub use rpc::*;
53pub use telemetry_events::Event;
54pub use user::*;
55
56#[derive(Debug, Clone, Eq, PartialEq)]
57pub struct DevServerToken(pub String);
58
59impl fmt::Display for DevServerToken {
60 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61 write!(f, "{}", self.0)
62 }
63}
64
65lazy_static! {
66 static ref ZED_SERVER_URL: Option<String> = std::env::var("ZED_SERVER_URL").ok();
67 static ref ZED_RPC_URL: Option<String> = std::env::var("ZED_RPC_URL").ok();
68 /// An environment variable whose presence indicates that the development auth
69 /// provider should be used.
70 ///
71 /// Only works in development. Setting this environment variable in other release
72 /// channels is a no-op.
73 pub static ref ZED_DEVELOPMENT_AUTH: bool =
74 std::env::var("ZED_DEVELOPMENT_AUTH").map_or(false, |value| !value.is_empty());
75 pub static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE")
76 .ok()
77 .and_then(|s| if s.is_empty() { None } else { Some(s) });
78 pub static ref ADMIN_API_TOKEN: Option<String> = std::env::var("ZED_ADMIN_API_TOKEN")
79 .ok()
80 .and_then(|s| if s.is_empty() { None } else { Some(s) });
81 pub static ref ZED_APP_PATH: Option<PathBuf> =
82 std::env::var("ZED_APP_PATH").ok().map(PathBuf::from);
83 pub static ref ZED_ALWAYS_ACTIVE: bool =
84 std::env::var("ZED_ALWAYS_ACTIVE").map_or(false, |e| !e.is_empty());
85}
86
87pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(500);
88pub const MAX_RECONNECTION_DELAY: Duration = Duration::from_secs(10);
89pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20);
90
91actions!(client, [SignIn, SignOut, Reconnect]);
92
93#[derive(Clone, Default, Serialize, Deserialize, JsonSchema)]
94pub struct ClientSettingsContent {
95 server_url: Option<String>,
96}
97
98#[derive(Deserialize)]
99pub struct ClientSettings {
100 pub server_url: String,
101}
102
103impl Settings for ClientSettings {
104 const KEY: Option<&'static str> = None;
105
106 type FileContent = ClientSettingsContent;
107
108 fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {
109 let mut result = sources.json_merge::<Self>()?;
110 if let Some(server_url) = &*ZED_SERVER_URL {
111 result.server_url.clone_from(&server_url)
112 }
113 Ok(result)
114 }
115}
116
117#[derive(Default, Clone, Serialize, Deserialize, JsonSchema)]
118pub struct ProxySettingsContent {
119 proxy: Option<String>,
120}
121
122#[derive(Deserialize, Default)]
123pub struct ProxySettings {
124 pub proxy: Option<String>,
125}
126
127impl Settings for ProxySettings {
128 const KEY: Option<&'static str> = None;
129
130 type FileContent = ProxySettingsContent;
131
132 fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {
133 Ok(Self {
134 proxy: sources
135 .user
136 .and_then(|value| value.proxy.clone())
137 .or(sources.default.proxy.clone()),
138 })
139 }
140}
141
142pub fn init_settings(cx: &mut AppContext) {
143 TelemetrySettings::register(cx);
144 ClientSettings::register(cx);
145 ProxySettings::register(cx);
146}
147
148pub fn init(client: &Arc<Client>, cx: &mut AppContext) {
149 let client = Arc::downgrade(client);
150 cx.on_action({
151 let client = client.clone();
152 move |_: &SignIn, cx| {
153 if let Some(client) = client.upgrade() {
154 cx.spawn(
155 |cx| async move { client.authenticate_and_connect(true, &cx).log_err().await },
156 )
157 .detach();
158 }
159 }
160 });
161
162 cx.on_action({
163 let client = client.clone();
164 move |_: &SignOut, cx| {
165 if let Some(client) = client.upgrade() {
166 cx.spawn(|cx| async move {
167 client.sign_out(&cx).await;
168 })
169 .detach();
170 }
171 }
172 });
173
174 cx.on_action({
175 let client = client.clone();
176 move |_: &Reconnect, cx| {
177 if let Some(client) = client.upgrade() {
178 cx.spawn(|cx| async move {
179 client.reconnect(&cx);
180 })
181 .detach();
182 }
183 }
184 });
185}
186
187struct GlobalClient(Arc<Client>);
188
189impl Global for GlobalClient {}
190
191pub struct Client {
192 id: AtomicU64,
193 peer: Arc<Peer>,
194 http: Arc<HttpClientWithUrl>,
195 telemetry: Arc<Telemetry>,
196 credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static>,
197 state: RwLock<ClientState>,
198
199 #[allow(clippy::type_complexity)]
200 #[cfg(any(test, feature = "test-support"))]
201 authenticate: RwLock<
202 Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,
203 >,
204
205 #[allow(clippy::type_complexity)]
206 #[cfg(any(test, feature = "test-support"))]
207 establish_connection: RwLock<
208 Option<
209 Box<
210 dyn 'static
211 + Send
212 + Sync
213 + Fn(
214 &Credentials,
215 &AsyncAppContext,
216 ) -> Task<Result<Connection, EstablishConnectionError>>,
217 >,
218 >,
219 >,
220}
221
222#[derive(Error, Debug)]
223pub enum EstablishConnectionError {
224 #[error("upgrade required")]
225 UpgradeRequired,
226 #[error("unauthorized")]
227 Unauthorized,
228 #[error("{0}")]
229 Other(#[from] anyhow::Error),
230 #[error("{0}")]
231 Http(#[from] http::Error),
232 #[error("{0}")]
233 Io(#[from] std::io::Error),
234 #[error("{0}")]
235 Websocket(#[from] async_tungstenite::tungstenite::http::Error),
236}
237
238impl From<WebsocketError> for EstablishConnectionError {
239 fn from(error: WebsocketError) -> Self {
240 if let WebsocketError::Http(response) = &error {
241 match response.status() {
242 StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,
243 StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,
244 _ => {}
245 }
246 }
247 EstablishConnectionError::Other(error.into())
248 }
249}
250
251impl EstablishConnectionError {
252 pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
253 Self::Other(error.into())
254 }
255}
256
257#[derive(Copy, Clone, Debug, PartialEq)]
258pub enum Status {
259 SignedOut,
260 UpgradeRequired,
261 Authenticating,
262 Connecting,
263 ConnectionError,
264 Connected {
265 peer_id: PeerId,
266 connection_id: ConnectionId,
267 },
268 ConnectionLost,
269 Reauthenticating,
270 Reconnecting,
271 ReconnectionError {
272 next_reconnection: Instant,
273 },
274}
275
276impl Status {
277 pub fn is_connected(&self) -> bool {
278 matches!(self, Self::Connected { .. })
279 }
280
281 pub fn is_signed_out(&self) -> bool {
282 matches!(self, Self::SignedOut | Self::UpgradeRequired)
283 }
284}
285
286struct ClientState {
287 credentials: Option<Credentials>,
288 status: (watch::Sender<Status>, watch::Receiver<Status>),
289 entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
290 _reconnect_task: Option<Task<()>>,
291 entities_by_type_and_remote_id: HashMap<(TypeId, u64), WeakSubscriber>,
292 models_by_message_type: HashMap<TypeId, AnyWeakModel>,
293 entity_types_by_message_type: HashMap<TypeId, TypeId>,
294 #[allow(clippy::type_complexity)]
295 message_handlers: HashMap<
296 TypeId,
297 Arc<
298 dyn Send
299 + Sync
300 + Fn(
301 AnyModel,
302 Box<dyn AnyTypedEnvelope>,
303 &Arc<Client>,
304 AsyncAppContext,
305 ) -> LocalBoxFuture<'static, Result<()>>,
306 >,
307 >,
308}
309
310enum WeakSubscriber {
311 Entity { handle: AnyWeakModel },
312 Pending(Vec<Box<dyn AnyTypedEnvelope>>),
313}
314
315#[derive(Clone, Debug, Eq, PartialEq)]
316pub enum Credentials {
317 DevServer { token: DevServerToken },
318 User { user_id: u64, access_token: String },
319}
320
321impl Credentials {
322 pub fn authorization_header(&self) -> String {
323 match self {
324 Credentials::DevServer { token } => format!("dev-server-token {}", token),
325 Credentials::User {
326 user_id,
327 access_token,
328 } => format!("{} {}", user_id, access_token),
329 }
330 }
331}
332
333/// A provider for [`Credentials`].
334///
335/// Used to abstract over reading and writing credentials to some form of
336/// persistence (like the system keychain).
337trait CredentialsProvider {
338 /// Reads the credentials from the provider.
339 fn read_credentials<'a>(
340 &'a self,
341 cx: &'a AsyncAppContext,
342 ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>>;
343
344 /// Writes the credentials to the provider.
345 fn write_credentials<'a>(
346 &'a self,
347 user_id: u64,
348 access_token: String,
349 cx: &'a AsyncAppContext,
350 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>>;
351
352 /// Deletes the credentials from the provider.
353 fn delete_credentials<'a>(
354 &'a self,
355 cx: &'a AsyncAppContext,
356 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>>;
357}
358
359impl Default for ClientState {
360 fn default() -> Self {
361 Self {
362 credentials: None,
363 status: watch::channel_with(Status::SignedOut),
364 entity_id_extractors: Default::default(),
365 _reconnect_task: None,
366 models_by_message_type: Default::default(),
367 entities_by_type_and_remote_id: Default::default(),
368 entity_types_by_message_type: Default::default(),
369 message_handlers: Default::default(),
370 }
371 }
372}
373
374pub enum Subscription {
375 Entity {
376 client: Weak<Client>,
377 id: (TypeId, u64),
378 },
379 Message {
380 client: Weak<Client>,
381 id: TypeId,
382 },
383}
384
385impl Drop for Subscription {
386 fn drop(&mut self) {
387 match self {
388 Subscription::Entity { client, id } => {
389 if let Some(client) = client.upgrade() {
390 let mut state = client.state.write();
391 let _ = state.entities_by_type_and_remote_id.remove(id);
392 }
393 }
394 Subscription::Message { client, id } => {
395 if let Some(client) = client.upgrade() {
396 let mut state = client.state.write();
397 let _ = state.entity_types_by_message_type.remove(id);
398 let _ = state.message_handlers.remove(id);
399 }
400 }
401 }
402 }
403}
404
405pub struct PendingEntitySubscription<T: 'static> {
406 client: Arc<Client>,
407 remote_id: u64,
408 _entity_type: PhantomData<T>,
409 consumed: bool,
410}
411
412impl<T: 'static> PendingEntitySubscription<T> {
413 pub fn set_model(mut self, model: &Model<T>, cx: &mut AsyncAppContext) -> Subscription {
414 self.consumed = true;
415 let mut state = self.client.state.write();
416 let id = (TypeId::of::<T>(), self.remote_id);
417 let Some(WeakSubscriber::Pending(messages)) =
418 state.entities_by_type_and_remote_id.remove(&id)
419 else {
420 unreachable!()
421 };
422
423 state.entities_by_type_and_remote_id.insert(
424 id,
425 WeakSubscriber::Entity {
426 handle: model.downgrade().into(),
427 },
428 );
429 drop(state);
430 for message in messages {
431 self.client.handle_message(message, cx);
432 }
433 Subscription::Entity {
434 client: Arc::downgrade(&self.client),
435 id,
436 }
437 }
438}
439
440impl<T: 'static> Drop for PendingEntitySubscription<T> {
441 fn drop(&mut self) {
442 if !self.consumed {
443 let mut state = self.client.state.write();
444 if let Some(WeakSubscriber::Pending(messages)) = state
445 .entities_by_type_and_remote_id
446 .remove(&(TypeId::of::<T>(), self.remote_id))
447 {
448 for message in messages {
449 log::info!("unhandled message {}", message.payload_type_name());
450 }
451 }
452 }
453 }
454}
455
456#[derive(Copy, Clone)]
457pub struct TelemetrySettings {
458 pub diagnostics: bool,
459 pub metrics: bool,
460}
461
462/// Control what info is collected by Zed.
463#[derive(Default, Clone, Serialize, Deserialize, JsonSchema)]
464pub struct TelemetrySettingsContent {
465 /// Send debug info like crash reports.
466 ///
467 /// Default: true
468 pub diagnostics: Option<bool>,
469 /// Send anonymized usage data like what languages you're using Zed with.
470 ///
471 /// Default: true
472 pub metrics: Option<bool>,
473}
474
475impl settings::Settings for TelemetrySettings {
476 const KEY: Option<&'static str> = Some("telemetry");
477
478 type FileContent = TelemetrySettingsContent;
479
480 fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {
481 Ok(Self {
482 diagnostics: sources.user.as_ref().and_then(|v| v.diagnostics).unwrap_or(
483 sources
484 .default
485 .diagnostics
486 .ok_or_else(Self::missing_default)?,
487 ),
488 metrics: sources
489 .user
490 .as_ref()
491 .and_then(|v| v.metrics)
492 .unwrap_or(sources.default.metrics.ok_or_else(Self::missing_default)?),
493 })
494 }
495}
496
497impl Client {
498 pub fn new(
499 clock: Arc<dyn SystemClock>,
500 http: Arc<HttpClientWithUrl>,
501 cx: &mut AppContext,
502 ) -> Arc<Self> {
503 let use_zed_development_auth = match ReleaseChannel::try_global(cx) {
504 Some(ReleaseChannel::Dev) => *ZED_DEVELOPMENT_AUTH,
505 Some(ReleaseChannel::Nightly | ReleaseChannel::Preview | ReleaseChannel::Stable)
506 | None => false,
507 };
508
509 let credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static> =
510 if use_zed_development_auth {
511 Arc::new(DevelopmentCredentialsProvider {
512 path: util::paths::CONFIG_DIR.join("development_auth"),
513 })
514 } else {
515 Arc::new(KeychainCredentialsProvider)
516 };
517
518 Arc::new(Self {
519 id: AtomicU64::new(0),
520 peer: Peer::new(0),
521 telemetry: Telemetry::new(clock, http.clone(), cx),
522 http,
523 credentials_provider,
524 state: Default::default(),
525
526 #[cfg(any(test, feature = "test-support"))]
527 authenticate: Default::default(),
528 #[cfg(any(test, feature = "test-support"))]
529 establish_connection: Default::default(),
530 })
531 }
532
533 pub fn production(cx: &mut AppContext) -> Arc<Self> {
534 let clock = Arc::new(clock::RealSystemClock);
535 let http = Arc::new(HttpClientWithUrl::new(
536 &ClientSettings::get_global(cx).server_url,
537 ProxySettings::get_global(cx).proxy.clone(),
538 ));
539 Self::new(clock, http.clone(), cx)
540 }
541
542 pub fn id(&self) -> u64 {
543 self.id.load(Ordering::SeqCst)
544 }
545
546 pub fn http_client(&self) -> Arc<HttpClientWithUrl> {
547 self.http.clone()
548 }
549
550 pub fn set_id(&self, id: u64) -> &Self {
551 self.id.store(id, Ordering::SeqCst);
552 self
553 }
554
555 #[cfg(any(test, feature = "test-support"))]
556 pub fn teardown(&self) {
557 let mut state = self.state.write();
558 state._reconnect_task.take();
559 state.message_handlers.clear();
560 state.models_by_message_type.clear();
561 state.entities_by_type_and_remote_id.clear();
562 state.entity_id_extractors.clear();
563 self.peer.teardown();
564 }
565
566 #[cfg(any(test, feature = "test-support"))]
567 pub fn override_authenticate<F>(&self, authenticate: F) -> &Self
568 where
569 F: 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>,
570 {
571 *self.authenticate.write() = Some(Box::new(authenticate));
572 self
573 }
574
575 #[cfg(any(test, feature = "test-support"))]
576 pub fn override_establish_connection<F>(&self, connect: F) -> &Self
577 where
578 F: 'static
579 + Send
580 + Sync
581 + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,
582 {
583 *self.establish_connection.write() = Some(Box::new(connect));
584 self
585 }
586
587 pub fn global(cx: &AppContext) -> Arc<Self> {
588 cx.global::<GlobalClient>().0.clone()
589 }
590 pub fn set_global(client: Arc<Client>, cx: &mut AppContext) {
591 cx.set_global(GlobalClient(client))
592 }
593
594 pub fn user_id(&self) -> Option<u64> {
595 if let Some(Credentials::User { user_id, .. }) = self.state.read().credentials.as_ref() {
596 Some(*user_id)
597 } else {
598 None
599 }
600 }
601
602 pub fn peer_id(&self) -> Option<PeerId> {
603 if let Status::Connected { peer_id, .. } = &*self.status().borrow() {
604 Some(*peer_id)
605 } else {
606 None
607 }
608 }
609
610 pub fn status(&self) -> watch::Receiver<Status> {
611 self.state.read().status.1.clone()
612 }
613
614 fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncAppContext) {
615 log::info!("set status on client {}: {:?}", self.id(), status);
616 let mut state = self.state.write();
617 *state.status.0.borrow_mut() = status;
618
619 match status {
620 Status::Connected { .. } => {
621 state._reconnect_task = None;
622 }
623 Status::ConnectionLost => {
624 let this = self.clone();
625 state._reconnect_task = Some(cx.spawn(move |cx| async move {
626 #[cfg(any(test, feature = "test-support"))]
627 let mut rng = StdRng::seed_from_u64(0);
628 #[cfg(not(any(test, feature = "test-support")))]
629 let mut rng = StdRng::from_entropy();
630
631 let mut delay = INITIAL_RECONNECTION_DELAY;
632 while let Err(error) = this.authenticate_and_connect(true, &cx).await {
633 log::error!("failed to connect {}", error);
634 if matches!(*this.status().borrow(), Status::ConnectionError) {
635 this.set_status(
636 Status::ReconnectionError {
637 next_reconnection: Instant::now() + delay,
638 },
639 &cx,
640 );
641 cx.background_executor().timer(delay).await;
642 delay = delay
643 .mul_f32(rng.gen_range(0.5..=2.5))
644 .max(INITIAL_RECONNECTION_DELAY)
645 .min(MAX_RECONNECTION_DELAY);
646 } else {
647 break;
648 }
649 }
650 }));
651 }
652 Status::SignedOut | Status::UpgradeRequired => {
653 self.telemetry.set_authenticated_user_info(None, false);
654 state._reconnect_task.take();
655 }
656 _ => {}
657 }
658 }
659
660 pub fn subscribe_to_entity<T>(
661 self: &Arc<Self>,
662 remote_id: u64,
663 ) -> Result<PendingEntitySubscription<T>>
664 where
665 T: 'static,
666 {
667 let id = (TypeId::of::<T>(), remote_id);
668
669 let mut state = self.state.write();
670 if state.entities_by_type_and_remote_id.contains_key(&id) {
671 return Err(anyhow!("already subscribed to entity"));
672 }
673
674 state
675 .entities_by_type_and_remote_id
676 .insert(id, WeakSubscriber::Pending(Default::default()));
677
678 Ok(PendingEntitySubscription {
679 client: self.clone(),
680 remote_id,
681 consumed: false,
682 _entity_type: PhantomData,
683 })
684 }
685
686 #[track_caller]
687 pub fn add_message_handler<M, E, H, F>(
688 self: &Arc<Self>,
689 entity: WeakModel<E>,
690 handler: H,
691 ) -> Subscription
692 where
693 M: EnvelopedMessage,
694 E: 'static,
695 H: 'static
696 + Sync
697 + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F
698 + Send
699 + Sync,
700 F: 'static + Future<Output = Result<()>>,
701 {
702 let message_type_id = TypeId::of::<M>();
703 let mut state = self.state.write();
704 state
705 .models_by_message_type
706 .insert(message_type_id, entity.into());
707
708 let prev_handler = state.message_handlers.insert(
709 message_type_id,
710 Arc::new(move |subscriber, envelope, client, cx| {
711 let subscriber = subscriber.downcast::<E>().unwrap();
712 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
713 handler(subscriber, *envelope, client.clone(), cx).boxed_local()
714 }),
715 );
716 if prev_handler.is_some() {
717 let location = std::panic::Location::caller();
718 panic!(
719 "{}:{} registered handler for the same message {} twice",
720 location.file(),
721 location.line(),
722 std::any::type_name::<M>()
723 );
724 }
725
726 Subscription::Message {
727 client: Arc::downgrade(self),
728 id: message_type_id,
729 }
730 }
731
732 pub fn add_request_handler<M, E, H, F>(
733 self: &Arc<Self>,
734 model: WeakModel<E>,
735 handler: H,
736 ) -> Subscription
737 where
738 M: RequestMessage,
739 E: 'static,
740 H: 'static
741 + Sync
742 + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F
743 + Send
744 + Sync,
745 F: 'static + Future<Output = Result<M::Response>>,
746 {
747 self.add_message_handler(model, move |handle, envelope, this, cx| {
748 Self::respond_to_request(
749 envelope.receipt(),
750 handler(handle, envelope, this.clone(), cx),
751 this,
752 )
753 })
754 }
755
756 pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
757 where
758 M: EntityMessage,
759 E: 'static,
760 H: 'static + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F + Send + Sync,
761 F: 'static + Future<Output = Result<()>>,
762 {
763 self.add_entity_message_handler::<M, E, _, _>(move |subscriber, message, client, cx| {
764 handler(subscriber.downcast::<E>().unwrap(), message, client, cx)
765 })
766 }
767
768 fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
769 where
770 M: EntityMessage,
771 E: 'static,
772 H: 'static + Fn(AnyModel, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F + Send + Sync,
773 F: 'static + Future<Output = Result<()>>,
774 {
775 let model_type_id = TypeId::of::<E>();
776 let message_type_id = TypeId::of::<M>();
777
778 let mut state = self.state.write();
779 state
780 .entity_types_by_message_type
781 .insert(message_type_id, model_type_id);
782 state
783 .entity_id_extractors
784 .entry(message_type_id)
785 .or_insert_with(|| {
786 |envelope| {
787 envelope
788 .as_any()
789 .downcast_ref::<TypedEnvelope<M>>()
790 .unwrap()
791 .payload
792 .remote_entity_id()
793 }
794 });
795 let prev_handler = state.message_handlers.insert(
796 message_type_id,
797 Arc::new(move |handle, envelope, client, cx| {
798 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
799 handler(handle, *envelope, client.clone(), cx).boxed_local()
800 }),
801 );
802 if prev_handler.is_some() {
803 panic!("registered handler for the same message twice");
804 }
805 }
806
807 pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
808 where
809 M: EntityMessage + RequestMessage,
810 E: 'static,
811 H: 'static + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F + Send + Sync,
812 F: 'static + Future<Output = Result<M::Response>>,
813 {
814 self.add_model_message_handler(move |entity, envelope, client, cx| {
815 Self::respond_to_request::<M, _>(
816 envelope.receipt(),
817 handler(entity, envelope, client.clone(), cx),
818 client,
819 )
820 })
821 }
822
823 async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
824 receipt: Receipt<T>,
825 response: F,
826 client: Arc<Self>,
827 ) -> Result<()> {
828 match response.await {
829 Ok(response) => {
830 client.respond(receipt, response)?;
831 Ok(())
832 }
833 Err(error) => {
834 client.respond_with_error(receipt, error.to_proto())?;
835 Err(error)
836 }
837 }
838 }
839
840 pub async fn has_credentials(&self, cx: &AsyncAppContext) -> bool {
841 self.credentials_provider
842 .read_credentials(cx)
843 .await
844 .is_some()
845 }
846
847 pub fn set_dev_server_token(&self, token: DevServerToken) -> &Self {
848 self.state.write().credentials = Some(Credentials::DevServer { token });
849 self
850 }
851
852 #[async_recursion(?Send)]
853 pub async fn authenticate_and_connect(
854 self: &Arc<Self>,
855 try_provider: bool,
856 cx: &AsyncAppContext,
857 ) -> anyhow::Result<()> {
858 let was_disconnected = match *self.status().borrow() {
859 Status::SignedOut => true,
860 Status::ConnectionError
861 | Status::ConnectionLost
862 | Status::Authenticating { .. }
863 | Status::Reauthenticating { .. }
864 | Status::ReconnectionError { .. } => false,
865 Status::Connected { .. } | Status::Connecting { .. } | Status::Reconnecting { .. } => {
866 return Ok(())
867 }
868 Status::UpgradeRequired => return Err(EstablishConnectionError::UpgradeRequired)?,
869 };
870 if was_disconnected {
871 self.set_status(Status::Authenticating, cx);
872 } else {
873 self.set_status(Status::Reauthenticating, cx)
874 }
875
876 let mut read_from_provider = false;
877 let mut credentials = self.state.read().credentials.clone();
878 if credentials.is_none() && try_provider {
879 credentials = self.credentials_provider.read_credentials(cx).await;
880 read_from_provider = credentials.is_some();
881 }
882
883 if credentials.is_none() {
884 let mut status_rx = self.status();
885 let _ = status_rx.next().await;
886 futures::select_biased! {
887 authenticate = self.authenticate(cx).fuse() => {
888 match authenticate {
889 Ok(creds) => credentials = Some(creds),
890 Err(err) => {
891 self.set_status(Status::ConnectionError, cx);
892 return Err(err);
893 }
894 }
895 }
896 _ = status_rx.next().fuse() => {
897 return Err(anyhow!("authentication canceled"));
898 }
899 }
900 }
901 let credentials = credentials.unwrap();
902 if let Credentials::User { user_id, .. } = &credentials {
903 self.set_id(*user_id);
904 }
905
906 if was_disconnected {
907 self.set_status(Status::Connecting, cx);
908 } else {
909 self.set_status(Status::Reconnecting, cx);
910 }
911
912 let mut timeout =
913 futures::FutureExt::fuse(cx.background_executor().timer(CONNECTION_TIMEOUT));
914 futures::select_biased! {
915 connection = self.establish_connection(&credentials, cx).fuse() => {
916 match connection {
917 Ok(conn) => {
918 self.state.write().credentials = Some(credentials.clone());
919 if !read_from_provider && IMPERSONATE_LOGIN.is_none() {
920 if let Credentials::User{user_id, access_token} = credentials {
921 self.credentials_provider.write_credentials(user_id, access_token, cx).await.log_err();
922 }
923 }
924
925 futures::select_biased! {
926 result = self.set_connection(conn, cx).fuse() => result,
927 _ = timeout => {
928 self.set_status(Status::ConnectionError, cx);
929 Err(anyhow!("timed out waiting on hello message from server"))
930 }
931 }
932 }
933 Err(EstablishConnectionError::Unauthorized) => {
934 self.state.write().credentials.take();
935 if read_from_provider {
936 self.credentials_provider.delete_credentials(cx).await.log_err();
937 self.set_status(Status::SignedOut, cx);
938 self.authenticate_and_connect(false, cx).await
939 } else {
940 self.set_status(Status::ConnectionError, cx);
941 Err(EstablishConnectionError::Unauthorized)?
942 }
943 }
944 Err(EstablishConnectionError::UpgradeRequired) => {
945 self.set_status(Status::UpgradeRequired, cx);
946 Err(EstablishConnectionError::UpgradeRequired)?
947 }
948 Err(error) => {
949 self.set_status(Status::ConnectionError, cx);
950 Err(error)?
951 }
952 }
953 }
954 _ = &mut timeout => {
955 self.set_status(Status::ConnectionError, cx);
956 Err(anyhow!("timed out trying to establish connection"))
957 }
958 }
959 }
960
961 async fn set_connection(
962 self: &Arc<Self>,
963 conn: Connection,
964 cx: &AsyncAppContext,
965 ) -> Result<()> {
966 let executor = cx.background_executor();
967 log::info!("add connection to peer");
968 let (connection_id, handle_io, mut incoming) = self.peer.add_connection(conn, {
969 let executor = executor.clone();
970 move |duration| executor.timer(duration)
971 });
972 let handle_io = executor.spawn(handle_io);
973
974 let peer_id = async {
975 log::info!("waiting for server hello");
976 let message = incoming
977 .next()
978 .await
979 .ok_or_else(|| anyhow!("no hello message received"))?;
980 log::info!("got server hello");
981 let hello_message_type_name = message.payload_type_name().to_string();
982 let hello = message
983 .into_any()
984 .downcast::<TypedEnvelope<proto::Hello>>()
985 .map_err(|_| {
986 anyhow!(
987 "invalid hello message received: {:?}",
988 hello_message_type_name
989 )
990 })?;
991 let peer_id = hello
992 .payload
993 .peer_id
994 .ok_or_else(|| anyhow!("invalid peer id"))?;
995 Ok(peer_id)
996 };
997
998 let peer_id = match peer_id.await {
999 Ok(peer_id) => peer_id,
1000 Err(error) => {
1001 self.peer.disconnect(connection_id);
1002 return Err(error);
1003 }
1004 };
1005
1006 log::info!(
1007 "set status to connected (connection id: {:?}, peer id: {:?})",
1008 connection_id,
1009 peer_id
1010 );
1011 self.set_status(
1012 Status::Connected {
1013 peer_id,
1014 connection_id,
1015 },
1016 cx,
1017 );
1018
1019 cx.spawn({
1020 let this = self.clone();
1021 |cx| {
1022 async move {
1023 while let Some(message) = incoming.next().await {
1024 this.handle_message(message, &cx);
1025 // Don't starve the main thread when receiving lots of messages at once.
1026 smol::future::yield_now().await;
1027 }
1028 }
1029 }
1030 })
1031 .detach();
1032
1033 cx.spawn({
1034 let this = self.clone();
1035 move |cx| async move {
1036 match handle_io.await {
1037 Ok(()) => {
1038 if *this.status().borrow()
1039 == (Status::Connected {
1040 connection_id,
1041 peer_id,
1042 })
1043 {
1044 this.set_status(Status::SignedOut, &cx);
1045 }
1046 }
1047 Err(err) => {
1048 log::error!("connection error: {:?}", err);
1049 this.set_status(Status::ConnectionLost, &cx);
1050 }
1051 }
1052 }
1053 })
1054 .detach();
1055
1056 Ok(())
1057 }
1058
1059 fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {
1060 #[cfg(any(test, feature = "test-support"))]
1061 if let Some(callback) = self.authenticate.read().as_ref() {
1062 return callback(cx);
1063 }
1064
1065 self.authenticate_with_browser(cx)
1066 }
1067
1068 fn establish_connection(
1069 self: &Arc<Self>,
1070 credentials: &Credentials,
1071 cx: &AsyncAppContext,
1072 ) -> Task<Result<Connection, EstablishConnectionError>> {
1073 #[cfg(any(test, feature = "test-support"))]
1074 if let Some(callback) = self.establish_connection.read().as_ref() {
1075 return callback(credentials, cx);
1076 }
1077
1078 self.establish_websocket_connection(credentials, cx)
1079 }
1080
1081 async fn get_rpc_url(
1082 http: Arc<HttpClientWithUrl>,
1083 release_channel: Option<ReleaseChannel>,
1084 ) -> Result<Url> {
1085 if let Some(url) = &*ZED_RPC_URL {
1086 return Url::parse(url).context("invalid rpc url");
1087 }
1088
1089 let mut url = http.build_url("/rpc");
1090 if let Some(preview_param) =
1091 release_channel.and_then(|channel| channel.release_query_param())
1092 {
1093 url += "?";
1094 url += preview_param;
1095 }
1096 let response = http.get(&url, Default::default(), false).await?;
1097 let collab_url = if response.status().is_redirection() {
1098 response
1099 .headers()
1100 .get("Location")
1101 .ok_or_else(|| anyhow!("missing location header in /rpc response"))?
1102 .to_str()
1103 .map_err(EstablishConnectionError::other)?
1104 .to_string()
1105 } else {
1106 Err(anyhow!(
1107 "unexpected /rpc response status {}",
1108 response.status()
1109 ))?
1110 };
1111
1112 Url::parse(&collab_url).context("invalid rpc url")
1113 }
1114
1115 fn establish_websocket_connection(
1116 self: &Arc<Self>,
1117 credentials: &Credentials,
1118 cx: &AsyncAppContext,
1119 ) -> Task<Result<Connection, EstablishConnectionError>> {
1120 let release_channel = cx
1121 .update(|cx| ReleaseChannel::try_global(cx))
1122 .ok()
1123 .flatten();
1124 let app_version = cx
1125 .update(|cx| AppVersion::global(cx).to_string())
1126 .ok()
1127 .unwrap_or_default();
1128
1129 let request = Request::builder()
1130 .header("Authorization", credentials.authorization_header())
1131 .header("x-zed-protocol-version", rpc::PROTOCOL_VERSION)
1132 .header("x-zed-app-version", app_version)
1133 .header(
1134 "x-zed-release-channel",
1135 release_channel.map(|r| r.dev_name()).unwrap_or("unknown"),
1136 );
1137
1138 let http = self.http.clone();
1139 cx.background_executor().spawn(async move {
1140 let mut rpc_url = Self::get_rpc_url(http, release_channel).await?;
1141 let rpc_host = rpc_url
1142 .host_str()
1143 .zip(rpc_url.port_or_known_default())
1144 .ok_or_else(|| anyhow!("missing host in rpc url"))?;
1145 let stream = smol::net::TcpStream::connect(rpc_host).await?;
1146
1147 log::info!("connected to rpc endpoint {}", rpc_url);
1148
1149 match rpc_url.scheme() {
1150 "https" => {
1151 rpc_url.set_scheme("wss").unwrap();
1152 let request = request.uri(rpc_url.as_str()).body(())?;
1153 let (stream, _) =
1154 async_tungstenite::async_std::client_async_tls(request, stream).await?;
1155 Ok(Connection::new(
1156 stream
1157 .map_err(|error| anyhow!(error))
1158 .sink_map_err(|error| anyhow!(error)),
1159 ))
1160 }
1161 "http" => {
1162 rpc_url.set_scheme("ws").unwrap();
1163 let request = request.uri(rpc_url.as_str()).body(())?;
1164 let (stream, _) = async_tungstenite::client_async(request, stream).await?;
1165 Ok(Connection::new(
1166 stream
1167 .map_err(|error| anyhow!(error))
1168 .sink_map_err(|error| anyhow!(error)),
1169 ))
1170 }
1171 _ => Err(anyhow!("invalid rpc url: {}", rpc_url))?,
1172 }
1173 })
1174 }
1175
1176 pub fn authenticate_with_browser(
1177 self: &Arc<Self>,
1178 cx: &AsyncAppContext,
1179 ) -> Task<Result<Credentials>> {
1180 let http = self.http.clone();
1181 cx.spawn(|cx| async move {
1182 let background = cx.background_executor().clone();
1183
1184 let (open_url_tx, open_url_rx) = oneshot::channel::<String>();
1185 cx.update(|cx| {
1186 cx.spawn(move |cx| async move {
1187 let url = open_url_rx.await?;
1188 cx.update(|cx| cx.open_url(&url))
1189 })
1190 .detach_and_log_err(cx);
1191 })
1192 .log_err();
1193
1194 let credentials = background
1195 .clone()
1196 .spawn(async move {
1197 // Generate a pair of asymmetric encryption keys. The public key will be used by the
1198 // zed server to encrypt the user's access token, so that it can'be intercepted by
1199 // any other app running on the user's device.
1200 let (public_key, private_key) =
1201 rpc::auth::keypair().expect("failed to generate keypair for auth");
1202 let public_key_string = String::try_from(public_key)
1203 .expect("failed to serialize public key for auth");
1204
1205 if let Some((login, token)) =
1206 IMPERSONATE_LOGIN.as_ref().zip(ADMIN_API_TOKEN.as_ref())
1207 {
1208 eprintln!("authenticate as admin {login}, {token}");
1209
1210 return Self::authenticate_as_admin(http, login.clone(), token.clone())
1211 .await;
1212 }
1213
1214 // Start an HTTP server to receive the redirect from Zed's sign-in page.
1215 let server =
1216 tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
1217 let port = server.server_addr().port();
1218
1219 // Open the Zed sign-in page in the user's browser, with query parameters that indicate
1220 // that the user is signing in from a Zed app running on the same device.
1221 let mut url = http.build_url(&format!(
1222 "/native_app_signin?native_app_port={}&native_app_public_key={}",
1223 port, public_key_string
1224 ));
1225
1226 if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
1227 log::info!("impersonating user @{}", impersonate_login);
1228 write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
1229 }
1230
1231 open_url_tx.send(url).log_err();
1232
1233 // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
1234 // access token from the query params.
1235 //
1236 // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
1237 // custom URL scheme instead of this local HTTP server.
1238 let (user_id, access_token) = background
1239 .spawn(async move {
1240 for _ in 0..100 {
1241 if let Some(req) = server.recv_timeout(Duration::from_secs(1))? {
1242 let path = req.url();
1243 let mut user_id = None;
1244 let mut access_token = None;
1245 let url = Url::parse(&format!("http://example.com{}", path))
1246 .context("failed to parse login notification url")?;
1247 for (key, value) in url.query_pairs() {
1248 if key == "access_token" {
1249 access_token = Some(value.to_string());
1250 } else if key == "user_id" {
1251 user_id = Some(value.to_string());
1252 }
1253 }
1254
1255 let post_auth_url =
1256 http.build_url("/native_app_signin_succeeded");
1257 req.respond(
1258 tiny_http::Response::empty(302).with_header(
1259 tiny_http::Header::from_bytes(
1260 &b"Location"[..],
1261 post_auth_url.as_bytes(),
1262 )
1263 .unwrap(),
1264 ),
1265 )
1266 .context("failed to respond to login http request")?;
1267 return Ok((
1268 user_id
1269 .ok_or_else(|| anyhow!("missing user_id parameter"))?,
1270 access_token.ok_or_else(|| {
1271 anyhow!("missing access_token parameter")
1272 })?,
1273 ));
1274 }
1275 }
1276
1277 Err(anyhow!("didn't receive login redirect"))
1278 })
1279 .await?;
1280
1281 let access_token = private_key
1282 .decrypt_string(&access_token)
1283 .context("failed to decrypt access token")?;
1284
1285 Ok(Credentials::User {
1286 user_id: user_id.parse()?,
1287 access_token,
1288 })
1289 })
1290 .await?;
1291
1292 cx.update(|cx| cx.activate(true))?;
1293 Ok(credentials)
1294 })
1295 }
1296
1297 async fn authenticate_as_admin(
1298 http: Arc<HttpClientWithUrl>,
1299 login: String,
1300 mut api_token: String,
1301 ) -> Result<Credentials> {
1302 #[derive(Deserialize)]
1303 struct AuthenticatedUserResponse {
1304 user: User,
1305 }
1306
1307 #[derive(Deserialize)]
1308 struct User {
1309 id: u64,
1310 }
1311
1312 // Use the collab server's admin API to retrieve the id
1313 // of the impersonated user.
1314 let mut url = Self::get_rpc_url(http.clone(), None).await?;
1315 url.set_path("/user");
1316 url.set_query(Some(&format!("github_login={login}")));
1317 let request = Request::get(url.as_str())
1318 .header("Authorization", format!("token {api_token}"))
1319 .body("".into())?;
1320
1321 let mut response = http.send(request).await?;
1322 let mut body = String::new();
1323 response.body_mut().read_to_string(&mut body).await?;
1324 if !response.status().is_success() {
1325 Err(anyhow!(
1326 "admin user request failed {} - {}",
1327 response.status().as_u16(),
1328 body,
1329 ))?;
1330 }
1331 let response: AuthenticatedUserResponse = serde_json::from_str(&body)?;
1332
1333 // Use the admin API token to authenticate as the impersonated user.
1334 api_token.insert_str(0, "ADMIN_TOKEN:");
1335 Ok(Credentials::User {
1336 user_id: response.user.id,
1337 access_token: api_token,
1338 })
1339 }
1340
1341 pub async fn sign_out(self: &Arc<Self>, cx: &AsyncAppContext) {
1342 self.state.write().credentials = None;
1343 self.disconnect(&cx);
1344
1345 if self.has_credentials(cx).await {
1346 self.credentials_provider
1347 .delete_credentials(cx)
1348 .await
1349 .log_err();
1350 }
1351 }
1352
1353 pub fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) {
1354 self.peer.teardown();
1355 self.set_status(Status::SignedOut, cx);
1356 }
1357
1358 pub fn reconnect(self: &Arc<Self>, cx: &AsyncAppContext) {
1359 self.peer.teardown();
1360 self.set_status(Status::ConnectionLost, cx);
1361 }
1362
1363 fn connection_id(&self) -> Result<ConnectionId> {
1364 if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1365 Ok(connection_id)
1366 } else {
1367 Err(anyhow!("not connected"))
1368 }
1369 }
1370
1371 pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1372 log::debug!("rpc send. client_id:{}, name:{}", self.id(), T::NAME);
1373 self.peer.send(self.connection_id()?, message)
1374 }
1375
1376 pub fn request<T: RequestMessage>(
1377 &self,
1378 request: T,
1379 ) -> impl Future<Output = Result<T::Response>> {
1380 self.request_envelope(request)
1381 .map_ok(|envelope| envelope.payload)
1382 }
1383
1384 pub fn request_stream<T: RequestMessage>(
1385 &self,
1386 request: T,
1387 ) -> impl Future<Output = Result<impl Stream<Item = Result<T::Response>>>> {
1388 let client_id = self.id.load(Ordering::SeqCst);
1389 log::debug!(
1390 "rpc request start. client_id:{}. name:{}",
1391 client_id,
1392 T::NAME
1393 );
1394 let response = self
1395 .connection_id()
1396 .map(|conn_id| self.peer.request_stream(conn_id, request));
1397 async move {
1398 let response = response?.await;
1399 log::debug!(
1400 "rpc request finish. client_id:{}. name:{}",
1401 client_id,
1402 T::NAME
1403 );
1404 response
1405 }
1406 }
1407
1408 pub fn request_envelope<T: RequestMessage>(
1409 &self,
1410 request: T,
1411 ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> {
1412 let client_id = self.id();
1413 log::debug!(
1414 "rpc request start. client_id:{}. name:{}",
1415 client_id,
1416 T::NAME
1417 );
1418 let response = self
1419 .connection_id()
1420 .map(|conn_id| self.peer.request_envelope(conn_id, request));
1421 async move {
1422 let response = response?.await;
1423 log::debug!(
1424 "rpc request finish. client_id:{}. name:{}",
1425 client_id,
1426 T::NAME
1427 );
1428 response
1429 }
1430 }
1431
1432 pub fn request_dynamic(
1433 &self,
1434 envelope: proto::Envelope,
1435 request_type: &'static str,
1436 ) -> impl Future<Output = Result<proto::Envelope>> {
1437 let client_id = self.id();
1438 log::debug!(
1439 "rpc request start. client_id:{}. name:{}",
1440 client_id,
1441 request_type
1442 );
1443 let response = self
1444 .connection_id()
1445 .map(|conn_id| self.peer.request_dynamic(conn_id, envelope, request_type));
1446 async move {
1447 let response = response?.await;
1448 log::debug!(
1449 "rpc request finish. client_id:{}. name:{}",
1450 client_id,
1451 request_type
1452 );
1453 Ok(response?.0)
1454 }
1455 }
1456
1457 fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
1458 log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
1459 self.peer.respond(receipt, response)
1460 }
1461
1462 fn respond_with_error<T: RequestMessage>(
1463 &self,
1464 receipt: Receipt<T>,
1465 error: proto::Error,
1466 ) -> Result<()> {
1467 log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
1468 self.peer.respond_with_error(receipt, error)
1469 }
1470
1471 fn handle_message(
1472 self: &Arc<Client>,
1473 message: Box<dyn AnyTypedEnvelope>,
1474 cx: &AsyncAppContext,
1475 ) {
1476 let mut state = self.state.write();
1477 let type_name = message.payload_type_name();
1478 let payload_type_id = message.payload_type_id();
1479 let sender_id = message.original_sender_id();
1480
1481 let mut subscriber = None;
1482
1483 if let Some(handle) = state
1484 .models_by_message_type
1485 .get(&payload_type_id)
1486 .and_then(|handle| handle.upgrade())
1487 {
1488 subscriber = Some(handle);
1489 } else if let Some((extract_entity_id, entity_type_id)) =
1490 state.entity_id_extractors.get(&payload_type_id).zip(
1491 state
1492 .entity_types_by_message_type
1493 .get(&payload_type_id)
1494 .copied(),
1495 )
1496 {
1497 let entity_id = (extract_entity_id)(message.as_ref());
1498
1499 match state
1500 .entities_by_type_and_remote_id
1501 .get_mut(&(entity_type_id, entity_id))
1502 {
1503 Some(WeakSubscriber::Pending(pending)) => {
1504 pending.push(message);
1505 return;
1506 }
1507 Some(weak_subscriber) => match weak_subscriber {
1508 WeakSubscriber::Entity { handle } => {
1509 subscriber = handle.upgrade();
1510 }
1511
1512 WeakSubscriber::Pending(_) => {}
1513 },
1514 _ => {}
1515 }
1516 }
1517
1518 let subscriber = if let Some(subscriber) = subscriber {
1519 subscriber
1520 } else {
1521 log::info!("unhandled message {}", type_name);
1522 self.peer.respond_with_unhandled_message(message).log_err();
1523 return;
1524 };
1525
1526 let handler = state.message_handlers.get(&payload_type_id).cloned();
1527 // Dropping the state prevents deadlocks if the handler interacts with rpc::Client.
1528 // It also ensures we don't hold the lock while yielding back to the executor, as
1529 // that might cause the executor thread driving this future to block indefinitely.
1530 drop(state);
1531
1532 if let Some(handler) = handler {
1533 let future = handler(subscriber, message, self, cx.clone());
1534 let client_id = self.id();
1535 log::debug!(
1536 "rpc message received. client_id:{}, sender_id:{:?}, type:{}",
1537 client_id,
1538 sender_id,
1539 type_name
1540 );
1541 cx.spawn(move |_| async move {
1542 match future.await {
1543 Ok(()) => {
1544 log::debug!(
1545 "rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
1546 client_id,
1547 sender_id,
1548 type_name
1549 );
1550 }
1551 Err(error) => {
1552 log::error!(
1553 "error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
1554 client_id,
1555 sender_id,
1556 type_name,
1557 error
1558 );
1559 }
1560 }
1561 })
1562 .detach();
1563 } else {
1564 log::info!("unhandled message {}", type_name);
1565 self.peer.respond_with_unhandled_message(message).log_err();
1566 }
1567 }
1568
1569 pub fn telemetry(&self) -> &Arc<Telemetry> {
1570 &self.telemetry
1571 }
1572}
1573
1574#[derive(Serialize, Deserialize)]
1575struct DevelopmentCredentials {
1576 user_id: u64,
1577 access_token: String,
1578}
1579
1580/// A credentials provider that stores credentials in a local file.
1581///
1582/// This MUST only be used in development, as this is not a secure way of storing
1583/// credentials on user machines.
1584///
1585/// Its existence is purely to work around the annoyance of having to constantly
1586/// re-allow access to the system keychain when developing Zed.
1587struct DevelopmentCredentialsProvider {
1588 path: PathBuf,
1589}
1590
1591impl CredentialsProvider for DevelopmentCredentialsProvider {
1592 fn read_credentials<'a>(
1593 &'a self,
1594 _cx: &'a AsyncAppContext,
1595 ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {
1596 async move {
1597 if IMPERSONATE_LOGIN.is_some() {
1598 return None;
1599 }
1600
1601 let json = std::fs::read(&self.path).log_err()?;
1602
1603 let credentials: DevelopmentCredentials = serde_json::from_slice(&json).log_err()?;
1604
1605 Some(Credentials::User {
1606 user_id: credentials.user_id,
1607 access_token: credentials.access_token,
1608 })
1609 }
1610 .boxed_local()
1611 }
1612
1613 fn write_credentials<'a>(
1614 &'a self,
1615 user_id: u64,
1616 access_token: String,
1617 _cx: &'a AsyncAppContext,
1618 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1619 async move {
1620 let json = serde_json::to_string(&DevelopmentCredentials {
1621 user_id,
1622 access_token,
1623 })?;
1624
1625 std::fs::write(&self.path, json)?;
1626
1627 Ok(())
1628 }
1629 .boxed_local()
1630 }
1631
1632 fn delete_credentials<'a>(
1633 &'a self,
1634 _cx: &'a AsyncAppContext,
1635 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1636 async move { Ok(std::fs::remove_file(&self.path)?) }.boxed_local()
1637 }
1638}
1639
1640/// A credentials provider that stores credentials in the system keychain.
1641struct KeychainCredentialsProvider;
1642
1643impl CredentialsProvider for KeychainCredentialsProvider {
1644 fn read_credentials<'a>(
1645 &'a self,
1646 cx: &'a AsyncAppContext,
1647 ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {
1648 async move {
1649 if IMPERSONATE_LOGIN.is_some() {
1650 return None;
1651 }
1652
1653 let (user_id, access_token) = cx
1654 .update(|cx| cx.read_credentials(&ClientSettings::get_global(cx).server_url))
1655 .log_err()?
1656 .await
1657 .log_err()??;
1658
1659 Some(Credentials::User {
1660 user_id: user_id.parse().ok()?,
1661 access_token: String::from_utf8(access_token).ok()?,
1662 })
1663 }
1664 .boxed_local()
1665 }
1666
1667 fn write_credentials<'a>(
1668 &'a self,
1669 user_id: u64,
1670 access_token: String,
1671 cx: &'a AsyncAppContext,
1672 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1673 async move {
1674 cx.update(move |cx| {
1675 cx.write_credentials(
1676 &ClientSettings::get_global(cx).server_url,
1677 &user_id.to_string(),
1678 access_token.as_bytes(),
1679 )
1680 })?
1681 .await
1682 }
1683 .boxed_local()
1684 }
1685
1686 fn delete_credentials<'a>(
1687 &'a self,
1688 cx: &'a AsyncAppContext,
1689 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1690 async move {
1691 cx.update(move |cx| cx.delete_credentials(&ClientSettings::get_global(cx).server_url))?
1692 .await
1693 }
1694 .boxed_local()
1695 }
1696}
1697
1698/// prefix for the zed:// url scheme
1699pub static ZED_URL_SCHEME: &str = "zed";
1700
1701/// Parses the given link into a Zed link.
1702///
1703/// Returns a [`Some`] containing the unprefixed link if the link is a Zed link.
1704/// Returns [`None`] otherwise.
1705pub fn parse_zed_link<'a>(link: &'a str, cx: &AppContext) -> Option<&'a str> {
1706 let server_url = &ClientSettings::get_global(cx).server_url;
1707 if let Some(stripped) = link
1708 .strip_prefix(server_url)
1709 .and_then(|result| result.strip_prefix('/'))
1710 {
1711 return Some(stripped);
1712 }
1713 if let Some(stripped) = link
1714 .strip_prefix(ZED_URL_SCHEME)
1715 .and_then(|result| result.strip_prefix("://"))
1716 {
1717 return Some(stripped);
1718 }
1719
1720 None
1721}
1722
1723#[cfg(test)]
1724mod tests {
1725 use super::*;
1726 use crate::test::FakeServer;
1727
1728 use clock::FakeSystemClock;
1729 use gpui::{BackgroundExecutor, Context, TestAppContext};
1730 use http::FakeHttpClient;
1731 use parking_lot::Mutex;
1732 use proto::TypedEnvelope;
1733 use settings::SettingsStore;
1734 use std::future;
1735
1736 #[gpui::test(iterations = 10)]
1737 async fn test_reconnection(cx: &mut TestAppContext) {
1738 init_test(cx);
1739 let user_id = 5;
1740 let client = cx.update(|cx| {
1741 Client::new(
1742 Arc::new(FakeSystemClock::default()),
1743 FakeHttpClient::with_404_response(),
1744 cx,
1745 )
1746 });
1747 let server = FakeServer::for_client(user_id, &client, cx).await;
1748 let mut status = client.status();
1749 assert!(matches!(
1750 status.next().await,
1751 Some(Status::Connected { .. })
1752 ));
1753 assert_eq!(server.auth_count(), 1);
1754
1755 server.forbid_connections();
1756 server.disconnect();
1757 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1758
1759 server.allow_connections();
1760 cx.executor().advance_clock(Duration::from_secs(10));
1761 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1762 assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1763
1764 server.forbid_connections();
1765 server.disconnect();
1766 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1767
1768 // Clear cached credentials after authentication fails
1769 server.roll_access_token();
1770 server.allow_connections();
1771 cx.executor().run_until_parked();
1772 cx.executor().advance_clock(Duration::from_secs(10));
1773 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1774 assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1775 }
1776
1777 #[gpui::test(iterations = 10)]
1778 async fn test_connection_timeout(executor: BackgroundExecutor, cx: &mut TestAppContext) {
1779 init_test(cx);
1780 let user_id = 5;
1781 let client = cx.update(|cx| {
1782 Client::new(
1783 Arc::new(FakeSystemClock::default()),
1784 FakeHttpClient::with_404_response(),
1785 cx,
1786 )
1787 });
1788 let mut status = client.status();
1789
1790 // Time out when client tries to connect.
1791 client.override_authenticate(move |cx| {
1792 cx.background_executor().spawn(async move {
1793 Ok(Credentials::User {
1794 user_id,
1795 access_token: "token".into(),
1796 })
1797 })
1798 });
1799 client.override_establish_connection(|_, cx| {
1800 cx.background_executor().spawn(async move {
1801 future::pending::<()>().await;
1802 unreachable!()
1803 })
1804 });
1805 let auth_and_connect = cx.spawn({
1806 let client = client.clone();
1807 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1808 });
1809 executor.run_until_parked();
1810 assert!(matches!(status.next().await, Some(Status::Connecting)));
1811
1812 executor.advance_clock(CONNECTION_TIMEOUT);
1813 assert!(matches!(
1814 status.next().await,
1815 Some(Status::ConnectionError { .. })
1816 ));
1817 auth_and_connect.await.unwrap_err();
1818
1819 // Allow the connection to be established.
1820 let server = FakeServer::for_client(user_id, &client, cx).await;
1821 assert!(matches!(
1822 status.next().await,
1823 Some(Status::Connected { .. })
1824 ));
1825
1826 // Disconnect client.
1827 server.forbid_connections();
1828 server.disconnect();
1829 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1830
1831 // Time out when re-establishing the connection.
1832 server.allow_connections();
1833 client.override_establish_connection(|_, cx| {
1834 cx.background_executor().spawn(async move {
1835 future::pending::<()>().await;
1836 unreachable!()
1837 })
1838 });
1839 executor.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1840 assert!(matches!(
1841 status.next().await,
1842 Some(Status::Reconnecting { .. })
1843 ));
1844
1845 executor.advance_clock(CONNECTION_TIMEOUT);
1846 assert!(matches!(
1847 status.next().await,
1848 Some(Status::ReconnectionError { .. })
1849 ));
1850 }
1851
1852 #[gpui::test(iterations = 10)]
1853 async fn test_authenticating_more_than_once(
1854 cx: &mut TestAppContext,
1855 executor: BackgroundExecutor,
1856 ) {
1857 init_test(cx);
1858 let auth_count = Arc::new(Mutex::new(0));
1859 let dropped_auth_count = Arc::new(Mutex::new(0));
1860 let client = cx.update(|cx| {
1861 Client::new(
1862 Arc::new(FakeSystemClock::default()),
1863 FakeHttpClient::with_404_response(),
1864 cx,
1865 )
1866 });
1867 client.override_authenticate({
1868 let auth_count = auth_count.clone();
1869 let dropped_auth_count = dropped_auth_count.clone();
1870 move |cx| {
1871 let auth_count = auth_count.clone();
1872 let dropped_auth_count = dropped_auth_count.clone();
1873 cx.background_executor().spawn(async move {
1874 *auth_count.lock() += 1;
1875 let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
1876 future::pending::<()>().await;
1877 unreachable!()
1878 })
1879 }
1880 });
1881
1882 let _authenticate = cx.spawn({
1883 let client = client.clone();
1884 move |cx| async move { client.authenticate_and_connect(false, &cx).await }
1885 });
1886 executor.run_until_parked();
1887 assert_eq!(*auth_count.lock(), 1);
1888 assert_eq!(*dropped_auth_count.lock(), 0);
1889
1890 let _authenticate = cx.spawn({
1891 let client = client.clone();
1892 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1893 });
1894 executor.run_until_parked();
1895 assert_eq!(*auth_count.lock(), 2);
1896 assert_eq!(*dropped_auth_count.lock(), 1);
1897 }
1898
1899 #[gpui::test]
1900 async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
1901 init_test(cx);
1902 let user_id = 5;
1903 let client = cx.update(|cx| {
1904 Client::new(
1905 Arc::new(FakeSystemClock::default()),
1906 FakeHttpClient::with_404_response(),
1907 cx,
1908 )
1909 });
1910 let server = FakeServer::for_client(user_id, &client, cx).await;
1911
1912 let (done_tx1, mut done_rx1) = smol::channel::unbounded();
1913 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1914 client.add_model_message_handler(
1915 move |model: Model<TestModel>, _: TypedEnvelope<proto::JoinProject>, _, mut cx| {
1916 match model.update(&mut cx, |model, _| model.id).unwrap() {
1917 1 => done_tx1.try_send(()).unwrap(),
1918 2 => done_tx2.try_send(()).unwrap(),
1919 _ => unreachable!(),
1920 }
1921 async { Ok(()) }
1922 },
1923 );
1924 let model1 = cx.new_model(|_| TestModel {
1925 id: 1,
1926 subscription: None,
1927 });
1928 let model2 = cx.new_model(|_| TestModel {
1929 id: 2,
1930 subscription: None,
1931 });
1932 let model3 = cx.new_model(|_| TestModel {
1933 id: 3,
1934 subscription: None,
1935 });
1936
1937 let _subscription1 = client
1938 .subscribe_to_entity(1)
1939 .unwrap()
1940 .set_model(&model1, &mut cx.to_async());
1941 let _subscription2 = client
1942 .subscribe_to_entity(2)
1943 .unwrap()
1944 .set_model(&model2, &mut cx.to_async());
1945 // Ensure dropping a subscription for the same entity type still allows receiving of
1946 // messages for other entity IDs of the same type.
1947 let subscription3 = client
1948 .subscribe_to_entity(3)
1949 .unwrap()
1950 .set_model(&model3, &mut cx.to_async());
1951 drop(subscription3);
1952
1953 server.send(proto::JoinProject { project_id: 1 });
1954 server.send(proto::JoinProject { project_id: 2 });
1955 done_rx1.next().await.unwrap();
1956 done_rx2.next().await.unwrap();
1957 }
1958
1959 #[gpui::test]
1960 async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
1961 init_test(cx);
1962 let user_id = 5;
1963 let client = cx.update(|cx| {
1964 Client::new(
1965 Arc::new(FakeSystemClock::default()),
1966 FakeHttpClient::with_404_response(),
1967 cx,
1968 )
1969 });
1970 let server = FakeServer::for_client(user_id, &client, cx).await;
1971
1972 let model = cx.new_model(|_| TestModel::default());
1973 let (done_tx1, _done_rx1) = smol::channel::unbounded();
1974 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1975 let subscription1 = client.add_message_handler(
1976 model.downgrade(),
1977 move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1978 done_tx1.try_send(()).unwrap();
1979 async { Ok(()) }
1980 },
1981 );
1982 drop(subscription1);
1983 let _subscription2 = client.add_message_handler(
1984 model.downgrade(),
1985 move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1986 done_tx2.try_send(()).unwrap();
1987 async { Ok(()) }
1988 },
1989 );
1990 server.send(proto::Ping {});
1991 done_rx2.next().await.unwrap();
1992 }
1993
1994 #[gpui::test]
1995 async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
1996 init_test(cx);
1997 let user_id = 5;
1998 let client = cx.update(|cx| {
1999 Client::new(
2000 Arc::new(FakeSystemClock::default()),
2001 FakeHttpClient::with_404_response(),
2002 cx,
2003 )
2004 });
2005 let server = FakeServer::for_client(user_id, &client, cx).await;
2006
2007 let model = cx.new_model(|_| TestModel::default());
2008 let (done_tx, mut done_rx) = smol::channel::unbounded();
2009 let subscription = client.add_message_handler(
2010 model.clone().downgrade(),
2011 move |model: Model<TestModel>, _: TypedEnvelope<proto::Ping>, _, mut cx| {
2012 model
2013 .update(&mut cx, |model, _| model.subscription.take())
2014 .unwrap();
2015 done_tx.try_send(()).unwrap();
2016 async { Ok(()) }
2017 },
2018 );
2019 model.update(cx, |model, _| {
2020 model.subscription = Some(subscription);
2021 });
2022 server.send(proto::Ping {});
2023 done_rx.next().await.unwrap();
2024 }
2025
2026 #[derive(Default)]
2027 struct TestModel {
2028 id: usize,
2029 subscription: Option<Subscription>,
2030 }
2031
2032 fn init_test(cx: &mut TestAppContext) {
2033 cx.update(|cx| {
2034 let settings_store = SettingsStore::test(cx);
2035 cx.set_global(settings_store);
2036 init_settings(cx);
2037 });
2038 }
2039}