chore: add telemetry on auth failure

Author: brandonskiser

Cargo.lock

@@ -125,6 +125,21 @@
       "type": "string",
       "description": "The oauth authentication flow executed by the user, e.g. device code or PKCE"
     },
+    {
+      "name": "codewhispererterminal_authMethod",
+      "type": "string",
+      "description": "The authentication method used, e.g. BuilderId or IdentityCenter"
+    },
+    {
+      "name": "codewhispererterminal_errorType",
+      "type": "string",
+      "description": "The type of error, currently only used for authentication errors: TokenRefresh or NewLogin"
+    },
+    {
+      "name": "codewhispererterminal_errorCode",
+      "type": "string",
+      "description": "The specific error code from the backend service"
+    },
     {
       "name": "result",
       "type": "string",
@@ -244,6 +259,19 @@
         { "type": "codewhispererterminal_inCloudshell" }
       ]
     },
+    {
+      "name": "codewhispererterminal_authFailed",
+      "description": "Emitted when authentication fails",
+      "passive": false,
+      "metadata": [
+        { "type": "credentialStartUrl" },
+        { "type": "codewhispererterminal_inCloudshell" },
+        { "type": "codewhispererterminal_authMethod" },
+        { "type": "oauthFlow" },
+        { "type": "codewhispererterminal_errorType" },
+        { "type": "codewhispererterminal_errorCode", "required": false }
+      ]
+    },
     {
       "name": "codewhispererterminal_refreshCredentials",
       "description": "Emitted when users refresh their credentials",

crates/aws-toolkit-telemetry-definitions/def.json

@@ -45,10 +45,16 @@ use aws_types::request_id::RequestId;
 use aws_types::sdk_config::StalledStreamProtectionConfig;
 use fig_aws_common::app_name;
 use fig_telemetry_core::{
+    AuthErrorType,
     Event,
     EventType,
     TelemetryResult,
 };
+use fig_util::auth::{
+    OAuthFlow,
+    START_URL,
+    TokenType,
+};
 use time::OffsetDateTime;
 use tracing::{
     debug,
@@ -69,22 +75,6 @@ use crate::{
     Result,
 };
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
-pub enum OAuthFlow {
-    DeviceCode,
-    #[serde(alias = "Pkce")]
-    PKCE,
-}
-
-impl std::fmt::Display for OAuthFlow {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match *self {
-            OAuthFlow::DeviceCode => write!(f, "DeviceCode"),
-            OAuthFlow::PKCE => write!(f, "PKCE"),
-        }
-    }
-}
-
 /// Indicates if an expiration time has passed, there is a small 1 min window that is removed
 /// so the token will not expire in transit
 fn is_expired(expiration_time: &OffsetDateTime) -> bool {
@@ -254,30 +244,42 @@ pub async fn start_device_authorization(
         ..
     } = DeviceRegistration::init_device_code_registration(&client, secret_store, &region).await?;
 
+    let start_url = start_url.as_deref().unwrap_or(START_URL);
+
     let output = client
         .start_device_authorization()
         .client_id(&client_id)
         .client_secret(&client_secret.0)
-        .start_url(start_url.as_deref().unwrap_or(START_URL))
+        .start_url(start_url)
         .send()
-        .await?;
-
-    Ok(StartDeviceAuthorizationResponse {
-        device_code: output.device_code.unwrap_or_default(),
-        user_code: output.user_code.unwrap_or_default(),
-        verification_uri: output.verification_uri.unwrap_or_default(),
-        verification_uri_complete: output.verification_uri_complete.unwrap_or_default(),
-        expires_in: output.expires_in,
-        interval: output.interval,
-        region: region.to_string(),
-        start_url: start_url.unwrap_or_else(|| START_URL.to_owned()),
-    })
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum TokenType {
-    BuilderId,
-    IamIdentityCenter,
+        .await;
+
+    match output {
+        Ok(output) => Ok(StartDeviceAuthorizationResponse {
+            device_code: output.device_code.unwrap_or_default(),
+            user_code: output.user_code.unwrap_or_default(),
+            verification_uri: output.verification_uri.unwrap_or_default(),
+            verification_uri_complete: output.verification_uri_complete.unwrap_or_default(),
+            expires_in: output.expires_in,
+            interval: output.interval,
+            region: region.to_string(),
+            start_url: start_url.to_string(),
+        }),
+        Err(err) => {
+            let err: Error = err.into();
+            fig_telemetry_core::send_event(
+                Event::new(EventType::AuthFailed {
+                    auth_method: TokenType::from_start_url(Some(start_url)),
+                    oauth_flow: OAuthFlow::DeviceCode,
+                    error_type: AuthErrorType::NewLogin,
+                    error_code: err.service_error_code(),
+                })
+                .with_credential_start_url(start_url.to_string()),
+            )
+            .await;
+            Err(err)
+        },
+    }
 }
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
@@ -512,6 +514,27 @@ pub async fn poll_create_token(
     device_code: String,
     start_url: Option<String>,
     region: Option<String>,
+) -> PollCreateToken {
+    match poll_create_token_impl(secret_store, device_code, start_url.clone(), region).await {
+        PollCreateToken::Error(err) => {
+            fig_telemetry_core::send_event(Event::new(EventType::AuthFailed {
+                auth_method: TokenType::from_start_url(start_url.as_deref()),
+                oauth_flow: OAuthFlow::DeviceCode,
+                error_type: AuthErrorType::NewLogin,
+                error_code: err.service_error_code(),
+            }))
+            .await;
+            PollCreateToken::Error(err)
+        },
+        other => other,
+    }
+}
+
+async fn poll_create_token_impl(
+    secret_store: &SecretStore,
+    device_code: String,
+    start_url: Option<String>,
+    region: Option<String>,
 ) -> PollCreateToken {
     let region = region.clone().map_or(OIDC_BUILDER_ID_REGION, Region::new);
     let client = client(region.clone());
@@ -634,23 +657,6 @@ mod tests {
     const US_EAST_1: Region = Region::from_static("us-east-1");
     const US_WEST_2: Region = Region::from_static("us-west-2");
 
-    macro_rules! test_ser_deser {
-        ($ty:ident, $variant:expr, $text:expr) => {
-            let quoted = format!("\"{}\"", $text);
-            assert_eq!(quoted, serde_json::to_string(&$variant).unwrap());
-            assert_eq!($variant, serde_json::from_str(&quoted).unwrap());
-
-            assert_eq!($text, format!("{}", $variant));
-        };
-    }
-
-    #[test]
-    fn test_oauth_flow_ser_deser() {
-        test_ser_deser!(OAuthFlow, OAuthFlow::DeviceCode, "DeviceCode");
-        test_ser_deser!(OAuthFlow, OAuthFlow::PKCE, "PKCE");
-        assert_eq!(OAuthFlow::PKCE, serde_json::from_str("\"Pkce\"").unwrap());
-    }
-
     #[test]
     fn test_client() {
         println!("{:?}", client(US_EAST_1));

crates/fig_auth/src/builder_id.rs

@@ -18,9 +18,6 @@ pub(crate) const SCOPES: &[&str] = &[
 
 pub(crate) const CLIENT_TYPE: &str = "public";
 
-// The start URL for public builder ID users
-pub const START_URL: &str = "https://view.awsapps.com/start";
-
 // The start URL for internal amzn users
 pub const AMZN_START_URL: &str = "https://amzn.awsapps.com/start";
 

crates/fig_auth/src/consts.rs

@@ -10,8 +10,6 @@ use thiserror::Error;
 #[allow(variant_size_differences)]
 #[derive(Debug, Error)]
 pub enum Error {
-    #[error(transparent)]
-    Ssooidc(#[from] Box<aws_sdk_ssooidc::Error>),
     #[error(transparent)]
     SdkRegisterClient(#[from] SdkError<RegisterClientError>),
     #[error(transparent)]
@@ -53,13 +51,27 @@ pub enum Error {
 impl Error {
     pub fn to_verbose_string(&self) -> String {
         match self {
-            Error::Ssooidc(s) => DisplayErrorContext(s).to_string(),
             Error::SdkRegisterClient(s) => DisplayErrorContext(s).to_string(),
             Error::SdkCreateToken(s) => DisplayErrorContext(s).to_string(),
             Error::SdkStartDeviceAuthorization(s) => DisplayErrorContext(s).to_string(),
             other => other.to_string(),
         }
     }
+
+    pub fn service_error_code(&self) -> Option<String> {
+        match self {
+            Error::SdkRegisterClient(err) => err
+                .as_service_error()
+                .and_then(|e| e.meta().code().map(|s| s.to_string())),
+            Error::SdkCreateToken(err) => err
+                .as_service_error()
+                .and_then(|e| e.meta().code().map(|s| s.to_string())),
+            Error::SdkStartDeviceAuthorization(err) => err
+                .as_service_error()
+                .and_then(|e| e.meta().code().map(|s| s.to_string())),
+            _ => None,
+        }
+    }
 }
 
 pub(crate) type Result<T, E = Error> = std::result::Result<T, E>;

crates/fig_auth/src/error.rs

@@ -12,9 +12,6 @@ pub use builder_id::{
     logout,
     refresh_token,
 };
-pub use consts::{
-    AMZN_START_URL,
-    START_URL,
-};
+pub use consts::AMZN_START_URL;
 pub use error::Error;
 pub(crate) use error::Result;

crates/fig_auth/src/lib.rs

@@ -30,6 +30,16 @@ pub use aws_types::region::Region;
 use base64::Engine;
 use base64::engine::general_purpose::URL_SAFE;
 use bytes::Bytes;
+use fig_telemetry_core::{
+    AuthErrorType,
+    Event,
+    EventType,
+};
+use fig_util::auth::{
+    OAuthFlow,
+    START_URL,
+    TokenType,
+};
 use http_body_util::Full;
 use hyper::body::Incoming;
 use hyper::server::conn::http1;
@@ -56,7 +66,6 @@ use crate::secret_store::SecretStore;
 use crate::{
     Error,
     Result,
-    START_URL,
 };
 
 const DEFAULT_AUTHORIZATION_TIMEOUT: Duration = Duration::from_secs(60 * 3);
@@ -70,8 +79,22 @@ pub async fn start_pkce_authorization(
     let issuer_url = start_url.as_deref().unwrap_or(START_URL);
     let region = region.clone().map_or(OIDC_BUILDER_ID_REGION, Region::new);
     let client = client(region.clone());
-    let registration = PkceRegistration::register(&client, region, issuer_url.to_string(), None).await?;
-    Ok((client, registration))
+    match PkceRegistration::register(&client, region, issuer_url.to_string(), None).await {
+        Ok(registration) => Ok((client, registration)),
+        Err(err) => {
+            fig_telemetry_core::send_event(
+                Event::new(EventType::AuthFailed {
+                    auth_method: TokenType::from_start_url(Some(issuer_url)),
+                    oauth_flow: OAuthFlow::PKCE,
+                    error_type: AuthErrorType::NewLogin,
+                    error_code: err.service_error_code(),
+                })
+                .with_credential_start_url(issuer_url.to_string()),
+            )
+            .await;
+            Err(err)
+        },
+    }
 }
 
 /// Represents a client used for registering with AWS IAM OIDC.
@@ -242,11 +265,28 @@ impl PkceRegistration {
                 code_verifier: self.code_verifier,
                 code,
             })
-            .await?;
+            .await;
 
         // Tokens are redacted in the log output.
         debug!(?response, "Received create_token response");
 
+        let response = match response {
+            Ok(res) => res,
+            Err(err) => {
+                fig_telemetry_core::send_event(
+                    Event::new(EventType::AuthFailed {
+                        auth_method: TokenType::from_start_url(Some(&self.issuer_url)),
+                        oauth_flow: OAuthFlow::PKCE,
+                        error_type: AuthErrorType::NewLogin,
+                        error_code: err.service_error_code(),
+                    })
+                    .with_credential_start_url(self.issuer_url),
+                )
+                .await;
+                return Err(err);
+            },
+        };
+
         let token = BuilderIdToken::from_output(
             response.output,
             self.region.clone(),

crates/fig_auth/src/pkce.rs

@@ -6,7 +6,6 @@ use std::sync::{
 use fig_auth::builder_id::{
     PollCreateToken,
     StartDeviceAuthorizationResponse,
-    TokenType,
 };
 use fig_auth::pkce::{
     Client,
@@ -31,6 +30,7 @@ use fig_proto::fig::{
     AuthStatusRequest,
     AuthStatusResponse,
 };
+use fig_util::auth::TokenType;
 use tokio::sync::Mutex;
 use tokio::sync::mpsc::{
     Receiver,

crates/fig_desktop_api/src/requests/auth.rs

@@ -16,3 +16,6 @@ fig_util.workspace = true
 serde.workspace = true
 strum.workspace = true
 tokio.workspace = true
+
+[dev-dependencies]
+serde_json.workspace = true