UefiHidDxeV2: Fix Keyboard LED Sync Issue (#641)

HidPkg/UefiHidDxeV2/src/keyboard.rs

@@ -258,15 +258,13 @@ impl KeyboardHidHandler {
     fn generate_led_output_reports(&mut self) -> Vec<(Option<ReportId>, Vec<u8>)> {
         let mut output_vec = Vec::new();
         let current_leds: BTreeSet<Usage> = self.key_queue.active_leds().iter().cloned().collect();
-        if current_leds != self.led_state {
-            self.led_state = current_leds;
-            for output_builder in self.output_builders.clone() {
-                let mut report_buffer = vec![0u8; output_builder.report_size];
-                for field_builder in &output_builder.relevant_variable_fields {
-                    (field_builder.field_builder)(self, field_builder.field.clone(), report_buffer.as_mut_slice());
-                }
-                output_vec.push((output_builder.report_id, report_buffer));
+        self.led_state = current_leds;
+        for output_builder in self.output_builders.clone() {
+            let mut report_buffer = vec![0u8; output_builder.report_size];
+            for field_builder in &output_builder.relevant_variable_fields {
+                (field_builder.field_builder)(self, field_builder.field.clone(), report_buffer.as_mut_slice());
             }
+            output_vec.push((output_builder.report_id, report_buffer));
         }
         output_vec
     }
@@ -401,13 +399,22 @@ impl KeyboardHidHandler {
         self.current_keys.clear();
         self.key_queue.reset(extended_verification);
         if extended_verification {
-            self.update_leds(hid_io)?;
+            self.send_led_reports(hid_io)?;
         }
         Ok(())
     }
 
     /// Called to send LED state to the device if there has been a change in LEDs.
     pub fn update_leds(&mut self, hid_io: &dyn HidIo) -> Result<(), efi::Status> {
+        let current_leds: BTreeSet<Usage> = self.key_queue.active_leds().iter().cloned().collect();
+        if current_leds != self.led_state {
+            self.send_led_reports(hid_io)?;
+        }
+        Ok(())
+    }
+
+    /// Called to send LED state to the device.
+    pub fn send_led_reports(&mut self, hid_io: &dyn HidIo) -> Result<(), efi::Status> {
         for (id, output_report) in self.generate_led_output_reports() {
             let result = hid_io.set_output_report(id.map(|x| u32::from(x) as u8), &output_report);
             if let Err(result) = result {
@@ -436,7 +443,6 @@ impl KeyboardHidHandler {
     /// Sets the current key state.
     pub fn set_key_toggle_state(&mut self, toggle_state: u8) {
         self.key_queue.set_key_toggle_state(toggle_state);
-        self.generate_led_output_reports();
     }
 
     /// Registers a new key notify callback function to be invoked on the specified `key_data` press.
@@ -523,8 +529,6 @@ impl HidReportReceiver for KeyboardHidHandler {
     fn initialize(&mut self, controller: efi::Handle, hid_io: &dyn HidIo) -> Result<(), efi::Status> {
         let descriptor = hid_io.get_report_descriptor()?;
         self.process_descriptor(descriptor)?;
-        // Set the key toggle state here so that the subsequent reset() can send the LED state to the device.
-        self.set_key_toggle_state(protocols::simple_text_input_ex::CAPS_LOCK_ACTIVE);
         self.reset(hid_io, true)?;
         self.install_protocol_interfaces(controller)?;
         self.initialize_keyboard_layout()?;

HidPkg/UefiHidDxeV2/src/keyboard/key_queue.rs

@@ -376,7 +376,7 @@ impl KeyQueue {
 
     // returns a copy of the key at the front of the notify queue
     pub(crate) fn peek_notify_key(&self) -> Option<KeyData> {
-        self.key_queue.front().cloned()
+        self.notified_key_queue.front().cloned()
     }
 
     // set the key toggle state. This allows control of scroll/caps/num locks, as well as whether partial key state is

HidPkg/UefiHidDxeV2/src/keyboard/simple_text_in_ex.rs

@@ -35,7 +35,6 @@ use crate::{
 pub struct SimpleTextInExFfi {
     simple_text_in_ex: protocols::simple_text_input_ex::Protocol,
     boot_services: &'static dyn UefiBootServices,
-    key_notify_event: efi::Event,
     keyboard_handler: *mut KeyboardHidHandler,
 }
 
@@ -57,7 +56,6 @@ impl SimpleTextInExFfi {
                 unregister_key_notify: Self::simple_text_in_ex_unregister_key_notify,
             },
             boot_services,
-            key_notify_event: ptr::null_mut(),
             keyboard_handler: keyboard_handler as *mut KeyboardHidHandler,
         };
 
@@ -95,8 +93,6 @@ impl SimpleTextInExFfi {
             drop(unsafe { Box::from_raw(simple_text_in_ex_ptr) });
         }
 
-        unsafe { (*simple_text_in_ex_ptr).key_notify_event = key_notify_event };
-
         //install the simple_text_in_ex protocol
         let mut controller = controller;
         let status = boot_services.install_protocol_interface(
@@ -113,6 +109,8 @@ impl SimpleTextInExFfi {
             return Err(status);
         }
 
+        keyboard_handler.key_notify_event = key_notify_event;
+
         Ok(())
     }
 
@@ -174,8 +172,14 @@ impl SimpleTextInExFfi {
             return Err(status);
         }
 
-        let key_notify_event: efi::Handle = unsafe { (*simple_text_in_ex_ptr).key_notify_event };
-        let status = boot_services.close_event(key_notify_event);
+        let status = 'uninstall_processing: {
+            let Some(keyboard_handler) = (unsafe { (*simple_text_in_ex_ptr).keyboard_handler.as_mut() }) else {
+                break 'uninstall_processing efi::Status::DEVICE_ERROR;
+            };
+            let key_notify_event = keyboard_handler.key_notify_event;
+            keyboard_handler.key_notify_event = ptr::null_mut();
+            boot_services.close_event(key_notify_event)
+        };
         if status.is_error() {
             //An error here means the event was not closed, so in theory the notification_callback on it could still be
             //fired.
@@ -849,11 +853,12 @@ mod test {
 
         keyboard_handler.set_layout(Some(hii_keyboard_layout::get_default_keyboard_layout()));
         keyboard_handler.initialize(2 as efi::Handle, &hid_io).unwrap();
-        keyboard_handler.set_notify_event(NOTIFY_EVENT);
 
         SimpleTextInExFfi::install(boot_services, 2 as efi::Handle, &mut keyboard_handler).unwrap();
         assert_ne!(CONTEXT_PTR.load(Ordering::SeqCst), ptr::null_mut());
 
+        keyboard_handler.set_notify_event(NOTIFY_EVENT);
+
         extern "efiapi" fn key_notify_callback_a(
             key_data: *mut protocols::simple_text_input_ex::KeyData,
         ) -> efi::Status {