2088 lines · 69.4 KB · Rust
RawDownload
#[cfg(any(test, feature = "test-support"))]pub mod test; mod socks;pub mod telemetry;pub mod user;pub mod zed_urls; use anyhow::{anyhow, bail, Context as _, Result};use async_recursion::async_recursion;use async_tungstenite::tungstenite::{    client::IntoClientRequest,    error::Error as WebsocketError,    http::{HeaderValue, Request, StatusCode},};use chrono::{DateTime, Utc};use clock::SystemClock;use futures::{    channel::oneshot, future::BoxFuture, AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt,    TryFutureExt as _, TryStreamExt,};use gpui::{actions, AppContext, AsyncAppContext, Global, Model, Task, WeakModel};use http_client::{AsyncBody, HttpClient, HttpClientWithUrl};use parking_lot::RwLock;use postage::watch;use rand::prelude::*;use release_channel::{AppVersion, ReleaseChannel};use rpc::proto::{AnyTypedEnvelope, EnvelopedMessage, PeerId, RequestMessage};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use settings::{Settings, SettingsSources};use socks::connect_socks_proxy_stream;use std::pin::Pin;use std::{    any::TypeId,    convert::TryFrom,    fmt::Write as _,    future::Future,    marker::PhantomData,    path::PathBuf,    sync::{        atomic::{AtomicU64, Ordering},        Arc, LazyLock, Weak,    },    time::{Duration, Instant},};use telemetry::Telemetry;use thiserror::Error;use url::Url;use util::{ResultExt, TryFutureExt}; 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()); /// An environment variable whose presence indicates that the development auth/// provider should be used.////// Only works in development. Setting this environment variable in other release/// channels is a no-op.pub static ZED_DEVELOPMENT_AUTH: LazyLock<bool> = LazyLock::new(|| {    std::env::var("ZED_DEVELOPMENT_AUTH").map_or(false, |value| !value.is_empty())});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 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").map_or(false, |e| !e.is_empty())); pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(500);pub const MAX_RECONNECTION_DELAY: Duration = Duration::from_secs(10);pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20); actions!(client, [SignIn, SignOut, Reconnect]); #[derive(Clone, Default, Serialize, Deserialize, JsonSchema)]pub struct ClientSettingsContent {    server_url: Option<String>,} #[derive(Deserialize)]pub struct ClientSettings {    pub server_url: String,} impl Settings for ClientSettings {    const KEY: Option<&'static str> = None;     type FileContent = ClientSettingsContent;     fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {        let mut result = sources.json_merge::<Self>()?;        if let Some(server_url) = &*ZED_SERVER_URL {            result.server_url.clone_from(server_url)        }        Ok(result)    }} #[derive(Default, Clone, Serialize, Deserialize, JsonSchema)]pub struct ProxySettingsContent {    proxy: Option<String>,} #[derive(Deserialize, Default)]pub struct ProxySettings {    pub proxy: Option<String>,} impl Settings for ProxySettings {    const KEY: Option<&'static str> = None;     type FileContent = ProxySettingsContent;     fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {        Ok(Self {            proxy: sources                .user                .or(sources.server)                .and_then(|value| value.proxy.clone())                .or(sources.default.proxy.clone()),        })    }} pub fn init_settings(cx: &mut AppContext) {    TelemetrySettings::register(cx);    ClientSettings::register(cx);    ProxySettings::register(cx);} pub fn init(client: &Arc<Client>, cx: &mut AppContext) {    let client = Arc::downgrade(client);    cx.on_action({        let client = client.clone();        move |_: &SignIn, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(                    |cx| async move { client.authenticate_and_connect(true, &cx).log_err().await },                )                .detach();            }        }    });     cx.on_action({        let client = client.clone();        move |_: &SignOut, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(|cx| async move {                    client.sign_out(&cx).await;                })                .detach();            }        }    });     cx.on_action({        let client = client.clone();        move |_: &Reconnect, cx| {            if let Some(client) = client.upgrade() {                cx.spawn(|cx| async move {                    client.reconnect(&cx);                })                .detach();            }        }    });} struct GlobalClient(Arc<Client>); impl Global for GlobalClient {} pub struct Client {    id: AtomicU64,    peer: Arc<Peer>,    http: Arc<HttpClientWithUrl>,    telemetry: Arc<Telemetry>,    credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static>,    state: RwLock<ClientState>,    handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    authenticate: RwLock<        Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,    >,     #[allow(clippy::type_complexity)]    #[cfg(any(test, feature = "test-support"))]    establish_connection: RwLock<        Option<            Box<                dyn 'static                    + Send                    + Sync                    + Fn(                        &Credentials,                        &AsyncAppContext,                    ) -> 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,    Connecting,    ConnectionError,    Connected {        peer_id: PeerId,        connection_id: ConnectionId,    },    ConnectionLost,    Reauthenticating,    Reconnecting,    ReconnectionError {        next_reconnection: Instant,    },} impl Status {    pub fn is_connected(&self) -> bool {        matches!(self, Self::Connected { .. })    }     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)    }} /// A provider for [`Credentials`].////// Used to abstract over reading and writing credentials to some form of/// persistence (like the system keychain).trait CredentialsProvider {    /// Reads the credentials from the provider.    fn read_credentials<'a>(        &'a self,        cx: &'a AsyncAppContext,    ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>>;     /// Writes the credentials to the provider.    fn write_credentials<'a>(        &'a self,        user_id: u64,        access_token: String,        cx: &'a AsyncAppContext,    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>>;     /// Deletes the credentials from the provider.    fn delete_credentials<'a>(        &'a self,        cx: &'a AsyncAppContext,    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>>;} 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_model(mut self, model: &Model<T>, cx: &AsyncAppContext) -> 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: model.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))            {                for message in messages {                    log::info!("unhandled message {}", message.payload_type_name());                }            }        }    }} #[derive(Copy, Clone)]pub struct TelemetrySettings {    pub diagnostics: bool,    pub metrics: bool,} /// Control what info is collected by Zed.#[derive(Default, Clone, Serialize, Deserialize, JsonSchema)]pub struct TelemetrySettingsContent {    /// Send debug info like crash reports.    ///    /// Default: true    pub diagnostics: Option<bool>,    /// Send anonymized usage data like what languages you're using Zed with.    ///    /// Default: true    pub metrics: Option<bool>,} impl settings::Settings for TelemetrySettings {    const KEY: Option<&'static str> = Some("telemetry");     type FileContent = TelemetrySettingsContent;     fn load(sources: SettingsSources<Self::FileContent>, _: &mut AppContext) -> Result<Self> {        Ok(Self {            diagnostics: sources                .user                .as_ref()                .or(sources.server.as_ref())                .and_then(|v| v.diagnostics)                .unwrap_or(                    sources                        .default                        .diagnostics                        .ok_or_else(Self::missing_default)?,                ),            metrics: sources                .user                .as_ref()                .or(sources.server.as_ref())                .and_then(|v| v.metrics)                .unwrap_or(sources.default.metrics.ok_or_else(Self::missing_default)?),        })    }} impl Client {    pub fn new(        clock: Arc<dyn SystemClock>,        http: Arc<HttpClientWithUrl>,        cx: &mut AppContext,    ) -> Arc<Self> {        let use_zed_development_auth = match ReleaseChannel::try_global(cx) {            Some(ReleaseChannel::Dev) => *ZED_DEVELOPMENT_AUTH,            Some(ReleaseChannel::Nightly | ReleaseChannel::Preview | ReleaseChannel::Stable)            | None => false,        };         let credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static> =            if use_zed_development_auth {                Arc::new(DevelopmentCredentialsProvider {                    path: paths::config_dir().join("development_auth"),                })            } else {                Arc::new(KeychainCredentialsProvider)

Lines 1–500 of 2088.