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