1#[cfg(any(test, feature = "test-support"))]
2pub mod test;
3
4mod socks;
5pub mod telemetry;
6pub mod user;
7
8use anyhow::{anyhow, 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 // Use the collab server's admin API to retrieve the id
1399 // of the impersonated user.
1400 let mut url = self.rpc_url(http.clone(), None).await?;
1401 url.set_path("/user");
1402 url.set_query(Some(&format!("github_login={login}")));
1403 let request: http_client::Request<AsyncBody> = Request::get(url.as_str())
1404 .header("Authorization", format!("token {api_token}"))
1405 .body("".into())?;
1406
1407 let mut response = http.send(request).await?;
1408 let mut body = String::new();
1409 response.body_mut().read_to_string(&mut body).await?;
1410 if !response.status().is_success() {
1411 Err(anyhow!(
1412 "admin user request failed {} - {}",
1413 response.status().as_u16(),
1414 body,
1415 ))?;
1416 }
1417 let response: AuthenticatedUserResponse = serde_json::from_str(&body)?;
1418
1419 // Use the admin API token to authenticate as the impersonated user.
1420 api_token.insert_str(0, "ADMIN_TOKEN:");
1421 Ok(Credentials::User {
1422 user_id: response.user.id,
1423 access_token: api_token,
1424 })
1425 }
1426
1427 pub async fn sign_out(self: &Arc<Self>, cx: &AsyncAppContext) {
1428 self.state.write().credentials = None;
1429 self.disconnect(&cx);
1430
1431 if self.has_credentials(cx).await {
1432 self.credentials_provider
1433 .delete_credentials(cx)
1434 .await
1435 .log_err();
1436 }
1437 }
1438
1439 pub fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) {
1440 self.peer.teardown();
1441 self.set_status(Status::SignedOut, cx);
1442 }
1443
1444 pub fn reconnect(self: &Arc<Self>, cx: &AsyncAppContext) {
1445 self.peer.teardown();
1446 self.set_status(Status::ConnectionLost, cx);
1447 }
1448
1449 fn connection_id(&self) -> Result<ConnectionId> {
1450 if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1451 Ok(connection_id)
1452 } else {
1453 Err(anyhow!("not connected"))
1454 }
1455 }
1456
1457 pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1458 log::debug!("rpc send. client_id:{}, name:{}", self.id(), T::NAME);
1459 self.peer.send(self.connection_id()?, message)
1460 }
1461
1462 pub fn send_dynamic(&self, envelope: proto::Envelope) -> Result<()> {
1463 let connection_id = self.connection_id()?;
1464 self.peer.send_dynamic(connection_id, envelope)
1465 }
1466
1467 pub fn request<T: RequestMessage>(
1468 &self,
1469 request: T,
1470 ) -> impl Future<Output = Result<T::Response>> {
1471 self.request_envelope(request)
1472 .map_ok(|envelope| envelope.payload)
1473 }
1474
1475 pub fn request_stream<T: RequestMessage>(
1476 &self,
1477 request: T,
1478 ) -> impl Future<Output = Result<impl Stream<Item = Result<T::Response>>>> {
1479 let client_id = self.id.load(Ordering::SeqCst);
1480 log::debug!(
1481 "rpc request start. client_id:{}. name:{}",
1482 client_id,
1483 T::NAME
1484 );
1485 let response = self
1486 .connection_id()
1487 .map(|conn_id| self.peer.request_stream(conn_id, request));
1488 async move {
1489 let response = response?.await;
1490 log::debug!(
1491 "rpc request finish. client_id:{}. name:{}",
1492 client_id,
1493 T::NAME
1494 );
1495 response
1496 }
1497 }
1498
1499 pub fn request_envelope<T: RequestMessage>(
1500 &self,
1501 request: T,
1502 ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> {
1503 let client_id = self.id();
1504 log::debug!(
1505 "rpc request start. client_id:{}. name:{}",
1506 client_id,
1507 T::NAME
1508 );
1509 let response = self
1510 .connection_id()
1511 .map(|conn_id| self.peer.request_envelope(conn_id, request));
1512 async move {
1513 let response = response?.await;
1514 log::debug!(
1515 "rpc request finish. client_id:{}. name:{}",
1516 client_id,
1517 T::NAME
1518 );
1519 response
1520 }
1521 }
1522
1523 pub fn request_dynamic(
1524 &self,
1525 envelope: proto::Envelope,
1526 request_type: &'static str,
1527 ) -> impl Future<Output = Result<proto::Envelope>> {
1528 let client_id = self.id();
1529 log::debug!(
1530 "rpc request start. client_id:{}. name:{}",
1531 client_id,
1532 request_type
1533 );
1534 let response = self
1535 .connection_id()
1536 .map(|conn_id| self.peer.request_dynamic(conn_id, envelope, request_type));
1537 async move {
1538 let response = response?.await;
1539 log::debug!(
1540 "rpc request finish. client_id:{}. name:{}",
1541 client_id,
1542 request_type
1543 );
1544 Ok(response?.0)
1545 }
1546 }
1547
1548 fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
1549 log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
1550 self.peer.respond(receipt, response)
1551 }
1552
1553 fn respond_with_error<T: RequestMessage>(
1554 &self,
1555 receipt: Receipt<T>,
1556 error: proto::Error,
1557 ) -> Result<()> {
1558 log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
1559 self.peer.respond_with_error(receipt, error)
1560 }
1561
1562 fn handle_message(
1563 self: &Arc<Client>,
1564 message: Box<dyn AnyTypedEnvelope>,
1565 cx: &AsyncAppContext,
1566 ) {
1567 let mut state = self.state.write();
1568 let type_name = message.payload_type_name();
1569 let payload_type_id = message.payload_type_id();
1570 let sender_id = message.original_sender_id();
1571
1572 let mut subscriber = None;
1573
1574 if let Some(handle) = state
1575 .models_by_message_type
1576 .get(&payload_type_id)
1577 .and_then(|handle| handle.upgrade())
1578 {
1579 subscriber = Some(handle);
1580 } else if let Some((extract_entity_id, entity_type_id)) =
1581 state.entity_id_extractors.get(&payload_type_id).zip(
1582 state
1583 .entity_types_by_message_type
1584 .get(&payload_type_id)
1585 .copied(),
1586 )
1587 {
1588 let entity_id = (extract_entity_id)(message.as_ref());
1589
1590 match state
1591 .entities_by_type_and_remote_id
1592 .get_mut(&(entity_type_id, entity_id))
1593 {
1594 Some(WeakSubscriber::Pending(pending)) => {
1595 pending.push(message);
1596 return;
1597 }
1598 Some(weak_subscriber) => match weak_subscriber {
1599 WeakSubscriber::Entity { handle } => {
1600 subscriber = handle.upgrade();
1601 }
1602
1603 WeakSubscriber::Pending(_) => {}
1604 },
1605 _ => {}
1606 }
1607 }
1608
1609 let subscriber = if let Some(subscriber) = subscriber {
1610 subscriber
1611 } else {
1612 log::info!("unhandled message {}", type_name);
1613 self.peer.respond_with_unhandled_message(message).log_err();
1614 return;
1615 };
1616
1617 let handler = state.message_handlers.get(&payload_type_id).cloned();
1618 // Dropping the state prevents deadlocks if the handler interacts with rpc::Client.
1619 // It also ensures we don't hold the lock while yielding back to the executor, as
1620 // that might cause the executor thread driving this future to block indefinitely.
1621 drop(state);
1622
1623 if let Some(handler) = handler {
1624 let future = handler(subscriber, message, self, cx.clone());
1625 let client_id = self.id();
1626 log::debug!(
1627 "rpc message received. client_id:{}, sender_id:{:?}, type:{}",
1628 client_id,
1629 sender_id,
1630 type_name
1631 );
1632 cx.spawn(move |_| async move {
1633 match future.await {
1634 Ok(()) => {
1635 log::debug!(
1636 "rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
1637 client_id,
1638 sender_id,
1639 type_name
1640 );
1641 }
1642 Err(error) => {
1643 log::error!(
1644 "error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
1645 client_id,
1646 sender_id,
1647 type_name,
1648 error
1649 );
1650 }
1651 }
1652 })
1653 .detach();
1654 } else {
1655 log::info!("unhandled message {}", type_name);
1656 self.peer.respond_with_unhandled_message(message).log_err();
1657 }
1658 }
1659
1660 pub fn telemetry(&self) -> &Arc<Telemetry> {
1661 &self.telemetry
1662 }
1663}
1664
1665impl ProtoClient for Client {
1666 fn request(
1667 &self,
1668 envelope: proto::Envelope,
1669 request_type: &'static str,
1670 ) -> BoxFuture<'static, Result<proto::Envelope>> {
1671 self.request_dynamic(envelope, request_type).boxed()
1672 }
1673
1674 fn send(&self, envelope: proto::Envelope) -> Result<()> {
1675 self.send_dynamic(envelope)
1676 }
1677}
1678
1679#[derive(Serialize, Deserialize)]
1680struct DevelopmentCredentials {
1681 user_id: u64,
1682 access_token: String,
1683}
1684
1685/// A credentials provider that stores credentials in a local file.
1686///
1687/// This MUST only be used in development, as this is not a secure way of storing
1688/// credentials on user machines.
1689///
1690/// Its existence is purely to work around the annoyance of having to constantly
1691/// re-allow access to the system keychain when developing Zed.
1692struct DevelopmentCredentialsProvider {
1693 path: PathBuf,
1694}
1695
1696impl CredentialsProvider for DevelopmentCredentialsProvider {
1697 fn read_credentials<'a>(
1698 &'a self,
1699 _cx: &'a AsyncAppContext,
1700 ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {
1701 async move {
1702 if IMPERSONATE_LOGIN.is_some() {
1703 return None;
1704 }
1705
1706 let json = std::fs::read(&self.path).log_err()?;
1707
1708 let credentials: DevelopmentCredentials = serde_json::from_slice(&json).log_err()?;
1709
1710 Some(Credentials::User {
1711 user_id: credentials.user_id,
1712 access_token: credentials.access_token,
1713 })
1714 }
1715 .boxed_local()
1716 }
1717
1718 fn write_credentials<'a>(
1719 &'a self,
1720 user_id: u64,
1721 access_token: String,
1722 _cx: &'a AsyncAppContext,
1723 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1724 async move {
1725 let json = serde_json::to_string(&DevelopmentCredentials {
1726 user_id,
1727 access_token,
1728 })?;
1729
1730 std::fs::write(&self.path, json)?;
1731
1732 Ok(())
1733 }
1734 .boxed_local()
1735 }
1736
1737 fn delete_credentials<'a>(
1738 &'a self,
1739 _cx: &'a AsyncAppContext,
1740 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1741 async move { Ok(std::fs::remove_file(&self.path)?) }.boxed_local()
1742 }
1743}
1744
1745/// A credentials provider that stores credentials in the system keychain.
1746struct KeychainCredentialsProvider;
1747
1748impl CredentialsProvider for KeychainCredentialsProvider {
1749 fn read_credentials<'a>(
1750 &'a self,
1751 cx: &'a AsyncAppContext,
1752 ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {
1753 async move {
1754 if IMPERSONATE_LOGIN.is_some() {
1755 return None;
1756 }
1757
1758 let (user_id, access_token) = cx
1759 .update(|cx| cx.read_credentials(&ClientSettings::get_global(cx).server_url))
1760 .log_err()?
1761 .await
1762 .log_err()??;
1763
1764 Some(Credentials::User {
1765 user_id: user_id.parse().ok()?,
1766 access_token: String::from_utf8(access_token).ok()?,
1767 })
1768 }
1769 .boxed_local()
1770 }
1771
1772 fn write_credentials<'a>(
1773 &'a self,
1774 user_id: u64,
1775 access_token: String,
1776 cx: &'a AsyncAppContext,
1777 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1778 async move {
1779 cx.update(move |cx| {
1780 cx.write_credentials(
1781 &ClientSettings::get_global(cx).server_url,
1782 &user_id.to_string(),
1783 access_token.as_bytes(),
1784 )
1785 })?
1786 .await
1787 }
1788 .boxed_local()
1789 }
1790
1791 fn delete_credentials<'a>(
1792 &'a self,
1793 cx: &'a AsyncAppContext,
1794 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1795 async move {
1796 cx.update(move |cx| cx.delete_credentials(&ClientSettings::get_global(cx).server_url))?
1797 .await
1798 }
1799 .boxed_local()
1800 }
1801}
1802
1803/// prefix for the zed:// url scheme
1804pub static ZED_URL_SCHEME: &str = "zed";
1805
1806/// Parses the given link into a Zed link.
1807///
1808/// Returns a [`Some`] containing the unprefixed link if the link is a Zed link.
1809/// Returns [`None`] otherwise.
1810pub fn parse_zed_link<'a>(link: &'a str, cx: &AppContext) -> Option<&'a str> {
1811 let server_url = &ClientSettings::get_global(cx).server_url;
1812 if let Some(stripped) = link
1813 .strip_prefix(server_url)
1814 .and_then(|result| result.strip_prefix('/'))
1815 {
1816 return Some(stripped);
1817 }
1818 if let Some(stripped) = link
1819 .strip_prefix(ZED_URL_SCHEME)
1820 .and_then(|result| result.strip_prefix("://"))
1821 {
1822 return Some(stripped);
1823 }
1824
1825 None
1826}
1827
1828#[cfg(test)]
1829mod tests {
1830 use super::*;
1831 use crate::test::FakeServer;
1832
1833 use clock::FakeSystemClock;
1834 use gpui::{BackgroundExecutor, Context, TestAppContext};
1835 use http_client::FakeHttpClient;
1836 use parking_lot::Mutex;
1837 use proto::TypedEnvelope;
1838 use settings::SettingsStore;
1839 use std::future;
1840
1841 #[gpui::test(iterations = 10)]
1842 async fn test_reconnection(cx: &mut TestAppContext) {
1843 init_test(cx);
1844 let user_id = 5;
1845 let client = cx.update(|cx| {
1846 Client::new(
1847 Arc::new(FakeSystemClock::default()),
1848 FakeHttpClient::with_404_response(),
1849 cx,
1850 )
1851 });
1852 let server = FakeServer::for_client(user_id, &client, cx).await;
1853 let mut status = client.status();
1854 assert!(matches!(
1855 status.next().await,
1856 Some(Status::Connected { .. })
1857 ));
1858 assert_eq!(server.auth_count(), 1);
1859
1860 server.forbid_connections();
1861 server.disconnect();
1862 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1863
1864 server.allow_connections();
1865 cx.executor().advance_clock(Duration::from_secs(10));
1866 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1867 assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1868
1869 server.forbid_connections();
1870 server.disconnect();
1871 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1872
1873 // Clear cached credentials after authentication fails
1874 server.roll_access_token();
1875 server.allow_connections();
1876 cx.executor().run_until_parked();
1877 cx.executor().advance_clock(Duration::from_secs(10));
1878 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1879 assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1880 }
1881
1882 #[gpui::test(iterations = 10)]
1883 async fn test_connection_timeout(executor: BackgroundExecutor, cx: &mut TestAppContext) {
1884 init_test(cx);
1885 let user_id = 5;
1886 let client = cx.update(|cx| {
1887 Client::new(
1888 Arc::new(FakeSystemClock::default()),
1889 FakeHttpClient::with_404_response(),
1890 cx,
1891 )
1892 });
1893 let mut status = client.status();
1894
1895 // Time out when client tries to connect.
1896 client.override_authenticate(move |cx| {
1897 cx.background_executor().spawn(async move {
1898 Ok(Credentials::User {
1899 user_id,
1900 access_token: "token".into(),
1901 })
1902 })
1903 });
1904 client.override_establish_connection(|_, cx| {
1905 cx.background_executor().spawn(async move {
1906 future::pending::<()>().await;
1907 unreachable!()
1908 })
1909 });
1910 let auth_and_connect = cx.spawn({
1911 let client = client.clone();
1912 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1913 });
1914 executor.run_until_parked();
1915 assert!(matches!(status.next().await, Some(Status::Connecting)));
1916
1917 executor.advance_clock(CONNECTION_TIMEOUT);
1918 assert!(matches!(
1919 status.next().await,
1920 Some(Status::ConnectionError { .. })
1921 ));
1922 auth_and_connect.await.unwrap_err();
1923
1924 // Allow the connection to be established.
1925 let server = FakeServer::for_client(user_id, &client, cx).await;
1926 assert!(matches!(
1927 status.next().await,
1928 Some(Status::Connected { .. })
1929 ));
1930
1931 // Disconnect client.
1932 server.forbid_connections();
1933 server.disconnect();
1934 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1935
1936 // Time out when re-establishing the connection.
1937 server.allow_connections();
1938 client.override_establish_connection(|_, cx| {
1939 cx.background_executor().spawn(async move {
1940 future::pending::<()>().await;
1941 unreachable!()
1942 })
1943 });
1944 executor.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1945 assert!(matches!(
1946 status.next().await,
1947 Some(Status::Reconnecting { .. })
1948 ));
1949
1950 executor.advance_clock(CONNECTION_TIMEOUT);
1951 assert!(matches!(
1952 status.next().await,
1953 Some(Status::ReconnectionError { .. })
1954 ));
1955 }
1956
1957 #[gpui::test(iterations = 10)]
1958 async fn test_authenticating_more_than_once(
1959 cx: &mut TestAppContext,
1960 executor: BackgroundExecutor,
1961 ) {
1962 init_test(cx);
1963 let auth_count = Arc::new(Mutex::new(0));
1964 let dropped_auth_count = Arc::new(Mutex::new(0));
1965 let client = cx.update(|cx| {
1966 Client::new(
1967 Arc::new(FakeSystemClock::default()),
1968 FakeHttpClient::with_404_response(),
1969 cx,
1970 )
1971 });
1972 client.override_authenticate({
1973 let auth_count = auth_count.clone();
1974 let dropped_auth_count = dropped_auth_count.clone();
1975 move |cx| {
1976 let auth_count = auth_count.clone();
1977 let dropped_auth_count = dropped_auth_count.clone();
1978 cx.background_executor().spawn(async move {
1979 *auth_count.lock() += 1;
1980 let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
1981 future::pending::<()>().await;
1982 unreachable!()
1983 })
1984 }
1985 });
1986
1987 let _authenticate = cx.spawn({
1988 let client = client.clone();
1989 move |cx| async move { client.authenticate_and_connect(false, &cx).await }
1990 });
1991 executor.run_until_parked();
1992 assert_eq!(*auth_count.lock(), 1);
1993 assert_eq!(*dropped_auth_count.lock(), 0);
1994
1995 let _authenticate = cx.spawn({
1996 let client = client.clone();
1997 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1998 });
1999 executor.run_until_parked();
2000 assert_eq!(*auth_count.lock(), 2);
2001 assert_eq!(*dropped_auth_count.lock(), 1);
2002 }
2003
2004 #[gpui::test]
2005 async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
2006 init_test(cx);
2007 let user_id = 5;
2008 let client = cx.update(|cx| {
2009 Client::new(
2010 Arc::new(FakeSystemClock::default()),
2011 FakeHttpClient::with_404_response(),
2012 cx,
2013 )
2014 });
2015 let server = FakeServer::for_client(user_id, &client, cx).await;
2016
2017 let (done_tx1, mut done_rx1) = smol::channel::unbounded();
2018 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
2019 client.add_model_message_handler(
2020 move |model: Model<TestModel>, _: TypedEnvelope<proto::JoinProject>, mut cx| {
2021 match model.update(&mut cx, |model, _| model.id).unwrap() {
2022 1 => done_tx1.try_send(()).unwrap(),
2023 2 => done_tx2.try_send(()).unwrap(),
2024 _ => unreachable!(),
2025 }
2026 async { Ok(()) }
2027 },
2028 );
2029 let model1 = cx.new_model(|_| TestModel {
2030 id: 1,
2031 subscription: None,
2032 });
2033 let model2 = cx.new_model(|_| TestModel {
2034 id: 2,
2035 subscription: None,
2036 });
2037 let model3 = cx.new_model(|_| TestModel {
2038 id: 3,
2039 subscription: None,
2040 });
2041
2042 let _subscription1 = client
2043 .subscribe_to_entity(1)
2044 .unwrap()
2045 .set_model(&model1, &mut cx.to_async());
2046 let _subscription2 = client
2047 .subscribe_to_entity(2)
2048 .unwrap()
2049 .set_model(&model2, &mut cx.to_async());
2050 // Ensure dropping a subscription for the same entity type still allows receiving of
2051 // messages for other entity IDs of the same type.
2052 let subscription3 = client
2053 .subscribe_to_entity(3)
2054 .unwrap()
2055 .set_model(&model3, &mut cx.to_async());
2056 drop(subscription3);
2057
2058 server.send(proto::JoinProject { project_id: 1 });
2059 server.send(proto::JoinProject { project_id: 2 });
2060 done_rx1.next().await.unwrap();
2061 done_rx2.next().await.unwrap();
2062 }
2063
2064 #[gpui::test]
2065 async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
2066 init_test(cx);
2067 let user_id = 5;
2068 let client = cx.update(|cx| {
2069 Client::new(
2070 Arc::new(FakeSystemClock::default()),
2071 FakeHttpClient::with_404_response(),
2072 cx,
2073 )
2074 });
2075 let server = FakeServer::for_client(user_id, &client, cx).await;
2076
2077 let model = cx.new_model(|_| TestModel::default());
2078 let (done_tx1, _done_rx1) = smol::channel::unbounded();
2079 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
2080 let subscription1 = client.add_message_handler(
2081 model.downgrade(),
2082 move |_, _: TypedEnvelope<proto::Ping>, _| {
2083 done_tx1.try_send(()).unwrap();
2084 async { Ok(()) }
2085 },
2086 );
2087 drop(subscription1);
2088 let _subscription2 = client.add_message_handler(
2089 model.downgrade(),
2090 move |_, _: TypedEnvelope<proto::Ping>, _| {
2091 done_tx2.try_send(()).unwrap();
2092 async { Ok(()) }
2093 },
2094 );
2095 server.send(proto::Ping {});
2096 done_rx2.next().await.unwrap();
2097 }
2098
2099 #[gpui::test]
2100 async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
2101 init_test(cx);
2102 let user_id = 5;
2103 let client = cx.update(|cx| {
2104 Client::new(
2105 Arc::new(FakeSystemClock::default()),
2106 FakeHttpClient::with_404_response(),
2107 cx,
2108 )
2109 });
2110 let server = FakeServer::for_client(user_id, &client, cx).await;
2111
2112 let model = cx.new_model(|_| TestModel::default());
2113 let (done_tx, mut done_rx) = smol::channel::unbounded();
2114 let subscription = client.add_message_handler(
2115 model.clone().downgrade(),
2116 move |model: Model<TestModel>, _: TypedEnvelope<proto::Ping>, mut cx| {
2117 model
2118 .update(&mut cx, |model, _| model.subscription.take())
2119 .unwrap();
2120 done_tx.try_send(()).unwrap();
2121 async { Ok(()) }
2122 },
2123 );
2124 model.update(cx, |model, _| {
2125 model.subscription = Some(subscription);
2126 });
2127 server.send(proto::Ping {});
2128 done_rx.next().await.unwrap();
2129 }
2130
2131 #[derive(Default)]
2132 struct TestModel {
2133 id: usize,
2134 subscription: Option<Subscription>,
2135 }
2136
2137 fn init_test(cx: &mut TestAppContext) {
2138 cx.update(|cx| {
2139 let settings_store = SettingsStore::test(cx);
2140 cx.set_global(settings_store);
2141 init_settings(cx);
2142 });
2143 }
2144}