Add device code auth in the dashboard, fix Linux themes (#16)

* Add device code auth in the dashboard, fix Linux themes

* Fix clipboard copy text

Author: brandonskiser

Cargo.lock

@@ -188,7 +188,7 @@ export const walkSubcommand = <
       (state, annotation) => {
         const { text, type } = annotation;
         const name =
-          "tokenName" in annotation ? annotation.tokenName ?? text : text;
+          "tokenName" in annotation ? (annotation.tokenName ?? text) : text;
 
         if (type === TokenType.Subcommand) {
           if (!state.subcommand.subcommands[name]) {

apps/autocomplete/src/history/helpers.ts

@@ -20,9 +20,20 @@ export default function LoginModal({ next }: { next: () => void }) {
     z.enum(["builderId", "iam"]),
     midway ? "iam" : "builderId",
   );
+
+  // Since PKCE requires the ability to open a browser, we also support falling
+  // back to device code in case of an error.
+  const [loginMethod, setLoginMethod] = useState<"pkce" | "deviceCode">("pkce");
+
+  // used for pkce
   const [currAuthRequestId, setAuthRequestId] = useAuthRequest();
 
-  // console.log(tab);
+  // used for device code
+  const [loginCode, setLoginCode] = useState<string | null>(null);
+  const [loginUrl, setLoginUrl] = useState<string | null>(null);
+  const [copyToClipboardText, setCopyToClipboardText] = useState<
+    "Copy to clipboard" | "Copied!"
+  >("Copy to clipboard");
 
   const [error, setError] = useState<string | null>(null);
   const [completedOnboarding] = useLocalStateZodDefault(
@@ -33,7 +44,15 @@ export default function LoginModal({ next }: { next: () => void }) {
   const auth = useAuth();
   const refreshAuth = useRefreshAuth();
 
-  async function handleLogin(issuerUrl?: string, region?: string) {
+  async function handleLogin(startUrl?: string, region?: string) {
+    if (loginMethod === "pkce") {
+      handlePkceAuth(startUrl, region);
+    } else {
+      handleDeviceCodeAuth(startUrl, region);
+    }
+  }
+
+  async function handlePkceAuth(issuerUrl?: string, region?: string) {
     setLoginState("loading");
     setError(null);
     // We need to reset the auth request state before attempting, otherwise
@@ -53,7 +72,13 @@ export default function LoginModal({ next }: { next: () => void }) {
     setAuthRequestId(init.authRequestId);
 
     Native.open(init.url).catch((err) => {
+      // Reset the auth request id so that we don't present the "OAuth cancelled" error
+      // to the user.
+      setAuthRequestId("");
       console.error(err);
+      setError(
+        "Failed to open the browser. As an alternative, try logging in with device code.",
+      );
     });
 
     await Auth.finishPkceAuthorization(init)
@@ -74,6 +99,44 @@ export default function LoginModal({ next }: { next: () => void }) {
       });
   }
 
+  async function handleDeviceCodeAuth(startUrl?: string, region?: string) {
+    setLoginState("loading");
+    setError(null);
+    setLoginUrl(null);
+    setCopyToClipboardText("Copy to clipboard");
+    await Auth.cancelPkceAuthorization().catch((err) => {
+      console.error(err);
+    });
+    const init = await Auth.builderIdStartDeviceAuthorization({
+      startUrl,
+      region,
+    }).catch((err) => {
+      setLoginState("not started");
+      setLoginCode(null);
+      setError(err.message);
+      console.error(err);
+    });
+
+    if (!init) return;
+
+    setLoginCode(init.code);
+    setLoginUrl(init.url);
+
+    await Auth.builderIdPollCreateToken(init)
+      .then(() => {
+        setLoginState("logged in");
+        Internal.sendWindowFocusRequest({});
+        refreshAuth();
+        next();
+      })
+      .catch((err) => {
+        setLoginState("not started");
+        setLoginCode(null);
+        setError(err.message);
+        console.error(err);
+      });
+  }
+
   useEffect(() => {
     setLoginState(auth.authed ? "logged in" : "not started");
   }, [auth]);
@@ -103,16 +166,31 @@ export default function LoginModal({ next }: { next: () => void }) {
           </div>
         )}
       </div>
+
       {error && (
-        <div className="w-full bg-red-200 border border-red-600 rounded py-1 px-1">
+        <div className="flex flex-col items-center gap-2 w-full bg-red-200 border border-red-600 rounded py-2 px-2">
           <p className="text-black dark:text-white font-semibold text-center">
             Failed to login
           </p>
           <p className="text-black dark:text-white text-center">{error}</p>
+          {loginMethod === "pkce" && loginState === "loading" && (
+            <Button
+              variant="ghost"
+              className="self-center mx-auto text-black hover:bg-white/40"
+              onClick={() => {
+                setLoginMethod("deviceCode");
+                setLoginState("not started");
+                setError(null);
+              }}
+            >
+              Login with Device Code
+            </Button>
+          )}
         </div>
       )}
-      <div className="flex flex-col gap-4 text-white text-sm">
-        {loginState === "loading" ? (
+
+      <div className="flex flex-col items-center gap-4 text-white text-sm">
+        {loginState === "loading" && loginMethod === "pkce" ? (
           <>
             <p className="text-center w-80">
               Waiting for authentication in the browser to complete...
@@ -127,6 +205,37 @@ export default function LoginModal({ next }: { next: () => void }) {
               Back
             </Button>
           </>
+        ) : loginState === "loading" &&
+          loginMethod === "deviceCode" &&
+          loginCode &&
+          loginUrl ? (
+          <>
+            <p className="text-center w-80">
+              Confirm code <span className="font-bold">{loginCode}</span> in the
+              login page at the following link:
+            </p>
+            <p className="text-center">{loginUrl}</p>
+            <Button
+              variant="ghost"
+              className="h-auto p-1 px-2 hover:bg-white/20 hover:text-white"
+              onClick={() => {
+                navigator.clipboard.writeText(loginUrl);
+                setCopyToClipboardText("Copied!");
+              }}
+            >
+              {copyToClipboardText}
+            </Button>
+            <Button
+              variant="glass"
+              className="self-center w-32"
+              onClick={() => {
+                setLoginState("not started");
+                setLoginCode(null);
+              }}
+            >
+              Back
+            </Button>
+          </>
         ) : (
           <Tab
             tab={tab}

apps/dashboard/src/components/installs/modal/login/index.tsx

@@ -41,6 +41,7 @@ whoami.workspace = true
 
 [dev-dependencies]
 reqwest.workspace = true
+tracing-subscriber.workspace = true
 
 [target.'cfg(target_os="macos")'.dependencies]
 macos-utils = { path = "../macos-utils" }

crates/fig_desktop_api/Cargo.toml

@@ -182,6 +182,7 @@ where
                 ApplicationUpdateStatusRequest,
                 AuthBuilderIdPollCreateTokenRequest,
                 AuthBuilderIdStartDeviceAuthorizationRequest,
+                AuthCancelPkceAuthorizationRequest,
                 AuthFinishPkceAuthorizationRequest,
                 AuthStartPkceAuthorizationRequest,
                 AuthStatusRequest,
@@ -285,6 +286,7 @@ where
                     event_handler.user_logged_in_callback(ctx).await;
                     result
                 },
+                AuthCancelPkceAuthorizationRequest(request) => auth::cancel_pkce_authorization(request).await,
                 AuthBuilderIdStartDeviceAuthorizationRequest(request) => {
                     auth::builder_id_start_device_authorization(request, &ctx).await
                 },

crates/fig_desktop_api/src/handler.rs

@@ -5,6 +5,7 @@ use std::env::{
 };
 
 use camino::Utf8PathBuf;
+use fig_os_shim::Context;
 use fig_util::directories::midway_cookie_path;
 #[cfg(target_os = "linux")]
 use fig_util::system_info::linux::{
@@ -66,7 +67,8 @@ pub struct Constants {
 
 impl Constants {
     fn new(support_api_proto: bool) -> Self {
-        let themes_folder = directories::themes_dir()
+        let ctx = Context::new();
+        let themes_folder = directories::themes_dir(&ctx)
             .ok()
             .and_then(|dir| Utf8PathBuf::try_from(dir).ok());
 

crates/fig_desktop_api/src/init_script.rs

@@ -19,6 +19,8 @@ use fig_proto::fig::{
     AuthBuilderIdPollCreateTokenResponse,
     AuthBuilderIdStartDeviceAuthorizationRequest,
     AuthBuilderIdStartDeviceAuthorizationResponse,
+    AuthCancelPkceAuthorizationRequest,
+    AuthCancelPkceAuthorizationResponse,
     AuthFinishPkceAuthorizationRequest,
     AuthFinishPkceAuthorizationResponse,
     AuthStartPkceAuthorizationRequest,
@@ -33,7 +35,10 @@ use tokio::sync::mpsc::{
     Sender,
     channel,
 };
-use tracing::error;
+use tracing::{
+    debug,
+    error,
+};
 
 use super::RequestResult;
 use crate::kv::KVStore;
@@ -89,6 +94,9 @@ struct PkceState<T> {
 
     /// [Receiver] for completed PKCE authorizations. Created per authorization attempt.
     finished_rx: Mutex<Receiver<(String, Result<(), fig_auth::Error>)>>,
+
+    /// [Sender] for cancelling any potentially ongoing PKCE authorizations.
+    cancel_tx: Mutex<Sender<()>>,
 }
 
 impl<T: PkceClient + Send + Sync + 'static> PkceState<T> {
@@ -97,26 +105,37 @@ impl<T: PkceClient + Send + Sync + 'static> PkceState<T> {
     fn new() -> Arc<Self> {
         let (new_tx, mut new_rx) = channel::<(String, PkceRegistration, T, Option<SecretStore>)>(1);
         let (finished_tx, finished_rx) = channel(1);
+        let (cancel_tx, mut cancel_rx) = channel(1);
 
         let new_tx_clone = new_tx.clone();
         let pkce_state = Arc::new(PkceState {
             request_id: Mutex::new("".into()),
             new_tx,
             finished_tx: Mutex::new(finished_tx),
             finished_rx: Mutex::new(finished_rx),
+            cancel_tx: Mutex::new(cancel_tx),
         });
         let pkce_state_clone = Arc::clone(&pkce_state);
         tokio::spawn(async move {
             let registration_tx = new_tx_clone;
             let pkce_state = pkce_state_clone;
             while let Some((auth_request_id, registration, client, secret_store)) = new_rx.recv().await {
+                while cancel_rx.try_recv().is_ok() {
+                    debug!("Ignoring buffered PKCE cancellation attempt");
+                }
+
                 tokio::select! {
                     // We received a new registration while waiting for the current one to complete,
                     // so resend it back around the loop.
                     Some(new_registration) = new_rx.recv() => {
-                         if let Err(err) = registration_tx.send(new_registration).await {
-                             error!(?err, "Error attempting to reprocess registration");
-                         }
+                        debug!(
+                            "Received a new registration for request_id: {} while already processing: {}",
+                            new_registration.0,
+                            auth_request_id
+                        );
+                        if let Err(err) = registration_tx.send(new_registration).await {
+                            error!(?err, "Error attempting to reprocess registration");
+                        }
                     }
                     // Registration successfully finished, send back the result.
                     result = registration.finish(&client, secret_store.as_ref()) => {
@@ -130,6 +149,18 @@ impl<T: PkceClient + Send + Sync + 'static> PkceState<T> {
                             error!(?err, "unknown error occurred finishing PKCE registration");
                         }
                     }
+                    _ = cancel_rx.recv() => {
+                        debug!("Cancelling current registration for request_id: {}", auth_request_id);
+                        if let Err(err) = pkce_state
+                            .finished_tx
+                            .lock()
+                            .await
+                            .send((auth_request_id, Err(fig_auth::Error::OAuthCustomError("cancelled".into()))))
+                            .await
+                        {
+                            error!(?err, "unknown error occurred cancelling PKCE registration");
+                        }
+                    }
                 }
             }
         });
@@ -176,6 +207,11 @@ impl<T: PkceClient + Send + Sync + 'static> PkceState<T> {
             None => Err(PkceError::RegistrationCancelled),
         }
     }
+
+    /// Cancel an ongoing PKCE registration, if one is present.
+    async fn cancel_registration(&self) {
+        self.cancel_tx.lock().await.try_send(()).ok();
+    }
 }
 
 pub async fn status(_request: AuthStatusRequest) -> RequestResult {
@@ -236,6 +272,11 @@ pub async fn finish_pkce_authorization(
     Ok(ServerOriginatedSubMessage::AuthFinishPkceAuthorizationResponse(AuthFinishPkceAuthorizationResponse {}).into())
 }
 
+pub async fn cancel_pkce_authorization(_: AuthCancelPkceAuthorizationRequest) -> RequestResult {
+    PKCE_REGISTRATION.cancel_registration().await;
+    Ok(ServerOriginatedSubMessage::AuthCancelPkceAuthorizationResponse(AuthCancelPkceAuthorizationResponse {}).into())
+}
+
 pub async fn builder_id_start_device_authorization(
     AuthBuilderIdStartDeviceAuthorizationRequest { start_url, region }: AuthBuilderIdStartDeviceAuthorizationRequest,
     ctx: &impl KVStore,
@@ -433,4 +474,40 @@ mod tests {
         let result_one = pkce_state.finish_registration(&id_one).await;
         assert!(matches!(result_one, Err(PkceError::RegistrationCancelled)));
     }
+
+    #[tokio::test]
+    async fn test_pkce_state_is_cancellable() {
+        tracing_subscriber::fmt::init();
+
+        let pkce_state = PkceState::new();
+
+        // Cancelling before we can finish the registration.
+        {
+            pkce_state.cancel_registration().await;
+            pkce_state.cancel_registration().await; // works multiple times
+
+            let client = TestPkceClient {};
+            let registration = PkceRegistration::register(&client, test_region(), test_issuer_url(), None)
+                .await
+                .unwrap();
+            let (uri, state) = (registration.redirect_uri.clone(), registration.state.clone());
+            let request_id = pkce_state.start_registration(registration, client, None).await;
+            send_test_auth_code(uri, state).await.unwrap();
+            assert!(pkce_state.finish_registration(&request_id).await.is_ok());
+        }
+
+        // Cancelling closes the server
+        {
+            let client = TestPkceClient {};
+            let registration = PkceRegistration::register(&client, test_region(), test_issuer_url(), None)
+                .await
+                .unwrap();
+            let (uri, state) = (registration.redirect_uri.clone(), registration.state.clone());
+            let _ = pkce_state.start_registration(registration, client, None).await;
+            // Give time for the HTTP server to be hosted
+            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
+            pkce_state.cancel_registration().await;
+            assert!(send_test_auth_code(uri, state).await.is_err());
+        }
+    }
 }

crates/fig_desktop_api/src/requests/auth.rs

@@ -269,11 +269,7 @@ pub fn host_sockets_dir() -> Result<PathBuf> {
 }
 
 /// The path to all of the themes
-pub fn themes_dir() -> Result<PathBuf> {
-    Ok(resources_path()?.join("themes"))
-}
-
-pub fn themes_dir_ctx(ctx: &Context) -> Result<PathBuf> {
+pub fn themes_dir(ctx: &Context) -> Result<PathBuf> {
     Ok(resources_path_ctx(ctx)?.join("themes"))
 }