fix: notification issues (#309)

* fix: notification issues
* refactor: add gatt connection api to simplify gatt processing

Closes #308

Author: chrissuozzo

examples/apps/src/ble_bas_peripheral.rs

@@ -52,7 +52,7 @@ where
 
     let _ = join(ble_task(runner), async {
         loop {
-            match advertise("Trouble Example", &mut peripheral).await {
+            match advertise("Trouble Example", &mut peripheral, &server).await {
                 Ok(conn) => {
                     // set up tasks when the connection is established to a central, so they don't run when no one is connected.
                     let a = gatt_events_task(&server, &conn);
@@ -101,67 +101,53 @@ async fn ble_task<C: Controller>(mut runner: Runner<'_, C>) {
 ///
 /// This function will handle the GATT events and process them.
 /// This is how we interact with read and write requests.
-async fn gatt_events_task(server: &Server<'_>, conn: &Connection<'_>) -> Result<(), Error> {
+async fn gatt_events_task(server: &Server<'_>, conn: &GattConnection<'_, '_>) -> Result<(), Error> {
     let level = server.battery_service.level;
     loop {
         match conn.next().await {
-            ConnectionEvent::Disconnected { reason } => {
+            GattConnectionEvent::Disconnected { reason } => {
                 info!("[gatt] disconnected: {:?}", reason);
                 break;
             }
-            ConnectionEvent::Gatt { data } => {
-                // We can choose to handle event directly without an attribute table
-                // let req = data.request();
-                // ..
-                // data.reply(conn, Ok(AttRsp::Error { .. }))
-
-                // But to simplify things, process it in the GATT server that handles
-                // the protocol details
-                match data.process(server).await {
-                    // Server processing emits
-                    Ok(Some(event)) => {
-                        match &event {
-                            GattEvent::Read(event) => {
-                                if event.handle() == level.handle {
-                                    let value = server.get(&level);
-                                    info!("[gatt] Read Event to Level Characteristic: {:?}", value);
-                                }
-                            }
-                            GattEvent::Write(event) => {
-                                if event.handle() == level.handle {
-                                    info!("[gatt] Write Event to Level Characteristic: {:?}", event.data());
-                                }
+            GattConnectionEvent::Gatt { event } => match event {
+                Ok(event) => {
+                    match &event {
+                        GattEvent::Read(event) => {
+                            if event.handle() == level.handle {
+                                let value = server.get(&level);
+                                info!("[gatt] Read Event to Level Characteristic: {:?}", value);
                             }
                         }
-
-                        // This step is also performed at drop(), but writing it explicitly is necessary
-                        // in order to ensure reply is sent.
-                        match event.accept() {
-                            Ok(reply) => {
-                                reply.send().await;
-                            }
-                            Err(e) => {
-                                warn!("[gatt] error sending response: {:?}", e);
+                        GattEvent::Write(event) => {
+                            if event.handle() == level.handle {
+                                info!("[gatt] Write Event to Level Characteristic: {:?}", event.data());
                             }
                         }
                     }
-                    Ok(_) => {}
-                    Err(e) => {
-                        warn!("[gatt] error processing event: {:?}", e);
+
+                    // This step is also performed at drop(), but writing it explicitly is necessary
+                    // in order to ensure reply is sent.
+                    match event.accept() {
+                        Ok(reply) => {
+                            reply.send().await;
+                        }
+                        Err(e) => warn!("[gatt] error sending response: {:?}", e),
                     }
                 }
-            }
+                Err(e) => warn!("[gatt] error processing event: {:?}", e),
+            },
         }
     }
     info!("[gatt] task finished");
     Ok(())
 }
 
 /// Create an advertiser to use to connect to a BLE Central, and wait for it to connect.
-async fn advertise<'a, C: Controller>(
+async fn advertise<'a, 'b, C: Controller>(
     name: &'a str,
     peripheral: &mut Peripheral<'a, C>,
-) -> Result<Connection<'a>, BleHostError<C::Error>> {
+    server: &'b Server<'_>,
+) -> Result<GattConnection<'a, 'b>, BleHostError<C::Error>> {
     let mut advertiser_data = [0; 31];
     AdStructure::encode_slice(
         &[
@@ -181,7 +167,7 @@ async fn advertise<'a, C: Controller>(
         )
         .await?;
     info!("[adv] advertising");
-    let conn = advertiser.accept().await?;
+    let conn = advertiser.accept().await?.with_attribute_server(server)?;
     info!("[adv] connection established");
     Ok(conn)
 }
@@ -190,18 +176,18 @@ async fn advertise<'a, C: Controller>(
 /// This task will notify the connected central of a counter value every 2 seconds.
 /// It will also read the RSSI value every 2 seconds.
 /// and will stop when the connection is closed by the central or an error occurs.
-async fn custom_task<C: Controller>(server: &Server<'_>, conn: &Connection<'_>, stack: &Stack<'_, C>) {
+async fn custom_task<C: Controller>(server: &Server<'_>, conn: &GattConnection<'_, '_>, stack: &Stack<'_, C>) {
     let mut tick: u8 = 0;
     let level = server.battery_service.level;
     loop {
         tick = tick.wrapping_add(1);
         info!("[custom_task] notifying connection of tick {}", tick);
-        if level.notify(server, conn, &tick).await.is_err() {
+        if level.notify(conn, &tick).await.is_err() {
             info!("[custom_task] error notifying connection");
             break;
         };
         // read RSSI (Received Signal Strength Indicator) of the connection.
-        if let Ok(rssi) = conn.rssi(stack).await {
+        if let Ok(rssi) = conn.raw().rssi(stack).await {
             info!("[custom_task] RSSI: {:?}", rssi);
         } else {
             info!("[custom_task] error getting RSSI");

host-macros/src/server.rs

@@ -8,12 +8,14 @@ use proc_macro2::TokenStream as TokenStream2;
 use quote::{quote, quote_spanned};
 use syn::meta::ParseNestedMeta;
 use syn::spanned::Spanned;
-use syn::{parse_quote, Expr, Result};
+use syn::{Expr, Result, parse_quote};
 
 #[derive(Default)]
 pub(crate) struct ServerArgs {
     mutex_type: Option<syn::Type>,
     attribute_table_size: Option<Expr>,
+    cccd_table_size: Option<Expr>,
+    connections_max: Option<Expr>,
 }
 
 impl ServerArgs {
@@ -42,11 +44,21 @@ impl ServerArgs {
                 })?;
                 self.attribute_table_size = Some(buffer.parse()?);
             }
-            other => {
-                return Err(meta.error(format!(
-                "Unsupported server property: '{other}'.\nSupported properties are: mutex_type, attribute_table_size"
-            )))
+            "cccd_table_size" => {
+                let buffer = meta.value().map_err(|_| {
+                    Error::custom("cccd_table_size must be followed by `= [size]`. e.g. cccd_table_size = 4".to_string())
+                })?;
+                self.cccd_table_size = Some(buffer.parse()?);
+            }
+            "connections_max" => {
+                let buffer = meta.value().map_err(|_| {
+                    Error::custom("connections_max must be followed by `= [size]`. e.g. connections_max = 1".to_string())
+                })?;
+                self.connections_max = Some(buffer.parse()?);
             }
+            other => return Err(meta.error(format!(
+                "Unsupported server property: '{other}'.\nSupported properties are: mutex_type, attribute_table_size, cccd_table_size, connections_max"
+            ))),
         }
         Ok(())
     }
@@ -75,6 +87,7 @@ impl ServerBuilder {
         let mut code_service_init = TokenStream2::new();
         let mut code_server_populate = TokenStream2::new();
         let mut code_attribute_summation = TokenStream2::new();
+        let mut code_cccd_summation = TokenStream2::new();
         for service in &self.properties.fields {
             let vis = &service.vis;
             let service_span = service.span();
@@ -95,6 +108,10 @@ impl ServerBuilder {
 
             code_attribute_summation.extend(quote_spanned! {service_span=>
                + #service_type::ATTRIBUTE_COUNT
+            });
+
+            code_cccd_summation.extend(quote_spanned! {service_span=>
+               + #service_type::CCCD_COUNT
             })
         }
 
@@ -104,16 +121,30 @@ impl ServerBuilder {
             parse_quote!(trouble_host::gap::GAP_SERVICE_ATTRIBUTE_COUNT #code_attribute_summation)
         };
 
+        let cccd_table_size = if let Some(value) = self.arguments.cccd_table_size {
+            value
+        } else {
+            parse_quote!(0 #code_cccd_summation)
+        };
+
+        let connections_max = if let Some(value) = self.arguments.connections_max {
+            value
+        } else {
+            parse_quote!(1)
+        };
+
         quote! {
             const _ATTRIBUTE_TABLE_SIZE: usize = #attribute_table_size;
             // This pattern causes the assertion to happen at compile time
             const _: () = {
                 core::assert!(_ATTRIBUTE_TABLE_SIZE >= trouble_host::gap::GAP_SERVICE_ATTRIBUTE_COUNT #code_attribute_summation, "Specified attribute table size is insufficient. Please increase attribute_table_size or remove the argument entirely to allow automatic sizing of the attribute table.");
             };
+            const _CCCD_TABLE_SIZE: usize = #cccd_table_size;
+            const _CONNECTIONS_MAX: usize = #connections_max;
 
             #visibility struct #name<'values>
             {
-                server: trouble_host::prelude::AttributeServer<'values, #mutex_type, _ATTRIBUTE_TABLE_SIZE>,
+                server: trouble_host::prelude::AttributeServer<'values, #mutex_type, _ATTRIBUTE_TABLE_SIZE, _CCCD_TABLE_SIZE, _CONNECTIONS_MAX>,
                 #code_service_definition
             }
 
@@ -179,7 +210,7 @@ impl ServerBuilder {
 
             impl<'values> core::ops::Deref for #name<'values>
             {
-                type Target = trouble_host::prelude::AttributeServer<'values, #mutex_type, _ATTRIBUTE_TABLE_SIZE>;
+                type Target = trouble_host::prelude::AttributeServer<'values, #mutex_type, _ATTRIBUTE_TABLE_SIZE, _CCCD_TABLE_SIZE, _CONNECTIONS_MAX>;
 
                 fn deref(&self) -> &Self::Target {
                     &self.server

host-macros/src/service.rs

@@ -109,6 +109,7 @@ pub(crate) struct ServiceBuilder {
     properties: syn::ItemStruct,
     args: ServiceArgs,
     attribute_count: usize,
+    cccd_count: usize,
     code_impl: TokenStream2,
     code_build_chars: TokenStream2,
     code_struct_init: TokenStream2,
@@ -121,6 +122,7 @@ impl ServiceBuilder {
             properties,
             args,
             attribute_count: 1, // Service counts as an attribute
+            cccd_count: 0,
             code_struct_init: TokenStream2::new(),
             code_impl: TokenStream2::new(),
             code_fields: TokenStream2::new(),
@@ -136,7 +138,12 @@ impl ServiceBuilder {
     /// If the characteristic has either the notify or indicate property,
     /// a Client Characteristic Configuration Descriptor (CCCD) declaration will also be added.
     fn increment_attributes(&mut self, access: &AccessArgs) -> usize {
-        self.attribute_count += if access.notify || access.indicate { 3 } else { 2 };
+        if access.notify || access.indicate {
+            self.cccd_count += 1;
+            self.attribute_count += 3;
+        } else {
+            self.attribute_count += 2;
+        }
         self.attribute_count
     }
     /// Construct the macro blueprint for the service struct.
@@ -150,6 +157,7 @@ impl ServiceBuilder {
         let code_build_chars = self.code_build_chars;
         let uuid = self.args.uuid;
         let attribute_count = self.attribute_count;
+        let cccd_count = self.cccd_count;
         quote! {
             #visibility struct #struct_name {
                 #fields
@@ -159,6 +167,7 @@ impl ServiceBuilder {
             #[allow(unused)]
             impl #struct_name {
                 #visibility const ATTRIBUTE_COUNT: usize = #attribute_count;
+                #visibility const CCCD_COUNT: usize = #cccd_count;
 
                 #visibility fn new<M, const MAX_ATTRIBUTES: usize>(table: &mut trouble_host::attribute::AttributeTable<'_, M, MAX_ATTRIBUTES>) -> Self
                 where

host/src/attribute.rs

@@ -13,7 +13,7 @@ use crate::Error;
 use crate::att::AttErrorCode;
 use crate::attribute_server::AttributeServer;
 use crate::cursor::{ReadCursor, WriteCursor};
-use crate::prelude::{AsGatt, Connection, FixedGattValue, FromGatt};
+use crate::prelude::{AsGatt, Connection, FixedGattValue, FromGatt, GattConnection};
 use crate::types::gatt_traits::FromGattError;
 pub use crate::types::uuid::Uuid;
 
@@ -629,17 +629,13 @@ impl<T: FromGatt> Characteristic<T> {
     /// If the provided connection has not subscribed for this characteristic, it will not be notified.
     ///
     /// If the characteristic does not support notifications, an error is returned.
-    pub async fn notify<M: RawMutex, const MAX: usize>(
-        &self,
-        server: &AttributeServer<'_, M, MAX>,
-        connection: &Connection<'_>,
-        value: &T,
-    ) -> Result<(), Error> {
+    pub async fn notify(&self, connection: &GattConnection<'_, '_>, value: &T) -> Result<(), Error> {
         let value = value.as_gatt();
-        server.table().set_raw(self.handle, value)?;
+        let server = connection.server;
+        server.set(self.handle, value)?;
 
         let cccd_handle = self.cccd_handle.ok_or(Error::NotFound)?;
-
+        let connection = connection.raw();
         if !server.should_notify(connection, cccd_handle) {
             // No reason to fail?
             return Ok(());
@@ -662,9 +658,9 @@ impl<T: FromGatt> Characteristic<T> {
     }
 
     /// Set the value of the characteristic in the provided attribute server.
-    pub fn set<M: RawMutex, const MAX: usize>(
+    pub fn set<M: RawMutex, const AT: usize, const CT: usize, const CN: usize>(
         &self,
-        server: &AttributeServer<'_, M, MAX>,
+        server: &AttributeServer<'_, M, AT, CT, CN>,
         value: &T,
     ) -> Result<(), Error> {
         let value = value.as_gatt();
@@ -676,7 +672,10 @@ impl<T: FromGatt> Characteristic<T> {
     ///
     /// If the characteristic for the handle cannot be found, an error is returned.
     ///
-    pub fn get<M: RawMutex, const MAX: usize>(&self, server: &AttributeServer<'_, M, MAX>) -> Result<T, Error> {
+    pub fn get<M: RawMutex, const AT: usize, const CT: usize, const CN: usize>(
+        &self,
+        server: &AttributeServer<'_, M, AT, CT, CN>,
+    ) -> Result<T, Error> {
         server.table().get(self)
     }
 
@@ -874,6 +873,8 @@ pub enum CCCDFlag {
 }
 
 /// CCCD flag.
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+#[derive(Clone, Copy, Default)]
 pub struct CCCD(pub(crate) u16);
 
 impl<const T: usize> From<[CCCDFlag; T]> for CCCD {
@@ -887,6 +888,16 @@ impl<const T: usize> From<[CCCDFlag; T]> for CCCD {
 }
 
 impl CCCD {
+    /// Get raw value
+    pub fn raw(&self) -> u16 {
+        self.0
+    }
+
+    /// Clear all properties
+    pub fn disable(&mut self) {
+        self.0 = 0;
+    }
+
     /// Check if any of the properties are set.
     pub fn any(&self, props: &[CCCDFlag]) -> bool {
         for p in props {
@@ -896,4 +907,15 @@ impl CCCD {
         }
         false
     }
+
+    /// Enable or disable notifications
+    pub fn set_notify(&mut self, is_enabled: bool) {
+        let mask: u16 = CCCDFlag::Notify as u16;
+        self.0 = if is_enabled { self.0 | mask } else { self.0 & !mask };
+    }
+
+    /// Check if notifications are enabled
+    pub fn should_notify(&self) -> bool {
+        (self.0 & (CCCDFlag::Notify as u16)) != 0
+    }
 }