Be more careful with sending notification

rustica-agent-cli/src/main.rs

@@ -140,6 +140,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 &mut config.updatable_configuration,
                 &mut config.signatory,
                 &config.certificate_options,
+                &None,
             )
             .await
             {

rustica-agent/src/lib.rs

@@ -248,18 +248,12 @@ impl SshAgentHandler for Handler {
             // All other cases require us to fetch a certificate from one
             // of the configured servers
             _ => {
-                // Send a notification in case the user is going to need to
-                // tap their key to prove ownership before receiving their
-                // certificate
-                if let Some(f) = &self.notification_function {
-                    f()
-                }
-
                 // Fetch a new certificate from one of the servers
                 fetch_new_certificate(
                     &mut configuration,
                     &self.signatory,
                     &self.certificate_options,
+                    &self.notification_function,
                 )
                 .await
                 .map(|cert| {
@@ -332,7 +326,10 @@ impl SshAgentHandler for Handler {
             })?;
 
             if let Some(f) = &self.notification_function {
+                println!("Trying to send a notification");
                 f()
+            } else {
+                println!("No notification function set");
             }
 
             if let Some(pin) = &descriptor.pin {
@@ -402,6 +399,13 @@ impl SshAgentHandler for Handler {
 
         match private_key {
             Some(key) => {
+                if let Some(f) = &self.notification_function {
+                    debug!("Trying to send a notification");
+                    f()
+                } else {
+                    debug!("No notification function set");
+                }
+
                 let signature = match key.sign(&data) {
                     None => return Err(AgentError::from("Signing Error")),
                     Some(signature) => signature,
@@ -585,9 +589,13 @@ pub async fn fetch_new_certificate(
     configuration: &mut UpdatableConfiguration,
     signatory: &Signatory,
     options: &CertificateConfig,
+    notification_function: &Option<Box<dyn Fn() + Send + Sync>>,
 ) -> Result<Certificate, RusticaAgentLibraryError> {
     for server in configuration.get_servers_mut() {
-        match server.refresh_certificate_async(signatory, options).await {
+        match server
+            .refresh_certificate_async(signatory, options, notification_function)
+            .await
+        {
             Ok((cert, mtls_credentials)) => {
                 let parsed_cert = Certificate::from_string(&cert.cert)
                     .map_err(|e| RusticaAgentLibraryError::ServerReturnedInvalidCertificate(e))?;

rustica-agent/src/rustica/cert.rs

@@ -12,8 +12,10 @@ impl RusticaServer {
         &self,
         signatory: &Signatory,
         options: &CertificateConfig,
+        notification_function: &Option<Box<dyn Fn() + Send + Sync>>,
     ) -> Result<(RusticaCert, Option<MtlsCredentials>), RefreshError> {
-        let (mut client, challenge) = super::complete_rustica_challenge(self, signatory).await?;
+        let (mut client, challenge) =
+            super::complete_rustica_challenge(self, signatory, notification_function).await?;
 
         let current_timestamp = match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
             Ok(ts) => ts.as_secs(),
@@ -68,7 +70,11 @@ impl RusticaServer {
         signatory: &mut Signatory,
         options: &CertificateConfig,
         handle: &Handle,
+        notification_function: &Option<Box<dyn Fn() + Send + Sync>>,
     ) -> Result<(RusticaCert, Option<MtlsCredentials>), RefreshError> {
-        handle.block_on(async { self.refresh_certificate_async(signatory, options).await })
+        handle.block_on(async {
+            self.refresh_certificate_async(signatory, options, notification_function)
+                .await
+        })
     }
 }

rustica-agent/src/rustica/key.rs

@@ -28,7 +28,8 @@ impl RusticaServer {
         signatory: &mut Signatory,
         attestation: &PIVAttestation,
     ) -> Result<(), RefreshError> {
-        let (mut client, challenge) = super::complete_rustica_challenge(self, signatory).await?;
+        let (mut client, challenge) =
+            super::complete_rustica_challenge(self, signatory, &None).await?;
 
         let request = RegisterKeyRequest {
             certificate: attestation.certificate.clone(),
@@ -57,7 +58,8 @@ impl RusticaServer {
         application: &str,
         attestation: &U2FAttestation,
     ) -> Result<(), RefreshError> {
-        let (mut client, challenge) = super::complete_rustica_challenge(self, signatory).await?;
+        let (mut client, challenge) =
+            super::complete_rustica_challenge(self, signatory, &None).await?;
 
         let request = RegisterU2fKeyRequest {
             auth_data: attestation.auth_data.clone(),

rustica-agent/src/rustica/mod.rs

@@ -57,6 +57,7 @@ pub async fn get_rustica_client(
 pub async fn complete_rustica_challenge(
     server: &RusticaServer,
     signatory: &Signatory,
+    notification_function: &Option<Box<dyn Fn() + Send + Sync>>,
 ) -> Result<(RusticaClient<tonic::transport::Channel>, Challenge), RefreshError> {
     let ssh_pubkey = match signatory {
         Signatory::Yubikey(signer) => {
@@ -111,6 +112,12 @@ pub async fn complete_rustica_challenge(
         return Err(RefreshError::ServerChallengeNotForClientKey);
     }
 
+    // We need to sign the challenge so let's notify the user they
+    // will need to interact with their device if (if a device is being used)
+    if let Some(f) = notification_function {
+        f();
+    }
+
     let resigned_certificate = match signatory {
         Signatory::Yubikey(signer) => {
             let signature = signer