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