2225 lines · 74.5 KB · Rust
RawDownload
#[cfg(any(test, feature = "test-support"))]pub mod test; mod proxy;pub mod telemetry;pub mod user;pub mod zed_urls; use anyhow::{Context as _, Result, anyhow};use async_tungstenite::tungstenite::{    client::IntoClientRequest,    error::Error as WebsocketError,    http::{HeaderValue, Request, StatusCode},};use clock::SystemClock;use cloud_api_client::CloudApiClient;use cloud_api_client::websocket_protocol::MessageToClient;use credentials_provider::CredentialsProvider;use feature_flags::FeatureFlagAppExt as _;use futures::{    AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,    channel::oneshot, future::BoxFuture,};use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity, actions};use http_client::{HttpClient, HttpClientWithUrl, http, read_proxy_from_env};use parking_lot::RwLock;use postage::watch;use proxy::connect_proxy_stream;use rand::prelude::*;use release_channel::{AppVersion, ReleaseChannel};use rpc::proto::{AnyTypedEnvelope, EnvelopedMessage, PeerId, RequestMessage};use serde::{Deserialize, Serialize};use settings::{RegisterSetting, Settings, SettingsContent};use std::{    any::TypeId,    convert::TryFrom,    fmt::Write as _,    future::Future,    marker::PhantomData,    path::PathBuf,    sync::{        Arc, LazyLock, Weak,        atomic::{AtomicU64, Ordering},    },    time::{Duration, Instant},};use std::{cmp, pin::Pin};use telemetry::Telemetry;use thiserror::Error;use tokio::net::TcpStream;use url::Url;use util::{ConnectionResult, ResultExt}; pub use rpc::*;pub use telemetry_events::Event;pub use user::*; static ZED_SERVER_URL: LazyLock<Option<String>> =    LazyLock::new(|| std::env::var("ZED_SERVER_URL").ok());static ZED_RPC_URL: LazyLock<Option<String>> = LazyLock::new(|| std::env::var("ZED_RPC_URL").ok()); pub static IMPERSONATE_LOGIN: LazyLock<Option<String>> = LazyLock::new(|| {    std::env::var("ZED_IMPERSONATE")        .ok()        .and_then(|s| if s.is_empty() { None } else { Some(s) })}); pub static USE_WEB_LOGIN: LazyLock<bool> = LazyLock::new(|| std::env::var("ZED_WEB_LOGIN").is_ok()); pub static ADMIN_API_TOKEN: LazyLock<Option<String>> = LazyLock::new(|| {    std::env::var("ZED_ADMIN_API_TOKEN")        .ok()        .and_then(|s| if s.is_empty() { None } else { Some(s) })}); pub static ZED_APP_PATH: LazyLock<Option<PathBuf>> =    LazyLock::new(|| std::env::var("ZED_APP_PATH").ok().map(PathBuf::from)); pub static ZED_ALWAYS_ACTIVE: LazyLock<bool> =    LazyLock::new(|| std::env::var("ZED_ALWAYS_ACTIVE").is_ok_and(|e| !e.is_empty())); pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(500);pub const MAX_RECONNECTION_DELAY: Duration = Duration::from_secs(30);pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20); actions!(    client,    [        /// Signs in to Zed account.        SignIn,        /// Signs out of Zed account.        SignOut,        /// Reconnects to the collaboration server.        Reconnect    ]); #[derive(Deserialize, RegisterSetting)]pub struct ClientSettings {    pub server_url: String,} impl Settings for ClientSettings {    fn from_settings(content: &settings::SettingsContent) -> Self {        if let Some(server_url) = &*ZED_SERVER_URL {            return Self {                server_url: server_url.clone(),            };        }        Self {            server_url: content.server_url.clone().unwrap(),        }    }} #[derive(Deserialize, Default, RegisterSetting)]pub struct ProxySettings {    pub proxy: Option<String>,} impl ProxySettings {    pub fn proxy_url(&self) -> Option<Url> {        self.proxy            .as_deref()            .map(str::trim)            .filter(|input| !input.is_empty())            .and_then(|input| {                input                    .parse::<Url>()                    .inspect_err(|e| log::error!("Error parsing proxy settings: {}", e))                    .ok()            })            .or_else(read_proxy_from_env)    }} impl Settings for ProxySettings {    fn from_settings(content: &settings::SettingsContent) -> Self {        Self {            proxy: content                .proxy                .as_deref()                .map(str::trim)                .filter(|proxy| !proxy.is_empty())                .map(ToOwned::to_owned),        }    }} pub fn init(client: &Arc<Client>, cx: &mut App) {    let client = Arc::downgrade(client);    cx.on_action({        let client = client.clone();        move |_: &SignIn, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(async move |cx| client.sign_in_with_optional_connect(true, cx).await)                    .detach_and_log_err(cx);            }        }    })    .on_action({        let client = client.clone();        move |_: &SignOut, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(async move |cx| {                    client.sign_out(cx).await;                })                .detach();            }        }    })    .on_action({        let client = client;        move |_: &Reconnect, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(async move |cx| {                    client.reconnect(cx);                })                .detach();            }        }    });} pub type MessageToClientHandler = Box<dyn Fn(&MessageToClient, &mut App) + Send + Sync + 'static>; struct GlobalClient(Arc<Client>); impl Global for GlobalClient {} pub struct Client {    id: AtomicU64,    peer: Arc<Peer>,    http: Arc<HttpClientWithUrl>,    cloud_client: Arc<CloudApiClient>,    telemetry: Arc<Telemetry>,    credentials_provider: ClientCredentialsProvider,    state: RwLock<ClientState>,    handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,    message_to_client_handlers: parking_lot::Mutex<Vec<MessageToClientHandler>>,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    authenticate:        RwLock<Option<Box<dyn 'static + Send + Sync + Fn(&AsyncApp) -> Task<Result<Credentials>>>>>,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    establish_connection: RwLock<        Option<            Box<                dyn 'static                    + Send                    + Sync                    + Fn(                        &Credentials,                        &AsyncApp,                    ) -> Task<Result<Connection, EstablishConnectionError>>,            >,        >,    >,     #[cfg(any(test, feature = "test-support"))]    rpc_url: RwLock<Option<Url>>,} #[derive(Error, Debug)]pub enum EstablishConnectionError {    #[error("upgrade required")]    UpgradeRequired,    #[error("unauthorized")]    Unauthorized,    #[error("{0}")]    Other(#[from] anyhow::Error),    #[error("{0}")]    InvalidHeaderValue(#[from] async_tungstenite::tungstenite::http::header::InvalidHeaderValue),    #[error("{0}")]    Io(#[from] std::io::Error),    #[error("{0}")]    Websocket(#[from] async_tungstenite::tungstenite::http::Error),} impl From<WebsocketError> for EstablishConnectionError {    fn from(error: WebsocketError) -> Self {        if let WebsocketError::Http(response) = &error {            match response.status() {                StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,                StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,                _ => {}            }        }        EstablishConnectionError::Other(error.into())    }} impl EstablishConnectionError {    pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {        Self::Other(error.into())    }} #[derive(Copy, Clone, Debug, PartialEq)]pub enum Status {    SignedOut,    UpgradeRequired,    Authenticating,    Authenticated,    AuthenticationError,    Connecting,    ConnectionError,    Connected {        peer_id: PeerId,        connection_id: ConnectionId,    },    ConnectionLost,    Reauthenticating,    Reauthenticated,    Reconnecting,    ReconnectionError {        next_reconnection: Instant,    },} impl Status {    pub fn is_connected(&self) -> bool {        matches!(self, Self::Connected { .. })    }     pub fn was_connected(&self) -> bool {        matches!(            self,            Self::ConnectionLost                | Self::Reauthenticating                | Self::Reauthenticated                | Self::Reconnecting        )    }     /// Returns whether the client is currently connected or was connected at some point.    pub fn is_or_was_connected(&self) -> bool {        self.is_connected() || self.was_connected()    }     pub fn is_signing_in(&self) -> bool {        matches!(            self,            Self::Authenticating | Self::Reauthenticating | Self::Connecting | Self::Reconnecting        )    }     pub fn is_signed_out(&self) -> bool {        matches!(self, Self::SignedOut | Self::UpgradeRequired)    }} struct ClientState {    credentials: Option<Credentials>,    status: (watch::Sender<Status>, watch::Receiver<Status>),    _reconnect_task: Option<Task<()>>,} #[derive(Clone, Debug, Eq, PartialEq)]pub struct Credentials {    pub user_id: u64,    pub access_token: String,} impl Credentials {    pub fn authorization_header(&self) -> String {        format!("{} {}", self.user_id, self.access_token)    }} pub struct ClientCredentialsProvider {    provider: Arc<dyn CredentialsProvider>,} impl ClientCredentialsProvider {    pub fn new(cx: &App) -> Self {        Self {            provider: <dyn CredentialsProvider>::global(cx),        }    }     fn server_url(&self, cx: &AsyncApp) -> Result<String> {        Ok(cx.update(|cx| ClientSettings::get_global(cx).server_url.clone()))    }     /// Reads the credentials from the provider.    fn read_credentials<'a>(        &'a self,        cx: &'a AsyncApp,    ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {        async move {            if IMPERSONATE_LOGIN.is_some() {                return None;            }             let server_url = self.server_url(cx).ok()?;            let (user_id, access_token) = self                .provider                .read_credentials(&server_url, cx)                .await                .log_err()                .flatten()?;             Some(Credentials {                user_id: user_id.parse().ok()?,                access_token: String::from_utf8(access_token).ok()?,            })        }        .boxed_local()    }     /// Writes the credentials to the provider.    fn write_credentials<'a>(        &'a self,        user_id: u64,        access_token: String,        cx: &'a AsyncApp,    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {        async move {            let server_url = self.server_url(cx)?;            self.provider                .write_credentials(                    &server_url,                    &user_id.to_string(),                    access_token.as_bytes(),                    cx,                )                .await        }        .boxed_local()    }     /// Deletes the credentials from the provider.    fn delete_credentials<'a>(        &'a self,        cx: &'a AsyncApp,    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {        async move {            let server_url = self.server_url(cx)?;            self.provider.delete_credentials(&server_url, cx).await        }        .boxed_local()    }} impl Default for ClientState {    fn default() -> Self {        Self {            credentials: None,            status: watch::channel_with(Status::SignedOut),            _reconnect_task: None,        }    }} pub enum Subscription {    Entity {        client: Weak<Client>,        id: (TypeId, u64),    },    Message {        client: Weak<Client>,        id: TypeId,    },} impl Drop for Subscription {    fn drop(&mut self) {        match self {            Subscription::Entity { client, id } => {                if let Some(client) = client.upgrade() {                    let mut state = client.handler_set.lock();                    let _ = state.entities_by_type_and_remote_id.remove(id);                }            }            Subscription::Message { client, id } => {                if let Some(client) = client.upgrade() {                    let mut state = client.handler_set.lock();                    let _ = state.entity_types_by_message_type.remove(id);                    let _ = state.message_handlers.remove(id);                }            }        }    }} pub struct PendingEntitySubscription<T: 'static> {    client: Arc<Client>,    remote_id: u64,    _entity_type: PhantomData<T>,    consumed: bool,} impl<T: 'static> PendingEntitySubscription<T> {    pub fn set_entity(mut self, entity: &Entity<T>, cx: &AsyncApp) -> Subscription {        self.consumed = true;        let mut handlers = self.client.handler_set.lock();        let id = (TypeId::of::<T>(), self.remote_id);        let Some(EntityMessageSubscriber::Pending(messages)) =            handlers.entities_by_type_and_remote_id.remove(&id)        else {            unreachable!()        };         handlers.entities_by_type_and_remote_id.insert(            id,            EntityMessageSubscriber::Entity {                handle: entity.downgrade().into(),            },        );        drop(handlers);        for message in messages {            let client_id = self.client.id();            let type_name = message.payload_type_name();            let sender_id = message.original_sender_id();            log::debug!(                "handling queued rpc message. client_id:{}, sender_id:{:?}, type:{}",                client_id,                sender_id,                type_name            );            self.client.handle_message(message, cx);        }        Subscription::Entity {            client: Arc::downgrade(&self.client),            id,        }    }} impl<T: 'static> Drop for PendingEntitySubscription<T> {    fn drop(&mut self) {        if !self.consumed {            let mut state = self.client.handler_set.lock();            if let Some(EntityMessageSubscriber::Pending(messages)) = state                .entities_by_type_and_remote_id                .remove(&(TypeId::of::<T>(), self.remote_id))

Lines 1–500 of 2225.