Merge branch 'main' into 1.10-fixes

Author: chaynabors

crates/chat-cli/src/auth/builder_id.rs

@@ -52,10 +52,9 @@ use crate::auth::AuthError;
 use crate::auth::consts::*;
 use crate::auth::scope::is_scopes;
 use crate::aws_common::app_name;
-use crate::database::Database;
-use crate::database::secret_store::{
+use crate::database::{
+    Database,
     Secret,
-    SecretStore,
 };
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
@@ -123,8 +122,8 @@ impl DeviceRegistration {
     }
 
     /// Loads the OIDC registered client from the secret store, deleting it if it is expired.
-    async fn load_from_secret_store(secret_store: &SecretStore, region: &Region) -> Result<Option<Self>, AuthError> {
-        let device_registration = secret_store.get(Self::SECRET_KEY).await?;
+    async fn load_from_secret_store(database: &Database, region: &Region) -> Result<Option<Self>, AuthError> {
+        let device_registration = database.get_secret(Self::SECRET_KEY).await?;
 
         if let Some(device_registration) = device_registration {
             // check that the data is not expired, assume it is invalid if not present
@@ -138,7 +137,7 @@ impl DeviceRegistration {
         }
 
         // delete the data if its expired or invalid
-        if let Err(err) = secret_store.delete(Self::SECRET_KEY).await {
+        if let Err(err) = database.delete_secret(Self::SECRET_KEY).await {
             error!(?err, "Failed to delete device registration from keychain");
         }
 
@@ -152,7 +151,7 @@ impl DeviceRegistration {
         client: &Client,
         region: &Region,
     ) -> Result<Self, AuthError> {
-        match Self::load_from_secret_store(&database.secret_store, region).await {
+        match Self::load_from_secret_store(database, region).await {
             Ok(Some(registration)) if registration.oauth_flow == OAuthFlow::DeviceCode => match &registration.scopes {
                 Some(scopes) if is_scopes(scopes) => return Ok(registration),
                 _ => warn!("Invalid scopes in device registration, ignoring"),
@@ -181,17 +180,17 @@ impl DeviceRegistration {
             SCOPES.iter().map(|s| (*s).to_owned()).collect(),
         );
 
-        if let Err(err) = device_registration.save(&database.secret_store).await {
+        if let Err(err) = device_registration.save(database).await {
             error!(?err, "Failed to write device registration to keychain");
         }
 
         Ok(device_registration)
     }
 
     /// Saves to the passed secret store.
-    pub async fn save(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
+    pub async fn save(&self, secret_store: &Database) -> Result<(), AuthError> {
         secret_store
-            .set(Self::SECRET_KEY, &serde_json::to_string(&self)?)
+            .set_secret(Self::SECRET_KEY, &serde_json::to_string(&self)?)
             .await?;
         Ok(())
     }
@@ -285,8 +284,8 @@ impl BuilderIdToken {
     }
 
     /// Load the token from the keychain, refresh the token if it is expired and return it
-    pub async fn load(database: &mut Database) -> Result<Option<Self>, AuthError> {
-        match database.secret_store.get(Self::SECRET_KEY).await {
+    pub async fn load(database: &Database) -> Result<Option<Self>, AuthError> {
+        match database.get_secret(Self::SECRET_KEY).await {
             Ok(Some(secret)) => {
                 let token: Option<Self> = serde_json::from_str(&secret.0)?;
                 match token {
@@ -296,7 +295,7 @@ impl BuilderIdToken {
                         let client = client(region.clone());
                         // if token is expired try to refresh
                         if token.is_expired() {
-                            token.refresh_token(&client, &database.secret_store, &region).await
+                            token.refresh_token(&client, database, &region).await
                         } else {
                             Ok(Some(token))
                         }
@@ -316,19 +315,19 @@ impl BuilderIdToken {
     pub async fn refresh_token(
         &self,
         client: &Client,
-        secret_store: &SecretStore,
+        database: &Database,
         region: &Region,
     ) -> Result<Option<Self>, AuthError> {
         let Some(refresh_token) = &self.refresh_token else {
             // if the token is expired and has no refresh token, delete it
-            if let Err(err) = self.delete(secret_store).await {
+            if let Err(err) = self.delete(database).await {
                 error!(?err, "Failed to delete builder id token");
             }
 
             return Ok(None);
         };
 
-        let registration = match DeviceRegistration::load_from_secret_store(secret_store, region).await? {
+        let registration = match DeviceRegistration::load_from_secret_store(database, region).await? {
             Some(registration) if registration.oauth_flow == self.oauth_flow => registration,
             // If the OIDC client registration is for a different oauth flow or doesn't exist, then
             // we can't refresh the token.
@@ -365,7 +364,7 @@ impl BuilderIdToken {
                 );
                 debug!("Refreshed access token, new token: {:?}", token);
 
-                if let Err(err) = token.save(secret_store).await {
+                if let Err(err) = token.save(database).await {
                     error!(?err, "Failed to store builder id access token");
                 };
 
@@ -378,7 +377,7 @@ impl BuilderIdToken {
                 // if the error is the client's fault, clear the token
                 if let SdkError::ServiceError(service_err) = &err {
                     if !service_err.err().is_slow_down_exception() {
-                        if let Err(err) = self.delete(secret_store).await {
+                        if let Err(err) = self.delete(database).await {
                             error!(?err, "Failed to delete builder id token");
                         }
                     }
@@ -398,16 +397,16 @@ impl BuilderIdToken {
     }
 
     /// Save the token to the keychain
-    pub async fn save(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
-        secret_store
-            .set(Self::SECRET_KEY, &serde_json::to_string(self)?)
+    pub async fn save(&self, database: &Database) -> Result<(), AuthError> {
+        database
+            .set_secret(Self::SECRET_KEY, &serde_json::to_string(self)?)
             .await?;
         Ok(())
     }
 
     /// Delete the token from the keychain
-    pub async fn delete(&self, secret_store: &SecretStore) -> Result<(), AuthError> {
-        secret_store.delete(Self::SECRET_KEY).await?;
+    pub async fn delete(&self, database: &Database) -> Result<(), AuthError> {
+        database.delete_secret(Self::SECRET_KEY).await?;
         Ok(())
     }
 
@@ -479,7 +478,7 @@ pub async fn poll_create_token(
             let token: BuilderIdToken =
                 BuilderIdToken::from_output(output, region, start_url, OAuthFlow::DeviceCode, scopes);
 
-            if let Err(err) = token.save(&database.secret_store).await {
+            if let Err(err) = token.save(database).await {
                 error!(?err, "Failed to store builder id token");
             };
 
@@ -500,13 +499,13 @@ pub async fn is_logged_in(database: &mut Database) -> bool {
 }
 
 pub async fn logout(database: &mut Database) -> Result<(), AuthError> {
-    let Ok(secret_store) = SecretStore::new().await else {
+    let Ok(secret_store) = Database::new().await else {
         return Ok(());
     };
 
     let (builder_res, device_res) = tokio::join!(
-        secret_store.delete(BuilderIdToken::SECRET_KEY),
-        secret_store.delete(DeviceRegistration::SECRET_KEY),
+        secret_store.delete_secret(BuilderIdToken::SECRET_KEY),
+        secret_store.delete_secret(DeviceRegistration::SECRET_KEY),
     );
 
     let profile_res = database.unset_auth_profile();

crates/chat-cli/src/auth/pkce.rs

@@ -231,7 +231,7 @@ impl PkceRegistration {
     /// then the access and refresh tokens will be saved.
     ///
     /// Only the first connection will be served.
-    pub async fn finish<C: PkceClient>(self, client: &C, database: Option<&Database>) -> Result<(), AuthError> {
+    pub async fn finish<C: PkceClient>(self, client: &C, database: Option<&mut Database>) -> Result<(), AuthError> {
         let code = tokio::select! {
             code = Self::recv_code(self.listener, self.state) => {
                 code?
@@ -270,11 +270,11 @@ impl PkceRegistration {
         );
 
         if let Some(database) = database {
-            if let Err(err) = device_registration.save(&database.secret_store).await {
+            if let Err(err) = device_registration.save(database).await {
                 error!(?err, "Failed to store pkce registration to secret store");
             }
 
-            if let Err(err) = token.save(&database.secret_store).await {
+            if let Err(err) = token.save(database).await {
                 error!(?err, "Failed to store builder id token");
             };
         }

crates/chat-cli/src/cli/debug.rs

@@ -1,7 +1,4 @@
-use clap::{
-    Subcommand,
-    ValueEnum,
-};
+use clap::ValueEnum;
 
 #[derive(Debug, ValueEnum, Clone, PartialEq, Eq)]
 pub enum Build {
@@ -65,7 +62,7 @@ pub enum TISAction {
 use std::path::PathBuf;
 
 #[cfg(target_os = "macos")]
-#[derive(Debug, Subcommand, Clone, PartialEq, Eq)]
+#[derive(Debug, clap::Subcommand, Clone, PartialEq, Eq)]
 pub enum InputMethodDebugAction {
     Install {
         bundle_path: Option<PathBuf>,

crates/chat-cli/src/database/mod.rs

@@ -1,4 +1,3 @@
-pub mod secret_store;
 pub mod settings;
 
 use std::ops::Deref;
@@ -17,7 +16,6 @@ use rusqlite::{
     ToSql,
     params,
 };
-use secret_store::SecretStore;
 use serde::de::DeserializeOwned;
 use serde::{
     Deserialize,
@@ -94,6 +92,25 @@ impl From<amzn_codewhisperer_client::types::Profile> for AuthProfile {
     }
 }
 
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
+#[serde(transparent)]
+pub struct Secret(pub String);
+
+impl std::fmt::Debug for Secret {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Secret").finish()
+    }
+}
+
+impl<T> From<T> for Secret
+where
+    T: Into<String>,
+{
+    fn from(value: T) -> Self {
+        Self(value.into())
+    }
+}
+
 // A cloneable error
 #[derive(Debug, Clone, thiserror::Error)]
 #[error("Failed to open database: {}", .0)]
@@ -117,9 +134,6 @@ pub enum DatabaseError {
     DbOpenError(#[from] DbOpenError),
     #[error("{}", .0)]
     PoisonError(String),
-    #[cfg(target_os = "macos")]
-    #[error("Security error: {}", .0)]
-    Security(String),
     #[error(transparent)]
     StringFromUtf8(#[from] std::string::FromUtf8Error),
     #[error(transparent)]
@@ -140,8 +154,7 @@ pub enum Table {
     State,
     /// The conversations tables contains user chat conversations.
     Conversations,
-    #[cfg(not(target_os = "macos"))]
-    /// The auth table contains
+    /// The auth table contains SSO and Builder ID credentials.
     Auth,
 }
 
@@ -150,7 +163,6 @@ impl std::fmt::Display for Table {
         match self {
             Table::State => write!(f, "state"),
             Table::Conversations => write!(f, "conversations"),
-            #[cfg(not(target_os = "macos"))]
             Table::Auth => write!(f, "auth_kv"),
         }
     }
@@ -166,7 +178,6 @@ struct Migration {
 pub struct Database {
     pool: Pool<SqliteConnectionManager>,
     pub settings: Settings,
-    pub secret_store: SecretStore,
 }
 
 impl Database {
@@ -176,7 +187,6 @@ impl Database {
                 return Self {
                     pool: Pool::builder().build(SqliteConnectionManager::memory()).unwrap(),
                     settings: Settings::new().await?,
-                    secret_store: SecretStore::new().await?,
                 }
                 .migrate();
             },
@@ -209,7 +219,6 @@ impl Database {
         Ok(Self {
             pool,
             settings: Settings::new().await?,
-            secret_store: SecretStore::new().await?,
         }
         .migrate()
         .map_err(|e| DbOpenError(e.to_string()))?)
@@ -419,6 +428,19 @@ impl Database {
 
         Ok(map)
     }
+
+    pub async fn get_secret(&self, key: &str) -> Result<Option<Secret>, DatabaseError> {
+        Ok(self.get_entry::<String>(Table::Auth, key)?.map(Into::into))
+    }
+
+    pub async fn set_secret(&self, key: &str, value: &str) -> Result<(), DatabaseError> {
+        self.set_entry(Table::Auth, key, value)?;
+        Ok(())
+    }
+
+    pub async fn delete_secret(&self, key: &str) -> Result<(), DatabaseError> {
+        self.delete_entry(Table::Auth, key)
+    }
 }
 
 fn max_migration<C: Deref<Target = Connection>>(conn: &C) -> Option<i64> {
@@ -502,4 +524,43 @@ mod tests {
         assert!(db.get_entry::<f32>(Table::State, "float").unwrap().is_some());
         assert!(db.get_entry::<bool>(Table::State, "bool").unwrap().is_some());
     }
+
+    #[tokio::test]
+    #[ignore = "not on ci"]
+    async fn test_set_password() {
+        let key = "test_set_password";
+        let store = Database::new().await.unwrap();
+        store.set_secret(key, "test").await.unwrap();
+        assert_eq!(store.get_secret(key).await.unwrap().unwrap().0, "test");
+        store.delete_secret(key).await.unwrap();
+    }
+
+    #[tokio::test]
+    #[ignore = "not on ci"]
+    async fn secret_get_time() {
+        let key = "test_secret_get_time";
+        let store = Database::new().await.unwrap();
+        store.set_secret(key, "1234").await.unwrap();
+
+        let now = std::time::Instant::now();
+        for _ in 0..100 {
+            store.get_secret(key).await.unwrap();
+        }
+
+        println!("duration: {:?}", now.elapsed() / 100);
+
+        store.delete_secret(key).await.unwrap();
+    }
+
+    #[tokio::test]
+    #[ignore = "not on ci"]
+    async fn secret_delete() {
+        let key = "test_secret_delete";
+
+        let store = Database::new().await.unwrap();
+        store.set_secret(key, "1234").await.unwrap();
+        assert_eq!(store.get_secret(key).await.unwrap().unwrap().0, "1234");
+        store.delete_secret(key).await.unwrap();
+        assert_eq!(store.get_secret(key).await.unwrap(), None);
+    }
 }