client.rs

  1#[cfg(any(test, feature = "test-support"))]
  2pub mod test;
  3
  4pub mod channel;
  5pub mod http;
  6pub mod user;
  7
  8use anyhow::{anyhow, Context, Result};
  9use async_recursion::async_recursion;
 10use async_tungstenite::tungstenite::{
 11    error::Error as WebsocketError,
 12    http::{Request, StatusCode},
 13};
 14use gpui::{action, AsyncAppContext, Entity, ModelContext, MutableAppContext, Task};
 15use lazy_static::lazy_static;
 16use parking_lot::RwLock;
 17use postage::{prelude::Stream, watch};
 18use rand::prelude::*;
 19use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, RequestMessage};
 20use std::{
 21    any::TypeId,
 22    collections::HashMap,
 23    convert::TryFrom,
 24    fmt::Write as _,
 25    future::Future,
 26    sync::{Arc, Weak},
 27    time::{Duration, Instant},
 28};
 29use surf::Url;
 30use thiserror::Error;
 31use util::{ResultExt, TryFutureExt};
 32
 33pub use channel::*;
 34pub use rpc::*;
 35pub use user::*;
 36
 37lazy_static! {
 38    static ref COLLAB_URL: String =
 39        std::env::var("ZED_COLLAB_URL").unwrap_or("https://collab.zed.dev:443".to_string());
 40    static ref SITE_URL: String =
 41        std::env::var("ZED_SITE_URL").unwrap_or("https://zed.dev".to_string());
 42    static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE")
 43        .ok()
 44        .and_then(|s| if s.is_empty() { None } else { Some(s) });
 45}
 46
 47action!(Authenticate);
 48
 49pub fn init(rpc: Arc<Client>, cx: &mut MutableAppContext) {
 50    cx.add_global_action(move |_: &Authenticate, cx| {
 51        let rpc = rpc.clone();
 52        cx.spawn(|cx| async move { rpc.authenticate_and_connect(&cx).log_err().await })
 53            .detach();
 54    });
 55}
 56
 57pub struct Client {
 58    peer: Arc<Peer>,
 59    state: RwLock<ClientState>,
 60    authenticate:
 61        Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,
 62    establish_connection: Option<
 63        Box<
 64            dyn 'static
 65                + Send
 66                + Sync
 67                + Fn(
 68                    &Credentials,
 69                    &AsyncAppContext,
 70                ) -> Task<Result<Connection, EstablishConnectionError>>,
 71        >,
 72    >,
 73}
 74
 75#[derive(Error, Debug)]
 76pub enum EstablishConnectionError {
 77    #[error("upgrade required")]
 78    UpgradeRequired,
 79    #[error("unauthorized")]
 80    Unauthorized,
 81    #[error("{0}")]
 82    Other(#[from] anyhow::Error),
 83    #[error("{0}")]
 84    Io(#[from] std::io::Error),
 85    #[error("{0}")]
 86    Http(#[from] async_tungstenite::tungstenite::http::Error),
 87}
 88
 89impl From<WebsocketError> for EstablishConnectionError {
 90    fn from(error: WebsocketError) -> Self {
 91        if let WebsocketError::Http(response) = &error {
 92            match response.status() {
 93                StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,
 94                StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,
 95                _ => {}
 96            }
 97        }
 98        EstablishConnectionError::Other(error.into())
 99    }
100}
101
102impl EstablishConnectionError {
103    pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
104        Self::Other(error.into())
105    }
106}
107
108#[derive(Copy, Clone, Debug)]
109pub enum Status {
110    SignedOut,
111    UpgradeRequired,
112    Authenticating,
113    Connecting,
114    ConnectionError,
115    Connected { connection_id: ConnectionId },
116    ConnectionLost,
117    Reauthenticating,
118    Reconnecting,
119    ReconnectionError { next_reconnection: Instant },
120}
121
122struct ClientState {
123    credentials: Option<Credentials>,
124    status: (watch::Sender<Status>, watch::Receiver<Status>),
125    entity_id_extractors: HashMap<TypeId, Box<dyn Send + Sync + Fn(&dyn AnyTypedEnvelope) -> u64>>,
126    model_handlers: HashMap<
127        (TypeId, u64),
128        Box<dyn Send + Sync + FnMut(Box<dyn AnyTypedEnvelope>, &mut AsyncAppContext)>,
129    >,
130    _maintain_connection: Option<Task<()>>,
131    heartbeat_interval: Duration,
132}
133
134#[derive(Clone)]
135pub struct Credentials {
136    pub user_id: u64,
137    pub access_token: String,
138}
139
140impl Default for ClientState {
141    fn default() -> Self {
142        Self {
143            credentials: None,
144            status: watch::channel_with(Status::SignedOut),
145            entity_id_extractors: Default::default(),
146            model_handlers: Default::default(),
147            _maintain_connection: None,
148            heartbeat_interval: Duration::from_secs(5),
149        }
150    }
151}
152
153pub struct Subscription {
154    client: Weak<Client>,
155    id: (TypeId, u64),
156}
157
158impl Drop for Subscription {
159    fn drop(&mut self) {
160        if let Some(client) = self.client.upgrade() {
161            drop(
162                client
163                    .state
164                    .write()
165                    .model_handlers
166                    .remove(&self.id)
167                    .unwrap(),
168            );
169        }
170    }
171}
172
173impl Client {
174    pub fn new() -> Arc<Self> {
175        Arc::new(Self {
176            peer: Peer::new(),
177            state: Default::default(),
178            authenticate: None,
179            establish_connection: None,
180        })
181    }
182
183    #[cfg(any(test, feature = "test-support"))]
184    pub fn override_authenticate<F>(&mut self, authenticate: F) -> &mut Self
185    where
186        F: 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>,
187    {
188        self.authenticate = Some(Box::new(authenticate));
189        self
190    }
191
192    #[cfg(any(test, feature = "test-support"))]
193    pub fn override_establish_connection<F>(&mut self, connect: F) -> &mut Self
194    where
195        F: 'static
196            + Send
197            + Sync
198            + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,
199    {
200        self.establish_connection = Some(Box::new(connect));
201        self
202    }
203
204    pub fn user_id(&self) -> Option<u64> {
205        self.state
206            .read()
207            .credentials
208            .as_ref()
209            .map(|credentials| credentials.user_id)
210    }
211
212    pub fn status(&self) -> watch::Receiver<Status> {
213        self.state.read().status.1.clone()
214    }
215
216    fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncAppContext) {
217        let mut state = self.state.write();
218        *state.status.0.borrow_mut() = status;
219
220        match status {
221            Status::Connected { .. } => {
222                let heartbeat_interval = state.heartbeat_interval;
223                let this = self.clone();
224                let foreground = cx.foreground();
225                state._maintain_connection = Some(cx.foreground().spawn(async move {
226                    loop {
227                        foreground.timer(heartbeat_interval).await;
228                        let _ = this.request(proto::Ping {}).await;
229                    }
230                }));
231            }
232            Status::ConnectionLost => {
233                let this = self.clone();
234                let foreground = cx.foreground();
235                let heartbeat_interval = state.heartbeat_interval;
236                state._maintain_connection = Some(cx.spawn(|cx| async move {
237                    let mut rng = StdRng::from_entropy();
238                    let mut delay = Duration::from_millis(100);
239                    while let Err(error) = this.authenticate_and_connect(&cx).await {
240                        log::error!("failed to connect {}", error);
241                        this.set_status(
242                            Status::ReconnectionError {
243                                next_reconnection: Instant::now() + delay,
244                            },
245                            &cx,
246                        );
247                        foreground.timer(delay).await;
248                        delay = delay
249                            .mul_f32(rng.gen_range(1.0..=2.0))
250                            .min(heartbeat_interval);
251                    }
252                }));
253            }
254            Status::SignedOut | Status::UpgradeRequired => {
255                state._maintain_connection.take();
256            }
257            _ => {}
258        }
259    }
260
261    pub fn subscribe<T, M, F>(
262        self: &Arc<Self>,
263        cx: &mut ModelContext<M>,
264        mut handler: F,
265    ) -> Subscription
266    where
267        T: EnvelopedMessage,
268        M: Entity,
269        F: 'static
270            + Send
271            + Sync
272            + FnMut(&mut M, TypedEnvelope<T>, Arc<Self>, &mut ModelContext<M>) -> Result<()>,
273    {
274        let subscription_id = (TypeId::of::<T>(), Default::default());
275        let client = self.clone();
276        let mut state = self.state.write();
277        let model = cx.weak_handle();
278        let prev_extractor = state
279            .entity_id_extractors
280            .insert(subscription_id.0, Box::new(|_| Default::default()));
281        if prev_extractor.is_some() {
282            panic!("registered a handler for the same entity twice")
283        }
284
285        state.model_handlers.insert(
286            subscription_id,
287            Box::new(move |envelope, cx| {
288                if let Some(model) = model.upgrade(cx) {
289                    let envelope = envelope.into_any().downcast::<TypedEnvelope<T>>().unwrap();
290                    model.update(cx, |model, cx| {
291                        if let Err(error) = handler(model, *envelope, client.clone(), cx) {
292                            log::error!("error handling message: {}", error)
293                        }
294                    });
295                }
296            }),
297        );
298
299        Subscription {
300            client: Arc::downgrade(self),
301            id: subscription_id,
302        }
303    }
304
305    pub fn subscribe_to_entity<T, M, F>(
306        self: &Arc<Self>,
307        remote_id: u64,
308        cx: &mut ModelContext<M>,
309        mut handler: F,
310    ) -> Subscription
311    where
312        T: EntityMessage,
313        M: Entity,
314        F: 'static
315            + Send
316            + Sync
317            + FnMut(&mut M, TypedEnvelope<T>, Arc<Self>, &mut ModelContext<M>) -> Result<()>,
318    {
319        let subscription_id = (TypeId::of::<T>(), remote_id);
320        let client = self.clone();
321        let mut state = self.state.write();
322        let model = cx.weak_handle();
323        state
324            .entity_id_extractors
325            .entry(subscription_id.0)
326            .or_insert_with(|| {
327                Box::new(|envelope| {
328                    let envelope = envelope
329                        .as_any()
330                        .downcast_ref::<TypedEnvelope<T>>()
331                        .unwrap();
332                    envelope.payload.remote_entity_id()
333                })
334            });
335        let prev_handler = state.model_handlers.insert(
336            subscription_id,
337            Box::new(move |envelope, cx| {
338                if let Some(model) = model.upgrade(cx) {
339                    let envelope = envelope.into_any().downcast::<TypedEnvelope<T>>().unwrap();
340                    model.update(cx, |model, cx| {
341                        if let Err(error) = handler(model, *envelope, client.clone(), cx) {
342                            log::error!("error handling message: {}", error)
343                        }
344                    });
345                }
346            }),
347        );
348        if prev_handler.is_some() {
349            panic!("registered a handler for the same entity twice")
350        }
351
352        Subscription {
353            client: Arc::downgrade(self),
354            id: subscription_id,
355        }
356    }
357
358    #[async_recursion(?Send)]
359    pub async fn authenticate_and_connect(
360        self: &Arc<Self>,
361        cx: &AsyncAppContext,
362    ) -> anyhow::Result<()> {
363        let was_disconnected = match *self.status().borrow() {
364            Status::SignedOut => true,
365            Status::ConnectionError | Status::ConnectionLost | Status::ReconnectionError { .. } => {
366                false
367            }
368            Status::Connected { .. }
369            | Status::Connecting { .. }
370            | Status::Reconnecting { .. }
371            | Status::Authenticating
372            | Status::Reauthenticating => return Ok(()),
373            Status::UpgradeRequired => return Err(EstablishConnectionError::UpgradeRequired)?,
374        };
375
376        if was_disconnected {
377            self.set_status(Status::Authenticating, cx);
378        } else {
379            self.set_status(Status::Reauthenticating, cx)
380        }
381
382        let mut used_keychain = false;
383        let credentials = self.state.read().credentials.clone();
384        let credentials = if let Some(credentials) = credentials {
385            credentials
386        } else if let Some(credentials) = read_credentials_from_keychain(cx) {
387            used_keychain = true;
388            credentials
389        } else {
390            let credentials = match self.authenticate(&cx).await {
391                Ok(credentials) => credentials,
392                Err(err) => {
393                    self.set_status(Status::ConnectionError, cx);
394                    return Err(err);
395                }
396            };
397            credentials
398        };
399
400        if was_disconnected {
401            self.set_status(Status::Connecting, cx);
402        } else {
403            self.set_status(Status::Reconnecting, cx);
404        }
405
406        match self.establish_connection(&credentials, cx).await {
407            Ok(conn) => {
408                log::info!("connected to rpc address {}", *COLLAB_URL);
409                self.state.write().credentials = Some(credentials.clone());
410                if !used_keychain && IMPERSONATE_LOGIN.is_none() {
411                    write_credentials_to_keychain(&credentials, cx).log_err();
412                }
413                self.set_connection(conn, cx).await;
414                Ok(())
415            }
416            Err(EstablishConnectionError::Unauthorized) => {
417                self.state.write().credentials.take();
418                if used_keychain {
419                    cx.platform().delete_credentials(&COLLAB_URL).log_err();
420                    self.set_status(Status::SignedOut, cx);
421                    self.authenticate_and_connect(cx).await
422                } else {
423                    self.set_status(Status::ConnectionError, cx);
424                    Err(EstablishConnectionError::Unauthorized)?
425                }
426            }
427            Err(EstablishConnectionError::UpgradeRequired) => {
428                self.set_status(Status::UpgradeRequired, cx);
429                Err(EstablishConnectionError::UpgradeRequired)?
430            }
431            Err(error) => {
432                self.set_status(Status::ConnectionError, cx);
433                Err(error)?
434            }
435        }
436    }
437
438    async fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncAppContext) {
439        let (connection_id, handle_io, mut incoming) = self.peer.add_connection(conn).await;
440        cx.foreground()
441            .spawn({
442                let mut cx = cx.clone();
443                let this = self.clone();
444                async move {
445                    while let Some(message) = incoming.recv().await {
446                        let mut state = this.state.write();
447                        if let Some(extract_entity_id) =
448                            state.entity_id_extractors.get(&message.payload_type_id())
449                        {
450                            let payload_type_id = message.payload_type_id();
451                            let entity_id = (extract_entity_id)(message.as_ref());
452                            let handler_key = (payload_type_id, entity_id);
453                            if let Some(mut handler) = state.model_handlers.remove(&handler_key) {
454                                drop(state); // Avoid deadlocks if the handler interacts with rpc::Client
455                                let start_time = Instant::now();
456                                log::info!("RPC client message {}", message.payload_type_name());
457                                (handler)(message, &mut cx);
458                                log::info!(
459                                    "RPC message handled. duration:{:?}",
460                                    start_time.elapsed()
461                                );
462                                this.state
463                                    .write()
464                                    .model_handlers
465                                    .insert(handler_key, handler);
466                            } else {
467                                log::info!("unhandled message {}", message.payload_type_name());
468                            }
469                        } else {
470                            log::info!("unhandled message {}", message.payload_type_name());
471                        }
472                    }
473                }
474            })
475            .detach();
476
477        self.set_status(Status::Connected { connection_id }, cx);
478
479        let handle_io = cx.background().spawn(handle_io);
480        let this = self.clone();
481        let cx = cx.clone();
482        cx.foreground()
483            .spawn(async move {
484                match handle_io.await {
485                    Ok(()) => this.set_status(Status::SignedOut, &cx),
486                    Err(err) => {
487                        log::error!("connection error: {:?}", err);
488                        this.set_status(Status::ConnectionLost, &cx);
489                    }
490                }
491            })
492            .detach();
493    }
494
495    fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {
496        if let Some(callback) = self.authenticate.as_ref() {
497            callback(cx)
498        } else {
499            self.authenticate_with_browser(cx)
500        }
501    }
502
503    fn establish_connection(
504        self: &Arc<Self>,
505        credentials: &Credentials,
506        cx: &AsyncAppContext,
507    ) -> Task<Result<Connection, EstablishConnectionError>> {
508        if let Some(callback) = self.establish_connection.as_ref() {
509            callback(credentials, cx)
510        } else {
511            self.establish_websocket_connection(credentials, cx)
512        }
513    }
514
515    fn establish_websocket_connection(
516        self: &Arc<Self>,
517        credentials: &Credentials,
518        cx: &AsyncAppContext,
519    ) -> Task<Result<Connection, EstablishConnectionError>> {
520        let request = Request::builder()
521            .header(
522                "Authorization",
523                format!("{} {}", credentials.user_id, credentials.access_token),
524            )
525            .header("X-Zed-Protocol-Version", rpc::PROTOCOL_VERSION);
526        cx.background().spawn(async move {
527            if let Some(host) = COLLAB_URL.strip_prefix("https://") {
528                let stream = smol::net::TcpStream::connect(host).await?;
529                let request = request.uri(format!("wss://{}/rpc", host)).body(())?;
530                let (stream, _) =
531                    async_tungstenite::async_tls::client_async_tls(request, stream).await?;
532                Ok(Connection::new(stream))
533            } else if let Some(host) = COLLAB_URL.strip_prefix("http://") {
534                let stream = smol::net::TcpStream::connect(host).await?;
535                let request = request.uri(format!("ws://{}/rpc", host)).body(())?;
536                let (stream, _) = async_tungstenite::client_async(request, stream).await?;
537                Ok(Connection::new(stream))
538            } else {
539                Err(anyhow!("invalid server url: {}", *COLLAB_URL))?
540            }
541        })
542    }
543
544    pub fn authenticate_with_browser(
545        self: &Arc<Self>,
546        cx: &AsyncAppContext,
547    ) -> Task<Result<Credentials>> {
548        let platform = cx.platform();
549        let executor = cx.background();
550        executor.clone().spawn(async move {
551            // Generate a pair of asymmetric encryption keys. The public key will be used by the
552            // zed server to encrypt the user's access token, so that it can'be intercepted by
553            // any other app running on the user's device.
554            let (public_key, private_key) =
555                rpc::auth::keypair().expect("failed to generate keypair for auth");
556            let public_key_string =
557                String::try_from(public_key).expect("failed to serialize public key for auth");
558
559            // Start an HTTP server to receive the redirect from Zed's sign-in page.
560            let server = tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
561            let port = server.server_addr().port();
562
563            // Open the Zed sign-in page in the user's browser, with query parameters that indicate
564            // that the user is signing in from a Zed app running on the same device.
565            let mut url = format!(
566                "{}/native_app_signin?native_app_port={}&native_app_public_key={}",
567                *SITE_URL, port, public_key_string
568            );
569
570            if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
571                log::info!("impersonating user @{}", impersonate_login);
572                write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
573            }
574
575            platform.open_url(&url);
576
577            // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
578            // access token from the query params.
579            //
580            // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
581            // custom URL scheme instead of this local HTTP server.
582            let (user_id, access_token) = executor
583                .spawn(async move {
584                    if let Some(req) = server.recv_timeout(Duration::from_secs(10 * 60))? {
585                        let path = req.url();
586                        let mut user_id = None;
587                        let mut access_token = None;
588                        let url = Url::parse(&format!("http://example.com{}", path))
589                            .context("failed to parse login notification url")?;
590                        for (key, value) in url.query_pairs() {
591                            if key == "access_token" {
592                                access_token = Some(value.to_string());
593                            } else if key == "user_id" {
594                                user_id = Some(value.to_string());
595                            }
596                        }
597
598                        let post_auth_url = format!("{}/native_app_signin_succeeded", *SITE_URL);
599                        req.respond(
600                            tiny_http::Response::empty(302).with_header(
601                                tiny_http::Header::from_bytes(
602                                    &b"Location"[..],
603                                    post_auth_url.as_bytes(),
604                                )
605                                .unwrap(),
606                            ),
607                        )
608                        .context("failed to respond to login http request")?;
609                        Ok((
610                            user_id.ok_or_else(|| anyhow!("missing user_id parameter"))?,
611                            access_token
612                                .ok_or_else(|| anyhow!("missing access_token parameter"))?,
613                        ))
614                    } else {
615                        Err(anyhow!("didn't receive login redirect"))
616                    }
617                })
618                .await?;
619
620            let access_token = private_key
621                .decrypt_string(&access_token)
622                .context("failed to decrypt access token")?;
623            platform.activate(true);
624
625            Ok(Credentials {
626                user_id: user_id.parse()?,
627                access_token,
628            })
629        })
630    }
631
632    pub async fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) -> Result<()> {
633        let conn_id = self.connection_id()?;
634        self.peer.disconnect(conn_id).await;
635        self.set_status(Status::SignedOut, cx);
636        Ok(())
637    }
638
639    fn connection_id(&self) -> Result<ConnectionId> {
640        if let Status::Connected { connection_id, .. } = *self.status().borrow() {
641            Ok(connection_id)
642        } else {
643            Err(anyhow!("not connected"))
644        }
645    }
646
647    pub async fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
648        self.peer.send(self.connection_id()?, message).await
649    }
650
651    pub async fn request<T: RequestMessage>(&self, request: T) -> Result<T::Response> {
652        self.peer.request(self.connection_id()?, request).await
653    }
654
655    pub fn respond<T: RequestMessage>(
656        &self,
657        receipt: Receipt<T>,
658        response: T::Response,
659    ) -> impl Future<Output = Result<()>> {
660        self.peer.respond(receipt, response)
661    }
662}
663
664fn read_credentials_from_keychain(cx: &AsyncAppContext) -> Option<Credentials> {
665    if IMPERSONATE_LOGIN.is_some() {
666        return None;
667    }
668
669    let (user_id, access_token) = cx
670        .platform()
671        .read_credentials(&COLLAB_URL)
672        .log_err()
673        .flatten()?;
674    Some(Credentials {
675        user_id: user_id.parse().ok()?,
676        access_token: String::from_utf8(access_token).ok()?,
677    })
678}
679
680fn write_credentials_to_keychain(credentials: &Credentials, cx: &AsyncAppContext) -> Result<()> {
681    cx.platform().write_credentials(
682        &COLLAB_URL,
683        &credentials.user_id.to_string(),
684        credentials.access_token.as_bytes(),
685    )
686}
687
688const WORKTREE_URL_PREFIX: &'static str = "zed://worktrees/";
689
690pub fn encode_worktree_url(id: u64, access_token: &str) -> String {
691    format!("{}{}/{}", WORKTREE_URL_PREFIX, id, access_token)
692}
693
694pub fn decode_worktree_url(url: &str) -> Option<(u64, String)> {
695    let path = url.trim().strip_prefix(WORKTREE_URL_PREFIX)?;
696    let mut parts = path.split('/');
697    let id = parts.next()?.parse::<u64>().ok()?;
698    let access_token = parts.next()?;
699    if access_token.is_empty() {
700        return None;
701    }
702    Some((id, access_token.to_string()))
703}
704
705#[cfg(test)]
706mod tests {
707    use super::*;
708    use crate::test::FakeServer;
709    use gpui::TestAppContext;
710
711    #[gpui::test(iterations = 10)]
712    async fn test_heartbeat(cx: TestAppContext) {
713        cx.foreground().forbid_parking();
714
715        let user_id = 5;
716        let mut client = Client::new();
717        let server = FakeServer::for_client(user_id, &mut client, &cx).await;
718
719        cx.foreground().advance_clock(Duration::from_secs(10));
720        let ping = server.receive::<proto::Ping>().await.unwrap();
721        server.respond(ping.receipt(), proto::Ack {}).await;
722
723        cx.foreground().advance_clock(Duration::from_secs(10));
724        let ping = server.receive::<proto::Ping>().await.unwrap();
725        server.respond(ping.receipt(), proto::Ack {}).await;
726
727        client.disconnect(&cx.to_async()).await.unwrap();
728        assert!(server.receive::<proto::Ping>().await.is_err());
729    }
730
731    #[gpui::test(iterations = 10)]
732    async fn test_reconnection(cx: TestAppContext) {
733        cx.foreground().forbid_parking();
734
735        let user_id = 5;
736        let mut client = Client::new();
737        let server = FakeServer::for_client(user_id, &mut client, &cx).await;
738        let mut status = client.status();
739        assert!(matches!(
740            status.recv().await,
741            Some(Status::Connected { .. })
742        ));
743        assert_eq!(server.auth_count(), 1);
744
745        server.forbid_connections();
746        server.disconnect().await;
747        while !matches!(status.recv().await, Some(Status::ReconnectionError { .. })) {}
748
749        server.allow_connections();
750        cx.foreground().advance_clock(Duration::from_secs(10));
751        while !matches!(status.recv().await, Some(Status::Connected { .. })) {}
752        assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
753
754        server.forbid_connections();
755        server.disconnect().await;
756        while !matches!(status.recv().await, Some(Status::ReconnectionError { .. })) {}
757
758        // Clear cached credentials after authentication fails
759        server.roll_access_token();
760        server.allow_connections();
761        cx.foreground().advance_clock(Duration::from_secs(10));
762        assert_eq!(server.auth_count(), 1);
763        cx.foreground().advance_clock(Duration::from_secs(10));
764        while !matches!(status.recv().await, Some(Status::Connected { .. })) {}
765        assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
766    }
767
768    #[test]
769    fn test_encode_and_decode_worktree_url() {
770        let url = encode_worktree_url(5, "deadbeef");
771        assert_eq!(decode_worktree_url(&url), Some((5, "deadbeef".to_string())));
772        assert_eq!(
773            decode_worktree_url(&format!("\n {}\t", url)),
774            Some((5, "deadbeef".to_string()))
775        );
776        assert_eq!(decode_worktree_url("not://the-right-format"), None);
777    }
778}