Repository: HaoboGu/rmk (language: rust)

Issue:
After the split keyboard's slave device disconnects and reconnects, the slave device has no key output.
RMK Version: 0.4.4
Microcontroller: rp2040 4Mb
Platform: GNU/Linux


Discussion / hints:

Fix patch:
diff --git a/rmk/src/ble/esp/mod.rs b/rmk/src/ble/esp/mod.rs
index 5837de4f..e7e6c013 100644
--- a/rmk/src/ble/esp/mod.rs
+++ b/rmk/src/ble/esp/mod.rs
@@ -9,6 +9,7 @@ use crate::storage::Storage;
 use crate::via::process::VialService;
 use crate::via::vial_task;
 use crate::KEYBOARD_STATE;
+use crate::CONNECTION_STATE;
 use crate::{
     action::KeyAction, ble::ble_communication_task, config::RmkConfig, keyboard::Keyboard,
     keymap::KeyMap,
@@ -75,6 +76,7 @@ pub(crate) async fn initialize_esp_ble_keyboard_with_config_and_run<
     let mut vial_service = VialService::new(&keymap, keyboard_config.vial_config);
     loop {
         KEYBOARD_STATE.store(false, core::sync::atomic::Ordering::Release);
+        CONNECTION_STATE.store(false, core::sync::atomic::Ordering::Release);
         info!("Advertising..");
         let mut ble_server = BleServer::new(keyboard_config.usb_config);
         ble_server.output_keyboard.lock().on_write(|args| {
@@ -86,6 +88,7 @@ pub(crate) async fn initialize_esp_ble_keyboard_with_config_and_run<
         ble_server.wait_for_connection().await;
 
         info!("BLE connected!");
+        CONNECTION_STATE.store(true, core::sync::atomic::Ordering::Release);
 
         // Create BLE HID writers
         let mut keyboard_writer = ble_server.input_keyboard;
@@ -114,7 +117,7 @@ pub(crate) async fn initialize_esp_ble_keyboard_with_config_and_run<
             hid_writer: ble_server.input_vial,
         };
         let via_fut = vial_task(&mut via_rw, &mut vial_service);
-        let matrix_fut = matrix.scan();
+        let matrix_fut = matrix.run();
         let storage_fut = storage.run();
         pin_mut!(storage_fut);
         pin_mut!(via_fut);
diff --git a/rmk/src/ble/mod.rs b/rmk/src/ble/mod.rs
index 965ebd95..d79828a4 100644
--- a/rmk/src/ble/mod.rs
+++ b/rmk/src/ble/mod.rs
@@ -16,6 +16,7 @@ use crate::{
     hid::HidWriterWrapper,
     keyboard::{write_other_report_to_host, KeyboardReportMessage, REPORT_CHANNEL_SIZE},
     usb::descriptor::CompositeReportType,
+    CONNECTION_STATE,
 };
 
 /// BLE communication task, send reports to host via BLE.
@@ -41,31 +42,39 @@ pub(crate) async fn ble_communication_task<
     // Wait 1 seconds, ensure that gatt server has been started
     Timer::after_secs(1).await;
     loop {
-        match keyboard_report_receiver.receive().await {
-            KeyboardReportMessage::KeyboardReport(report) => {
-                debug!(
-                    "Send keyboard report via BLE: {=[u8]:#X}, modifier: {:b}",
-                    report.keycodes, report.modifier
-                );
-                match ble_keyboard_writer.write_serialize(&report).await {
-                    Ok(()) => {}
-                    Err(e) => error!("Send keyboard report error: {}", e),
-                };
-            }
-            KeyboardReportMessage::CompositeReport(report, report_type) => {
-                match report_type {
-                    CompositeReportType::Media => {
-                        write_other_report_to_host(report, report_type, ble_media_writer).await
-                    }
-                    CompositeReportType::Mouse => {
-                        write_other_report_to_host(report, report_type, ble_mouse_writer).await
-                    }
-                    CompositeReportType::System => {
-                        write_other_report_to_host(report, report_type, ble_system_control_writer)
+        let report = keyboard_report_receiver.receive().await;
+        // Only send the report after the connection is established.
+        if CONNECTION_STATE.load(core::sync::atomic::Ordering::Acquire) {
+            match report {
+                KeyboardReportMessage::KeyboardReport(report) => {
+                    debug!(
+                        "Send keyboard report via BLE: {=[u8]:#X}, modifier: {:b}",
+                        report.keycodes, report.modifier
+                    );
+                    match ble_keyboard_writer.write_serialize(&report).await {
+                        Ok(()) => {}
+                        Err(e) => error!("Send keyboard report error: {}", e),
+                    };
+                }
+                KeyboardReportMessage::CompositeReport(report, report_type) => {
+                    match report_type {
+                        CompositeReportType::Media => {
+                            write_other_report_to_host(report, report_type, ble_media_writer).await
+                        }
+                        CompositeReportType::Mouse => {
+                            write_other_report_to_host(report, report_type, ble_mouse_writer).await
+                        }
+                        CompositeReportType::System => {
+                            write_other_report_to_host(
+                                report,
+                                report_type,
+                                ble_system_control_writer,
+                            )
                             .await
-                    }
-                    CompositeReportType::None => (),
-                };
+                        }
+                        CompositeReportType::None => (),
+                    };
+                }
             }
         }
     }
diff --git a/rmk/src/ble/nrf/bonder.rs b/rmk/src/ble/nrf/bonder.rs
index 99fa7e37..f906c06d 100644
--- a/rmk/src/ble/nrf/bonder.rs
+++ b/rmk/src/ble/nrf/bonder.rs
@@ -2,6 +2,7 @@ use super::BONDED_DEVICE_NUM;
 use crate::{
     ble::nrf::ACTIVE_PROFILE,
     storage::{FlashOperationMessage, FLASH_CHANNEL},
+    CONNECTION_STATE,
 };
 use core::{cell::RefCell, sync::atomic::Ordering};
 use defmt::{debug, error, info, warn, Format};
@@ -112,6 +113,8 @@ impl SecurityHandler for MultiBonder {
 
     fn on_security_update(&self, _conn: &Connection, security_mode: SecurityMode) {
         info!("on_security_update, new security mode: {}", security_mode);
+        // Security updated, indicating that the connection is established?
+        CONNECTION_STATE.store(true, Ordering::Release);
     }
 
     fn on_bonded(
@@ -158,11 +161,12 @@ impl SecurityHandler for MultiBonder {
         }
     }
 
-    fn get_key(&self, _conn: &Connection, master_id: MasterId) -> Option<EncryptionInfo> {
+    fn get_key(&self, conn: &Connection, master_id: MasterId) -> Option<EncryptionInfo> {
         // Reconnecting with an existing bond
         debug!("Getting bond for {}", master_id);
 
-        self.bond_info
+        let info = self
+            .bond_info
             .borrow()
             .iter()
             .find(|(_, info)| {
@@ -172,7 +176,10 @@ impl SecurityHandler for MultiBonder {
                     && info.peer.master_id == master_id
                     && info.removed == false
             })
-            .and_then(|(_, d)| Some(d.peer.key))
+            .and_then(|(_, d)| Some(d.peer.key));
+
+        self.load_sys_attrs(conn);
+        info
     }
 
     fn save_sys_attrs(&self, conn: &Connection) {
@@ -341,15 +348,18 @@ impl SecurityHandler for Bonder {
         }
     }
 
-    fn get_key(&self, _conn: &Connection, master_id: MasterId) -> Option<EncryptionInfo> {
+    fn get_key(&self, conn: &Connection, master_id: MasterId) -> Option<EncryptionInfo> {
         // Reconnecting with an existing bond
         debug!("Getting bond for {}", master_id);
 
-        self.bond_info
+        let info = self
+            .bond_info
             .borrow()
             .iter()
             .find(|(_, info)| info.peer.master_id == master_id && info.removed == false)
-            .and_then(|(_, d)| Some(d.peer.key))
+            .and_then(|(_, d)| Some(d.peer.key));
+        self.load_sys_attrs(conn);
+        info
     }
 
     fn save_sys_attrs(&self, conn: &Connection) {
diff --git a/rmk/src/ble/nrf/mod.rs b/rmk/src/ble/nrf/mod.rs
index 4d30f6e4..408b5be0 100644
--- a/rmk/src/ble/nrf/mod.rs
+++ b/rmk/src/ble/nrf/mod.rs
@@ -13,7 +13,6 @@ use crate::config::BleBatteryConfig;
 use crate::keyboard::{keyboard_report_channel, REPORT_CHANNEL_SIZE};
 use crate::matrix::MatrixTrait;
 use crate::storage::StorageKeys;
-use crate::KEYBOARD_STATE;
 use crate::{
     ble::{
         ble_communication_task,
@@ -28,6 +27,7 @@ use crate::{
     storage::{get_bond_info_key, Storage, StorageData},
     vial_task, KeyAction, KeyMap, LightService, RmkConfig, VialService, CONNECTION_TYPE,
 };
+use crate::{CONNECTION_STATE, KEYBOARD_STATE};
 use bonder::MultiBonder;
 use core::sync::atomic::{AtomicU8, Ordering};
 use core::{cell::RefCell, mem};
@@ -531,6 +531,8 @@ pub(crate) async fn run_dummy_keyboard<
         REPORT_CHANNEL_SIZE,
     >,
 ) {
+    CONNECTION_STATE.store(false, Ordering::Release);
+    // Don't need to wait for connection, just do scanning to detect if there's a profile update
     let matrix_fut = matrix.scan();
     let keyboard_fut = keyboard.run();
     let storage_fut = storage.run();
@@ -587,6 +589,7 @@ pub(crate) async fn run_ble_keyboard<
         REPORT_CHANNEL_SIZE,
     >,
 ) {
+    CONNECTION_STATE.store(false, Ordering::Release);
     info!("Starting GATT server 20 ms later");
     Timer::after_millis(20).await;
     let mut ble_keyboard_writer = BleHidWriter::<'_, 8>::new(&conn, ble_server.hid.input_keyboard);
@@ -601,7 +604,7 @@ pub(crate) async fn run_ble_keyboard<
     // Tasks
     let battery_fut = bas.run(battery_config, &conn);
     let led_fut = led_service_task(light_service);
-    let matrix_fut = matrix.scan();
+    let matrix_fut = matrix.run();
     // Run the GATT server on the connection. This returns when the connection gets disconnected.
     let ble_fut = gatt_server::run(&conn, ble_server, |_| {});
     let keyboard_fut = keyboard.run();
diff --git a/rmk/src/ble/nrf/server.rs b/rmk/src/ble/nrf/server.rs
index addbdef9..2b23eeb7 100644
--- a/rmk/src/ble/nrf/server.rs
+++ b/rmk/src/ble/nrf/server.rs
@@ -175,7 +175,6 @@ impl gatt_server::Server for BleServer {
                 }
             }
         }
-
         if let Some(event) = self.vial.on_write(handle, data) {
             match event {
                 VialServiceEvent::InputVialKeyCccdWrite => {
diff --git a/rmk/src/keyboard.rs b/rmk/src/keyboard.rs
index a29de15e..35bad243 100644
--- a/rmk/src/keyboard.rs
+++ b/rmk/src/keyboard.rs
@@ -1,4 +1,5 @@
 use crate::config::BehaviorConfig;
+use crate::CONNECTION_STATE;
 use crate::{
     action::{Action, KeyAction},
     hid::{ConnectionType, HidWriterWrapper},
@@ -82,15 +83,19 @@ pub(crate) async fn communication_task<'a, W: HidWriterWrapper, W2: HidWriterWra
     // This delay is necessary otherwise this task will stuck at the first send when the USB is suspended
     Timer::after_secs(2).await;
     loop {
-        match receiver.receive().await {
-            KeyboardReportMessage::KeyboardReport(report) => {
-                match keybooard_hid_writer.write_serialize(&report).await {
-                    Ok(()) => {}
-                    Err(e) => error!("Send keyboard report error: {}", e),
-                };
-            }
-            KeyboardReportMessage::CompositeReport(report, report_type) => {
-                write_other_report_to_host(report, report_type, other_hid_writer).await;
+        let report = receiver.receive().await;
+        // Only send the report after the connection is established.
+        if CONNECTION_STATE.load(core::sync::atomic::Ordering::Acquire) {
+            match report {
+                KeyboardReportMessage::KeyboardReport(report) => {
+                    match keybooard_hid_writer.write_serialize(&report).await {
+                        Ok(()) => {}
+                        Err(e) => error!("Send keyboard report error: {}", e),
+                    };
+                }
+                KeyboardReportMessage::CompositeReport(report, report_type) => {
+                    write_other_report_to_host(report, report_type, other_hid_writer).await;
+                }
             }
         }
     }
diff --git a/rmk/src/lib.rs b/rmk/src/lib.rs
index 69218137..05ef8fec 100644
--- a/rmk/src/lib.rs
+++ b/rmk/src/lib.rs
@@ -85,6 +85,9 @@ pub(crate) static KEYBOARD_STATE: AtomicBool = AtomicBool::new(false);
 /// - 1: BLE
 /// - Other: reserved
 pub(crate) static CONNECTION_TYPE: AtomicU8 = AtomicU8::new(0);
+/// Whethe the connection is ready.
+/// After the connection is ready, the matrix starts scanning
+pub(crate) static CONNECTION_STATE: AtomicBool = AtomicBool::new(false);
 
 /// Run RMK keyboard service. This function should never return.
 ///
@@ -307,9 +310,10 @@ pub(crate) async fn run_usb_keyboard<
     >,
 ) -> ! {
     loop {
+        CONNECTION_STATE.store(false, core::sync::atomic::Ordering::Release);
         let usb_fut = usb_device.device.run();
         let keyboard_fut = keyboard.run();
-        let matrix_fut = matrix.scan();
+        let matrix_fut = matrix.run();
         let communication_fut = communication_task(
             keyboard_report_receiver,
             &mut usb_device.keyboard_hid_writer,
diff --git a/rmk/src/matrix.rs b/rmk/src/matrix.rs
index 90bb4856..614f0266 100644
--- a/rmk/src/matrix.rs
+++ b/rmk/src/matrix.rs
@@ -1,8 +1,9 @@
 use crate::{
     debounce::{DebounceState, DebouncerTrait},
     keyboard::{key_event_channel, KeyEvent},
+    CONNECTION_STATE,
 };
-use defmt::Format;
+use defmt::{error, info, Format};
 use embassy_time::{Instant, Timer};
 use embedded_hal::digital::{InputPin, OutputPin};
 #[cfg(feature = "async_matrix")]
@@ -17,20 +18,42 @@ pub(crate) trait MatrixTrait {
     const ROW: usize;
     const COL: usize;
 
+    // Wait for USB or BLE really connected
+    async fn wait_for_connected(&self) {
+        while !CONNECTION_STATE.load(core::sync::atomic::Ordering::Acquire) {
+            embassy_time::Timer::after_millis(100).await;
+        }
+        info!("Connected, start scanning matrix");
+    }
+
+    // Run the matrix
+    async fn run(&mut self) {
+        // We don't check disconnected state because disconnection means the task will be dropped 
+        loop {
+            self.wait_for_connected().await;
+            self.scan().await;
+        }
+    }
+
     // Do matrix scanning, save the result in matrix's key_state field.
     async fn scan(&mut self);
+
     // Read key state at position (row, col)
     fn get_key_state(&mut self, row: usize, col: usize) -> KeyState;
+
     // Update key state at position (row, col)
     fn update_key_state(&mut self, row: usize, col: usize, f: impl FnOnce(&mut KeyState));
+
     // Get matrix row num
     fn get_row_num(&self) -> usize {
         Self::ROW
     }
+
     // Get matrix col num
     fn get_col_num(&self) -> usize {
         Self::COL
     }
+
     #[cfg(feature = "async_matrix")]
     async fn wait_for_key(&mut self);
 }
diff --git a/rmk/src/split/driver.rs b/rmk/src/split/driver.rs
index 6216dad1..3d78944f 100644
--- a/rmk/src/split/driver.rs
+++ b/rmk/src/split/driver.rs
@@ -1,9 +1,9 @@
 ///! The abstracted driver layer of the split keyboard.
 ///!
-use crate::keyboard::{key_event_channel, KeyEvent};
-
 use super::SplitMessage;
-use defmt::{debug, error};
+use crate::keyboard::{key_event_channel, KeyEvent};
+use crate::CONNECTION_STATE;
+use defmt::{debug, error, warn};
 
 #[derive(Debug, Clone, Copy, defmt::Format)]
 pub(crate) enum SplitDriverError {
@@ -70,13 +70,19 @@ impl<
                             error!("Invalid peripheral row/col: {} {}", e.row, e.col);
                             continue;
                         }
-                        key_event_channel
-                            .send(KeyEvent {
-                                row: e.row + ROW_OFFSET as u8,
-                                col: e.col + COL_OFFSET as u8,
-                                pressed: e.pressed,
-                            })
-                            .await;
+
+                        if CONNECTION_STATE.load(core::sync::atomic::Ordering::Acquire) {
+                            // Only when the connection is established, send the key event.
+                            key_event_channel
+                                .send(KeyEvent {
+                                    row: e.row + ROW_OFFSET as u8,
+                                    col: e.col + COL_OFFSET as u8,
+                                    pressed: e.pressed,
+                                })
+                                .await;
+                        } else {
+                            warn!("Key event from peripheral is ignored because the connection is not established.");
+                        }
                     }
                 }
                 Err(e) => error!("Peripheral message read error: {:?}", e),
diff --git a/rmk/src/split/mod.rs b/rmk/src/split/mod.rs
index 43c5fe70..0396807a 100644
--- a/rmk/src/split/mod.rs
+++ b/rmk/src/split/mod.rs
@@ -21,22 +21,9 @@ pub const SPLIT_MESSAGE_MAX_SIZE: usize = SplitMessage::POSTCARD_MAX_SIZE + 4;
 pub(crate) enum SplitMessage {
     /// Key event from peripheral to central
     Key(KeyEvent),
-    /// Led state, on/off
+    /// Led state, on/off, from central to peripheral
     LedState(bool),
-}
-
-/// Message used for synchronization between central thread and peripheral receiver(both in central board)
-#[derive(Debug, Clone, Copy, defmt::Format)]
-pub(crate) enum KeySyncMessage {
-    /// Response of `SyncSignal`, sent key state matrix from peripheral monitor to main
-    /// u8 is the number of sent key states
-    StartSend(u16),
-    /// Key state: (row, col, key_pressing_state)
-    Key(u8, u8, bool),
-}
-
-/// Signal used for inform that the matrix starts receives key states from peripheral key receiver
-#[derive(Debug, Clone, Copy, defmt::Format)]
-pub(crate) enum KeySyncSignal {
-    Start,
+    /// The central connection state, true if central has been connected to host.
+    /// This message is sync from central to peripheral
+    ConnectionState(bool),
 }
diff --git a/rmk/src/split/nrf/central.rs b/rmk/src/split/nrf/central.rs
index 0dd8a3fb..39253138 100644
--- a/rmk/src/split/nrf/central.rs
+++ b/rmk/src/split/nrf/central.rs
@@ -1,16 +1,19 @@
 use core::sync::atomic::{AtomicBool, Ordering};
 
 use defmt::{error, info};
-use embassy_futures::join::join;
+use embassy_futures::{join::join, select::select};
 use embassy_sync::{
     blocking_mutex::raw::CriticalSectionRawMutex,
     channel::{Channel, Receiver, Sender},
 };
 use nrf_softdevice::ble::{central, gatt_client, Address, AddressType};
 
-use crate::split::{
-    driver::{PeripheralMatrixMonitor, SplitDriverError, SplitReader},
-    SplitMessage, SPLIT_MESSAGE_MAX_SIZE,
+use crate::{
+    split::{
+        driver::{PeripheralMatrixMonitor, SplitDriverError, SplitReader},
+        SplitMessage, SPLIT_MESSAGE_MAX_SIZE,
+    },
+    CONNECTION_STATE,
 };
 
 /// Gatt client used in split central to receive split message from peripherals
@@ -106,7 +109,7 @@ pub(crate) async fn run_ble_client(
         }
 
         // Receive peripheral's notifications
-        let disconnect_error = gatt_client::run(&conn, &ble_client, |event| match event {
+        let receive_peripheral = gatt_client::run(&conn, &ble_client, |event| match event {
             BleSplitCentralClientEvent::MessageToCentralNotification(message) => {
                 match postcard::from_bytes(&message) {
                     Ok(split_message) => {
@@ -119,10 +122,49 @@ pub(crate) async fn run_ble_client(
                     }
                 };
             }
-        })
-        .await;
+        });
+
+        // Notify messages to peripheral
+        let notify_peripheral = async {
+            let mut conn_state = CONNECTION_STATE.load(Ordering::Acquire);
+            // Send once on start
+            let mut buf = [0_u8; SPLIT_MESSAGE_MAX_SIZE];
+            match postcard::to_slice(&SplitMessage::ConnectionState(conn_state), &mut buf) {
+                Ok(_bytes) => {
+                    if let Err(e) = ble_client.message_to_peripheral_write(&buf).await {
+                        error!("BLE message_to_peripheral_write error: {}", e);
+                    }
+                }
+                Err(e) => error!("Postcard serialize split message error: {}", e),
+            };
+            loop {
+                // Check the central state every 200ms
+                embassy_time::Timer::after_millis(200).await;
+                // Current, only ConnectionState will be notified to peripheral
+                let current_conn_state = CONNECTION_STATE.load(Ordering::Acquire);
+                if conn_state != current_conn_state {
+                    // ConnectionState changed, notify peripheral
+                    conn_state = current_conn_state;
+                    let mut buf = [0_u8; SPLIT_MESSAGE_MAX_SIZE];
+                    match postcard::to_slice(&SplitMessage::ConnectionState(conn_state), &mut buf) {
+                        Ok(_bytes) => {
+                            if let Err(e) = ble_client.message_to_peripheral_write(&buf).await {
+                                error!("BLE message_to_peripheral_write error: {}", e);
+                            }
+                        }
+                        Err(e) => error!("Postcard serialize split message error: {}", e),
+                    };
+                }
+            }
+        };
+
+        match select(receive_peripheral, notify_peripheral).await {
+            embassy_futures::select::Either::First(e) => {
+                error!("BLE peripheral disconnect error: {:?}", e);
+            }
+            embassy_futures::select::Either::Second(_) => (),
+        }
 
-        error!("BLE peripheral disconnect error: {:?}", disconnect_error);
         // Wait for 1s before trying to connect (again)
         embassy_time::Timer::after_secs(1).await;
     }
diff --git a/rmk/src/split/nrf/peripheral.rs b/rmk/src/split/nrf/peripheral.rs
index 13e8a04c..69244192 100644
--- a/rmk/src/split/nrf/peripheral.rs
+++ b/rmk/src/split/nrf/peripheral.rs
@@ -1,10 +1,19 @@
+use crate::ble::nrf::softdevice_task;
+use crate::split::driver::{SplitDriverError, SplitReader, SplitWriter};
+use crate::split::peripheral::SplitPeripheral;
+use crate::split::{SplitMessage, SPLIT_MESSAGE_MAX_SIZE};
+use crate::MatrixTrait;
+use core::mem;
 use defmt::error;
+use embassy_executor::Spawner;
+use embassy_futures::block_on;
+use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
+use embassy_sync::channel::{Channel, Receiver};
+use nrf_softdevice::ble::gatt_server::set_sys_attrs;
+use nrf_softdevice::ble::peripheral::{advertise_connectable, ConnectableAdvertisement};
 use nrf_softdevice::ble::{gatt_server, Connection};
-
-use crate::split::{
-    driver::{SplitDriverError, SplitReader, SplitWriter},
-    SplitMessage, SPLIT_MESSAGE_MAX_SIZE,
-};
+use nrf_softdevice::ble::{set_address, Address, AddressType};
+use nrf_softdevice::{raw, Config, Softdevice};
 
 /// Gatt service used in split peripheral to send split message to central
 #[nrf_softdevice::gatt_service(uuid = "4dd5fbaa-18e5-4b07-bf0a-353698659946")]
@@ -26,29 +35,26 @@ pub(crate) struct BleSplitPeripheralServer {
 pub(crate) struct BleSplitPeripheralDriver<'a> {
     server: &'a BleSplitPeripheralServer,
     conn: &'a Connection,
+    receiver: Receiver<'a, CriticalSectionRawMutex, SplitMessage, 4>,
 }
 
 impl<'a> BleSplitPeripheralDriver<'a> {
-    pub(crate) fn new(server: &'a BleSplitPeripheralServer, conn: &'a Connection) -> Self {
-        Self { server, conn }
+    pub(crate) fn new(
+        server: &'a BleSplitPeripheralServer,
+        conn: &'a Connection,
+        receiver: Receiver<'a, CriticalSectionRawMutex, SplitMessage, 4>,
+    ) -> Self {
+        Self {
+            server,
+            conn,
+            receiver,
+        }
     }
 }
 
 impl<'a> SplitReader for BleSplitPeripheralDriver<'a> {
     async fn read(&mut self) -> Result<SplitMessage, SplitDriverError> {
-        let message = self
-            .server
-            .service
-            .message_to_peripheral_get()
-            .map_err(|e| {
-                error!("BLE read error: {:?}", e);
-                SplitDriverError::BleError(1)
-            })?;
-        let message: SplitMessage = postcard::from_bytes(&message).map_err(|e| {
-            error!("Postcard deserialize split message error: {}", e);
-            SplitDriverError::DeserializeError
-        })?;
-        Ok(message)
+        Ok(self.receiver.receive().await)
     }
 }
 
@@ -71,3 +77,146 @@ impl<'a> SplitWriter for BleSplitPeripheralDriver<'a> {
         Ok(bytes.len())
     }
 }
+
+/// Initialize and run the nRF peripheral keyboard service via BLE.
+///
+/// # Arguments
+///
+/// * `input_pins` - input gpio pins
+/// * `output_pins` - output gpio pins
+/// * `spawner` - embassy task spawner, used to spawn nrf_softdevice background task
+pub(crate) async fn initialize_nrf_ble_split_peripheral_and_run<
+    M: MatrixTrait,
+    const ROW: usize,
+    const COL: usize,
+>(
+    mut matrix: M,
+    central_addr: [u8; 6],
+    peripheral_addr: [u8; 6],
+    spawner: Spawner,
+) -> ! {
+    use defmt::info;
+    use embassy_futures::select::select3;
+    use nrf_softdevice::ble::gatt_server;
+
+    use crate::{
+        split::nrf::peripheral::{
+            BleSplitPeripheralDriver, BleSplitPeripheralServer, BleSplitPeripheralServerEvent,
+            SplitBleServiceEvent,
+        },
+        CONNECTION_STATE,
+    };
+
+    let ble_config = Config {
+        clock: Some(raw::nrf_clock_lf_cfg_t {
+            source: raw::NRF_CLOCK_LF_SRC_RC as u8,
+            rc_ctiv: 16,
+            rc_temp_ctiv: 2,
+            accuracy: raw::NRF_CLOCK_LF_ACCURACY_500_PPM as u8,
+            // External osc
+            // source: raw::NRF_CLOCK_LF_SRC_XTAL as u8,
+            // rc_ctiv: 0,
+            // rc_temp_ctiv: 0,
+            // accuracy: raw::NRF_CLOCK_LF_ACCURACY_20_PPM as u8,
+        }),
+        conn_gap: Some(raw::ble_gap_conn_cfg_t {
+            conn_count: 6,
+            event_length: 24,
+        }),
+        conn_gatt: Some(raw::ble_gatt_conn_cfg_t { att_mtu: 256 }),
+        gatts_attr_tab_size: Some(raw::ble_gatts_cfg_attr_tab_size_t {
+            attr_tab_size: raw::BLE_GATTS_ATTR_TAB_SIZE_DEFAULT,
+        }),
+        gap_role_count: Some(raw::ble_gap_cfg_role_count_t {
+            adv_set_count: 1,
+            periph_role_count: 4,
+            central_role_count: 4,
+            central_sec_count: 4,
+            _bitfield_1: raw::ble_gap_cfg_role_count_t::new_bitfield_1(0),
+        }),
+        gap_device_name: Some(raw::ble_gap_cfg_device_name_t {
+            p_value: "rmk_peripheral_board".as_ptr() as _,
+            current_len: "rmk_peripheral_board".len() as u16,
+            max_len: "rmk_peripheral_board".len() as u16,
+            write_perm: unsafe { mem::zeroed() },
+            _bitfield_1: raw::ble_gap_cfg_device_name_t::new_bitfield_1(
+                raw::BLE_GATTS_VLOC_STACK as u8,
+            ),
+        }),
+        ..Default::default()
+    };
+
+    let sd = Softdevice::enable(&ble_config);
+    set_address(
+        sd,
+        &Address::new(AddressType::RandomStatic, peripheral_addr),
+    );
+
+    {
+        // Use the immutable ref of `Softdevice` to run the softdevice_task
+        // The mumtable ref is used for configuring Flash and BleServer
+        let sdv = unsafe { nrf_softdevice::Softdevice::steal() };
+        defmt::unwrap!(spawner.spawn(softdevice_task(sdv)))
+    };
+
+    let server = defmt::unwrap!(BleSplitPeripheralServer::new(sd));
+
+    loop {
+        CONNECTION_STATE.store(false, core::sync::atomic::Ordering::Release);
+        let advertisement: ConnectableAdvertisement<'_> =
+            ConnectableAdvertisement::NonscannableDirected {
+                peer: Address::new(AddressType::RandomStatic, central_addr),
+            };
+        let conn = match advertise_connectable(sd, advertisement, &Default::default()).await {
+            Ok(conn) => conn,
+            Err(e) => {
+                defmt::error!("Split peripheral advertise error: {}", e);
+                continue;
+            }
+        };
+
+        // Channel used for receiving messages from central
+        let receive_channel: Channel<CriticalSectionRawMutex, SplitMessage, 4> = Channel::new();
+        let receiver = receive_channel.receiver();
+        let sender = receive_channel.sender();
+
+        // Set sys attr of peripheral
+        set_sys_attrs(&conn, None).ok();
+
+        let server_fut = gatt_server::run(&conn, &server, |event| match event {
+            BleSplitPeripheralServerEvent::Service(split_event) => match split_event {
+                SplitBleServiceEvent::MessageToCentralCccdWrite { notifications } => {
+                    info!("Split value CCCD updated: {}", notifications)
+                }
+                SplitBleServiceEvent::MessageToPeripheralWrite(message) => {
+                    match postcard::from_bytes::<SplitMessage>(&message) {
+                        Ok(message) => {
+                            info!("Message from central: {:?}", message);
+                            // Retry 3 times
+                            let mut success = false;
+                            for _i in 0..3 {
+                                if let Err(e) = sender.try_send(message) {
+                                    error!("Send split message to reader error: {}", e);
+                                    continue;
+                                }
+                                success = true;
+                                break;
+                            }
+                            // Should we block on it?
+                            if !success {
+                                block_on(sender.send(message));
+                            }
+                        }
+                        Err(e) => defmt::error!("Postcard deserialize split message error: {}", e),
+                    }
+                }
+            },
+        });
+
+        let mut peripheral =
+            SplitPeripheral::new(BleSplitPeripheralDriver::new(&server, &conn, receiver));
+        let peripheral_fut = peripheral.run();
+        let matrix_fut = matrix.run();
+        select3(matrix_fut, server_fut, peripheral_fut).await;
+    }
+}
diff --git a/rmk/src/split/peripheral.rs b/rmk/src/split/peripheral.rs
index b3dc268c..49161b81 100644
--- a/rmk/src/split/peripheral.rs
+++ b/rmk/src/split/peripheral.rs
@@ -1,36 +1,23 @@
+use super::driver::{SplitReader, SplitWriter};
+use super::SplitMessage;
 #[cfg(not(feature = "rapid_debouncer"))]
 use crate::debounce::default_bouncer::DefaultDebouncer;
 #[cfg(feature = "rapid_debouncer")]
 use crate::debounce::fast_debouncer::RapidDebouncer;
 use crate::debounce::DebouncerTrait;
+use crate::direct_pin::DirectPinMatrix;
 use crate::keyboard::key_event_channel;
 use crate::matrix::Matrix;
-use crate::matrix::MatrixTrait;
+use crate::CONNECTION_STATE;
+use defmt::error;
+#[cfg(feature = "_nrf_ble")]
+use embassy_executor::Spawner;
+use embassy_futures::select::select;
 use embedded_hal::digital::{InputPin, OutputPin};
 #[cfg(feature = "async_matrix")]
 use embedded_hal_async::digital::Wait;
-#[cfg(feature = "_nrf_ble")]
-use {
-    crate::ble::nrf::softdevice_task,
-    core::mem,
-    embassy_executor::Spawner,
-    nrf_softdevice::ble::gatt_server::set_sys_attrs,
-    nrf_softdevice::ble::peripheral::{advertise_connectable, ConnectableAdvertisement},
-    nrf_softdevice::ble::{set_address, Address, AddressType},
-    nrf_softdevice::{raw, Config, Softdevice},
-};
-
-use super::{
-    driver::{SplitReader, SplitWriter},
-    SplitMessage,
-};
-use crate::direct_pin::DirectPinMatrix;
-
 #[cfg(not(feature = "_nrf_ble"))]
-use {
-    super::serial::SerialSplitDriver,
-    embedded_io_async::{Read, Write},
-};
+use embedded_io_async::{Read, Write};
 
 /// Run the split peripheral service.
 ///
@@ -76,10 +63,13 @@ pub async fn run_rmk_split_peripheral<
     let matrix = Matrix::<_, _, _, COL, ROW>::new(input_pins, output_pins, debouncer);
 
     #[cfg(not(feature = "_nrf_ble"))]
-    initialize_serial_split_peripheral_and_run::<_, S, ROW, COL>(matrix, serial).await;
+    crate::split::serial::initialize_serial_split_peripheral_and_run::<_, S, ROW, COL>(
+        matrix, serial,
+    )
+    .await;
 
     #[cfg(feature = "_nrf_ble")]
-    initialize_nrf_ble_split_peripheral_and_run::<_, ROW, COL>(
+    crate::split::nrf::peripheral::initialize_nrf_ble_split_peripheral_and_run::<_, ROW, COL>(
         matrix,
         central_addr,
         peripheral_addr,
@@ -124,10 +114,13 @@ pub async fn run_rmk_split_peripheral_direct_pin<
     let matrix = DirectPinMatrix::<_, _, ROW, COL, SIZE>::new(direct_pins, debouncer, low_active);
 
     #[cfg(not(feature = "_nrf_ble"))]
-    initialize_serial_split_peripheral_and_run::<_, S, ROW, COL>(matrix, serial).await;
+    crate::split::serial::initialize_serial_split_peripheral_and_run::<_, S, ROW, COL>(
+        matrix, serial,
+    )
+    .await;
 
     #[cfg(feature = "_nrf_ble")]
-    initialize_nrf_ble_split_peripheral_and_run::<_, ROW, COL>(
+    crate::split::nrf::peripheral::initialize_nrf_ble_split_peripheral_and_run::<_, ROW, COL>(
         matrix,
         central_addr,
         peripheral_addr,
@@ -136,146 +129,6 @@ pub async fn run_rmk_split_peripheral_direct_pin<
     .await;
 }
 
-/// Initialize and run the nRF peripheral keyboard service via BLE.
-///
-/// # Arguments
-///
-/// * `input_pins` - input gpio pins
-/// * `output_pins` - output gpio pins
-/// * `spawner` - embassy task spawner, used to spawn nrf_softdevice background task
-#[cfg(feature = "_nrf_ble")]
-pub(crate) async fn initialize_nrf_ble_split_peripheral_and_run<
-    M: MatrixTrait,
-    const ROW: usize,
-    const COL: usize,
->(
-    mut matrix: M,
-    central_addr: [u8; 6],
-    peripheral_addr: [u8; 6],
-    spawner: Spawner,
-) -> ! {
-    use defmt::info;
-    use embassy_futures::select::select3;
-    use nrf_softdevice::ble::gatt_server;
-
-    use crate::split::nrf::peripheral::{
-        BleSplitPeripheralDriver, BleSplitPeripheralServer, BleSplitPeripheralServerEvent,
-        SplitBleServiceEvent,
-    };
-
-    let ble_config = Config {
-        clock: Some(raw::nrf_clock_lf_cfg_t {
-            source: raw::NRF_CLOCK_LF_SRC_RC as u8,
-            rc_ctiv: 16,
-            rc_temp_ctiv: 2,
-            accuracy: raw::NRF_CLOCK_LF_ACCURACY_500_PPM as u8,
-            // External osc
-            // source: raw::NRF_CLOCK_LF_SRC_XTAL as u8,
-            // rc_ctiv: 0,
-            // rc_temp_ctiv: 0,
-            // accuracy: raw::NRF_CLOCK_LF_ACCURACY_20_PPM as u8,
-        }),
-        conn_gap: Some(raw::ble_gap_conn_cfg_t {
-            conn_count: 6,
-            event_length: 24,
-        }),
-        conn_gatt: Some(raw::ble_gatt_conn_cfg_t { att_mtu: 256 }),
-        gatts_attr_tab_size: Some(raw::ble_gatts_cfg_attr_tab_size_t {
-            attr_tab_size: raw::BLE_GATTS_ATTR_TAB_SIZE_DEFAULT,
-        }),
-        gap_role_count: Some(raw::ble_gap_cfg_role_count_t {
-            adv_set_count: 1,
-            periph_role_count: 4,
-            central_role_count: 4,
-            central_sec_count: 4,
-            _bitfield_1: raw::ble_gap_cfg_role_count_t::new_bitfield_1(0),
-        }),
-        gap_device_name: Some(raw::ble_gap_cfg_device_name_t {
-            p_value: "rmk_peripheral_board".as_ptr() as _,
-            current_len: "rmk_peripheral_board".len() as u16,
-            max_len: "rmk_peripheral_board".len() as u16,
-            write_perm: unsafe { mem::zeroed() },
-            _bitfield_1: raw::ble_gap_cfg_device_name_t::new_bitfield_1(
-                raw::BLE_GATTS_VLOC_STACK as u8,
-            ),
-        }),
-        ..Default::default()
-    };
-
-    let sd = Softdevice::enable(&ble_config);
-    set_address(
-        sd,
-        &Address::new(AddressType::RandomStatic, peripheral_addr),
-    );
-
-    {
-        // Use the immutable ref of `Softdevice` to run the softdevice_task
-        // The mumtable ref is used for configuring Flash and BleServer
-        let sdv = unsafe { nrf_softdevice::Softdevice::steal() };
-        defmt::unwrap!(spawner.spawn(softdevice_task(sdv)))
-    };
-
-    let server = defmt::unwrap!(BleSplitPeripheralServer::new(sd));
-
-    loop {
-        let advertisement = ConnectableAdvertisement::NonscannableDirected {
-            peer: Address::new(AddressType::RandomStatic, central_addr),
-        };
-        let conn = match advertise_connectable(sd, advertisement, &Default::default()).await {
-            Ok(conn) => conn,
-            Err(e) => {
-                defmt::error!("Split peripheral advertise error: {}", e);
-                continue;
-            }
-        };
-
-        // Set sys attr of peripheral
-        set_sys_attrs(&conn, None).unwrap();
-
-        let server_fut = gatt_server::run(&conn, &server, |event| match event {
-            BleSplitPeripheralServerEvent::Service(split_event) => match split_event {
-                SplitBleServiceEvent::MessageToCentralCccdWrite { notifications } => {
-                    info!("Split value CCCD updated: {}", notifications)
-                }
-                SplitBleServiceEvent::MessageToPeripheralWrite(message) => {
-                    // TODO: Handle message from central to peripheral
-                    info!("Message from central: {:?}", message);
-                }
-            },
-        });
-
-        let mut peripheral = SplitPeripheral::new(BleSplitPeripheralDriver::new(&server, &conn));
-        let peripheral_fut = peripheral.run();
-        let matrix_fut = matrix.scan();
-        select3(matrix_fut, server_fut, peripheral_fut).await;
-    }
-}
-
-/// Initialize and run the peripheral keyboard service via serial.
-///
-/// # Arguments
-///
-/// * `input_pins` - input gpio pins
-/// * `output_pins` - output gpio pins
-/// * `serial` - serial port to send key events to central board
-#[cfg(not(feature = "_nrf_ble"))]
-pub(crate) async fn initialize_serial_split_peripheral_and_run<
-    M: MatrixTrait,
-    S: Write + Read,
-    const ROW: usize,
-    const COL: usize,
->(
-    mut matrix: M,
-    serial: S,
-) -> ! {
-    use embassy_futures::select::select;
-
-    let mut peripheral = SplitPeripheral::new(SerialSplitDriver::new(serial));
-    loop {
-        select(matrix.scan(), peripheral.run()).await;
-    }
-}
-
 /// The split peripheral instance.
 pub(crate) struct SplitPeripheral<S: SplitWriter + SplitReader> {
     split_driver: S,
@@ -292,12 +145,26 @@ impl<S: SplitWriter + SplitReader> SplitPeripheral<S> {
     /// If also receives split messages from the central through `SplitReader`.
     pub(crate) async fn run(&mut self) -> ! {
         loop {
-            let e = key_event_channel.receive().await;
-
-            self.split_driver.write(&SplitMessage::Key(e)).await.ok();
-
-            // 10KHZ scan rate
-            embassy_time::Timer::after_micros(10).await;
+            match select(self.split_driver.read(), key_event_channel.receive()).await {
+                embassy_futures::select::Either::First(m) => match m {
+                    // Currently only handle the central state message
+                    Ok(split_message) => match split_message {
+                        SplitMessage::ConnectionState(state) => {
+                            CONNECTION_STATE.store(state, core::sync::atomic::Ordering::Release);
+                        }
+                        _ => (),
+                    },
+                    Err(e) => {
+                        error!("Split message read error: {:?}", e);
+                    }
+                },
+                embassy_futures::select::Either::Second(e) => {
+                    // Only send the key event if the connection is established
+                    if CONNECTION_STATE.load(core::sync::atomic::Ordering::Acquire) {
+                        self.split_driver.write(&SplitMessage::Key(e)).await.ok();
+                    }
+                }
+            }
         }
     }
 }
diff --git a/rmk/src/split/serial/mod.rs b/rmk/src/split/serial/mod.rs
index 798c5284..e316728f 100644
--- a/rmk/src/split/serial/mod.rs
+++ b/rmk/src/split/serial/mod.rs
@@ -1,10 +1,9 @@
 use defmt::{error, info};
 use embedded_io_async::{Read, Write};
 
-use crate::split::{
-    driver::{PeripheralMatrixMonitor, SplitReader, SplitWriter},
-    SplitMessage, SPLIT_MESSAGE_MAX_SIZE,
-};
+use crate::{matrix::MatrixTrait, split::{
+    driver::{PeripheralMatrixMonitor, SplitReader, SplitWriter}, peripheral::SplitPeripheral, SplitMessage, SPLIT_MESSAGE_MAX_SIZE
+}};
 
 use super::driver::SplitDriverError;
 
@@ -78,3 +77,27 @@ impl<S: Read + Write> SplitWriter for SerialSplitDriver<S> {
             .map_err(|_e| SplitDriverError::SerialError)
     }
 }
+
+/// Initialize and run the peripheral keyboard service via serial.
+///
+/// # Arguments
+///
+/// * `input_pins` - input gpio pins
+/// * `output_pins` - output gpio pins
+/// * `serial` - serial port to send key events to central board
+pub(crate) async fn initialize_serial_split_peripheral_and_run<
+    M: MatrixTrait,
+    S: Write + Read,
+    const ROW: usize,
+    const COL: usize,
+>(
+    mut matrix: M,
+    serial: S,
+) -> ! {
+    use embassy_futures::select::select;
+
+    let mut peripheral = SplitPeripheral::new(SerialSplitDriver::new(serial));
+    loop {
+        select(matrix.run(), peripheral.run()).await;
+    }
+}
diff --git a/rmk/src/usb/mod.rs b/rmk/src/usb/mod.rs
index 342c73e5..d7b39357 100644
--- a/rmk/src/usb/mod.rs
+++ b/rmk/src/usb/mod.rs
@@ -15,7 +15,7 @@ use usbd_hid::descriptor::SerializedDescriptor;
 use crate::{
     config::KeyboardUsbConfig,
     hid::{UsbHidReader, UsbHidReaderWriter, UsbHidWriter},
-    usb::descriptor::{CompositeReport, ViaReport},
+    usb::descriptor::{CompositeReport, ViaReport}, CONNECTION_STATE,
 };
 
 pub(crate) static USB_STATE: AtomicU8 = AtomicU8::new(UsbState::Disabled as u8);
@@ -213,6 +213,7 @@ impl Handler for UsbDeviceHandler {
     fn configured(&mut self, configured: bool) {
         if configured {
             USB_STATE.store(UsbState::Configured as u8, Ordering::Relaxed);
+            CONNECTION_STATE.store(true, Ordering::Release);
             info!("Device configured, it may now draw up to the configured current from Vbus.")
         } else {
             USB_STATE.store(UsbState::Enabled as u8, Ordering::Relaxed);


Test patch: