Add `ZED_IMPERSONATE` env var, for testing

Max Brunsfeld and Nathan Sobo created

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

Change summary

server/src/auth.rs | 28 +++++++++++++++++++++++-----
server/src/db.rs   |  9 ++++++---
zed/src/rpc.rs     | 27 ++++++++++++++++++++-------
3 files changed, 49 insertions(+), 15 deletions(-)

Detailed changes

server/src/auth.rs 🔗

@@ -18,7 +18,7 @@ use scrypt::{
 use serde::{Deserialize, Serialize};
 use std::{borrow::Cow, convert::TryFrom, sync::Arc};
 use surf::{StatusCode, Url};
-use tide::Server;
+use tide::{log, Server};
 use zrpc::auth as zed_auth;
 
 static CURRENT_GITHUB_USER: &'static str = "current_github_user";
@@ -121,6 +121,7 @@ pub fn add_routes(app: &mut Server<Arc<AppState>>) {
 struct NativeAppSignInParams {
     native_app_port: String,
     native_app_public_key: String,
+    impersonate: Option<String>,
 }
 
 async fn get_sign_in(mut request: Request) -> tide::Result {
@@ -142,11 +143,15 @@ async fn get_sign_in(mut request: Request) -> tide::Result {
 
     let app_sign_in_params: Option<NativeAppSignInParams> = request.query().ok();
     if let Some(query) = app_sign_in_params {
-        redirect_url
-            .query_pairs_mut()
+        let mut redirect_query = redirect_url.query_pairs_mut();
+        redirect_query
             .clear()
             .append_pair("native_app_port", &query.native_app_port)
             .append_pair("native_app_public_key", &query.native_app_public_key);
+
+        if let Some(impersonate) = &query.impersonate {
+            redirect_query.append_pair("impersonate", impersonate);
+        }
     }
 
     let (auth_url, csrf_token) = request
@@ -222,7 +227,20 @@ async fn get_auth_callback(mut request: Request) -> tide::Result {
     // When signing in from the native app, generate a new access token for the current user. Return
     // a redirect so that the user's browser sends this access token to the locally-running app.
     if let Some((user, app_sign_in_params)) = user.zip(query.native_app_sign_in_params) {
-        let access_token = create_access_token(request.db(), user.id).await?;
+        let mut user_id = user.id;
+        if let Some(impersonated_login) = app_sign_in_params.impersonate {
+            log::info!("attempting to impersonate user @{}", impersonated_login);
+            if let Some(user) = request.db().get_users_by_ids([user_id]).await?.first() {
+                if user.admin {
+                    user_id = request.db().create_user(&impersonated_login, false).await?;
+                    log::info!("impersonating user {}", user_id.0);
+                } else {
+                    log::info!("refusing to impersonate user");
+                }
+            }
+        }
+
+        let access_token = create_access_token(request.db(), user_id).await?;
         let native_app_public_key =
             zed_auth::PublicKey::try_from(app_sign_in_params.native_app_public_key.clone())
                 .context("failed to parse app public key")?;
@@ -232,7 +250,7 @@ async fn get_auth_callback(mut request: Request) -> tide::Result {
 
         return Ok(tide::Redirect::new(&format!(
             "http://127.0.0.1:{}?user_id={}&access_token={}",
-            app_sign_in_params.native_app_port, user.id.0, encrypted_access_token,
+            app_sign_in_params.native_app_port, user_id.0, encrypted_access_token,
         ))
         .into());
     }

server/src/db.rs 🔗

@@ -108,8 +108,11 @@ impl Db {
         })
     }
 
-    pub async fn get_users_by_ids(&self, ids: impl Iterator<Item = UserId>) -> Result<Vec<User>> {
-        let ids = ids.map(|id| id.0).collect::<Vec<_>>();
+    pub async fn get_users_by_ids(
+        &self,
+        ids: impl IntoIterator<Item = UserId>,
+    ) -> Result<Vec<User>> {
+        let ids = ids.into_iter().map(|id| id.0).collect::<Vec<_>>();
         test_support!(self, {
             let query = "
                 SELECT users.*
@@ -547,7 +550,7 @@ pub mod tests {
         let friend3 = db.create_user("friend-3", false).await.unwrap();
 
         assert_eq!(
-            db.get_users_by_ids([user, friend1, friend2, friend3].iter().copied())
+            db.get_users_by_ids([user, friend1, friend2, friend3])
                 .await
                 .unwrap(),
             vec![

zed/src/rpc.rs 🔗

@@ -14,6 +14,7 @@ use std::{
     any::TypeId,
     collections::HashMap,
     convert::TryFrom,
+    fmt::Write as _,
     future::Future,
     sync::{Arc, Weak},
     time::{Duration, Instant},
@@ -29,6 +30,7 @@ use zrpc::{
 lazy_static! {
     static ref ZED_SERVER_URL: String =
         std::env::var("ZED_SERVER_URL").unwrap_or("https://zed.dev:443".to_string());
+    static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE").ok();
 }
 
 pub struct Client {
@@ -350,12 +352,12 @@ impl Client {
             self.set_status(Status::Reauthenticating, cx)
         }
 
-        let mut read_from_keychain = false;
+        let mut used_keychain = false;
         let credentials = self.state.read().credentials.clone();
         let credentials = if let Some(credentials) = credentials {
             credentials
         } else if let Some(credentials) = read_credentials_from_keychain(cx) {
-            read_from_keychain = true;
+            used_keychain = true;
             credentials
         } else {
             let credentials = match self.authenticate(&cx).await {
@@ -378,7 +380,7 @@ impl Client {
             Ok(conn) => {
                 log::info!("connected to rpc address {}", *ZED_SERVER_URL);
                 self.state.write().credentials = Some(credentials.clone());
-                if !read_from_keychain {
+                if !used_keychain && IMPERSONATE_LOGIN.is_none() {
                     write_credentials_to_keychain(&credentials, cx).log_err();
                 }
                 self.set_connection(conn, cx).await;
@@ -387,8 +389,8 @@ impl Client {
             Err(err) => {
                 if matches!(err, EstablishConnectionError::Unauthorized) {
                     self.state.write().credentials.take();
-                    cx.platform().delete_credentials(&ZED_SERVER_URL).log_err();
-                    if read_from_keychain {
+                    if used_keychain {
+                        cx.platform().delete_credentials(&ZED_SERVER_URL).log_err();
                         self.set_status(Status::SignedOut, cx);
                         self.authenticate_and_connect(cx).await
                     } else {
@@ -524,10 +526,17 @@ impl Client {
 
             // Open the Zed sign-in page in the user's browser, with query parameters that indicate
             // that the user is signing in from a Zed app running on the same device.
-            platform.open_url(&format!(
+            let mut url = format!(
                 "{}/sign_in?native_app_port={}&native_app_public_key={}",
                 *ZED_SERVER_URL, port, public_key_string
-            ));
+            );
+
+            if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
+                log::info!("impersonating user @{}", impersonate_login);
+                write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
+            }
+
+            platform.open_url(&url);
 
             // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
             // access token from the query params.
@@ -611,6 +620,10 @@ impl Client {
 }
 
 fn read_credentials_from_keychain(cx: &AsyncAppContext) -> Option<Credentials> {
+    if IMPERSONATE_LOGIN.is_some() {
+        return None;
+    }
+
     let (user_id, access_token) = cx
         .platform()
         .read_credentials(&ZED_SERVER_URL)