feat: Add database tracking and report for Push Reliability

autoconnect/autoconnect-common/src/protocol.rs

@@ -80,11 +80,19 @@ impl FromStr for ClientMessage {
     }
 }
 
+/// Returned ACKnowledgement of the received message by the User Agent.
+/// This is the payload for the `messageType:ack` packet.
+///
 #[derive(Debug, Deserialize)]
 pub struct ClientAck {
+    // The channel_id which received messages
     #[serde(rename = "channelID")]
     pub channel_id: Uuid,
+    // The corresponding version number for the message.
     pub version: String,
+    // The Reliability ID for the messages (if present)
+    #[serde(default)]
+    pub reliability_id: Option<String>,
 }
 
 #[derive(Debug, Serialize)]

autoconnect/autoconnect-settings/src/app_state.rs

@@ -68,6 +68,7 @@ impl AppState {
             db_settings: settings.db_settings.clone(),
         };
         let storage_type = StorageType::from_dsn(&db_settings.dsn);
+
         #[allow(unused)]
         let db: Box<dyn DbClient> = match storage_type {
             #[cfg(feature = "bigtable")]

autoconnect/autoconnect-web/Cargo.toml

@@ -34,3 +34,5 @@ ctor.workspace = true
 tokio.workspace = true
 
 autoconnect_common = { workspace = true, features = ["test-support"] }
+
+[features]

autoconnect/autoconnect-ws/Cargo.toml

@@ -33,3 +33,5 @@ async-stream = "0.3"
 ctor.workspace = true
 
 autoconnect_common = { workspace = true, features = ["test-support"] }
+
+[features]

autoconnect/autoconnect-ws/autoconnect-ws-sm/Cargo.toml

@@ -30,3 +30,5 @@ tokio.workspace = true
 serde_json.workspace = true
 
 autoconnect_common = { workspace = true, features = ["test-support"] }
+
+[features]

autoconnect/autoconnect-ws/autoconnect-ws-sm/src/identified/on_client_msg.rs

@@ -216,6 +216,7 @@ impl WebPushClient {
                 // Get the stored notification record.
                 let n = &self.ack_state.unacked_stored_notifs[pos];
                 debug!("✅ Ack notif: {:?}", &n);
+                // TODO: Record "ack'd" reliability_id, if present.
                 // Only force delete Topic messages, since they don't have a timestamp.
                 // Other messages persist in the database, to be, eventually, cleaned up by their
                 // TTL. We will need to update the `CurrentTimestamp` field for the channel

autoendpoint/src/error.rs

@@ -70,6 +70,9 @@ pub enum ApiErrorKind {
     #[error(transparent)]
     Jwt(#[from] jsonwebtoken::errors::Error),
 
+    #[error(transparent)]
+    Serde(#[from] serde_json::Error),
+
     #[error("Error while validating token")]
     TokenHashValidation(#[source] openssl::error::ErrorStack),
 
@@ -143,6 +146,7 @@ impl ApiErrorKind {
 
             ApiErrorKind::VapidError(_)
             | ApiErrorKind::Jwt(_)
+            | ApiErrorKind::Serde(_)
             | ApiErrorKind::TokenHashValidation(_)
             | ApiErrorKind::InvalidAuthentication
             | ApiErrorKind::InvalidLocalAuth(_) => StatusCode::UNAUTHORIZED,
@@ -179,7 +183,7 @@ impl ApiErrorKind {
             ApiErrorKind::InvalidMessageId => "invalid_message_id",
 
             ApiErrorKind::VapidError(_) => "vapid_error",
-            ApiErrorKind::Jwt(_) => "jwt",
+            ApiErrorKind::Jwt(_) | ApiErrorKind::Serde(_) => "jwt",
             ApiErrorKind::TokenHashValidation(_) => "token_hash_validation",
             ApiErrorKind::InvalidAuthentication => "invalid_authentication",
             ApiErrorKind::InvalidLocalAuth(_) => "invalid_local_auth",
@@ -251,6 +255,7 @@ impl ApiErrorKind {
             ApiErrorKind::VapidError(_)
             | ApiErrorKind::TokenHashValidation(_)
             | ApiErrorKind::Jwt(_)
+            | ApiErrorKind::Serde(_)
             | ApiErrorKind::InvalidAuthentication
             | ApiErrorKind::InvalidLocalAuth(_) => Some(109),
 

autoendpoint/src/extractors/notification.rs

@@ -1,4 +1,4 @@
-use crate::error::{ApiError, ApiErrorKind};
+use crate::error::{ApiError, ApiErrorKind, ApiResult};
 use crate::extractors::{
     message_id::MessageId, notification_headers::NotificationHeaders, subscription::Subscription,
 };
@@ -103,6 +103,7 @@ impl From<Notification> for autopush_common::notification::Notification {
             timestamp: notification.timestamp,
             data: notification.data,
             sortkey_timestamp,
+            reliability_id: notification.subscription.reliability_id,
             headers: {
                 let headers: HashMap<String, String> = notification.headers.into();
                 if headers.is_empty() {
@@ -160,25 +161,28 @@ impl Notification {
     /// fields are still required when delivering to the connection server, so
     /// we can't simply convert this notification type to that one and serialize
     /// via serde.
-    pub fn serialize_for_delivery(&self) -> HashMap<&'static str, serde_json::Value> {
+    pub fn serialize_for_delivery(&self) -> ApiResult<HashMap<&'static str, serde_json::Value>> {
         let mut map = HashMap::new();
 
         map.insert(
             "channelID",
-            serde_json::to_value(self.subscription.channel_id).unwrap(),
+            serde_json::to_value(self.subscription.channel_id)?,
         );
-        map.insert("version", serde_json::to_value(&self.message_id).unwrap());
-        map.insert("ttl", serde_json::to_value(self.headers.ttl).unwrap());
-        map.insert("topic", serde_json::to_value(&self.headers.topic).unwrap());
-        map.insert("timestamp", serde_json::to_value(self.timestamp).unwrap());
+        map.insert("version", serde_json::to_value(&self.message_id)?);
+        map.insert("ttl", serde_json::to_value(self.headers.ttl)?);
+        map.insert("topic", serde_json::to_value(&self.headers.topic)?);
+        map.insert("timestamp", serde_json::to_value(self.timestamp)?);
+        if let Some(reliability_id) = &self.subscription.reliability_id {
+            map.insert("reliability_id", serde_json::to_value(reliability_id)?);
+        }
 
         if let Some(data) = &self.data {
-            map.insert("data", serde_json::to_value(data).unwrap());
+            map.insert("data", serde_json::to_value(data)?);
 
             let headers: HashMap<_, _> = self.headers.clone().into();
-            map.insert("headers", serde_json::to_value(headers).unwrap());
+            map.insert("headers", serde_json::to_value(headers)?);
         }
 
-        map
+        Ok(map)
     }
 }

autoendpoint/src/extractors/subscription.rs

@@ -38,7 +38,7 @@ pub struct Subscription {
     /// (This should ONLY be applied for messages that match known
     /// Mozilla provided VAPID public keys.)
     ///
-    pub tracking_id: Option<String>,
+    pub reliability_id: Option<String>,
 }
 
 impl FromRequest for Subscription {
@@ -73,11 +73,13 @@ impl FromRequest for Subscription {
                 .transpose()?;
 
             trace!("raw vapid: {:?}", &vapid);
-            let trackable = if let Some(vapid) = &vapid {
-                app_state.reliability.is_trackable(vapid)
-            } else {
-                false
-            };
+            let reliability_id: Option<String> = vapid.clone().and_then(|v| {
+                app_state
+                    .vapid_tracker
+                    .is_trackable(&v)
+                    .then(|| app_state.vapid_tracker.get_id(req.headers()))
+            });
+            debug!("🔍 Assigning Reliability ID: {:?}", reliability_id);
 
             // Capturing the vapid sub right now will cause too much cardinality. Instead,
             // let's just capture if we have a valid VAPID, as well as what sort of bad sub
@@ -132,14 +134,11 @@ impl FromRequest for Subscription {
                     .incr(&format!("updates.vapid.draft{:02}", vapid.vapid.version()))?;
             }
 
-            let tracking_id =
-                trackable.then(|| app_state.reliability.get_tracking_id(req.headers()));
-
             Ok(Subscription {
                 user,
                 channel_id,
                 vapid,
-                tracking_id,
+                reliability_id,
             })
         }
         .boxed_local()

autoendpoint/src/routers/common.rs

@@ -21,6 +21,12 @@ pub fn build_message_data(notification: &Notification) -> ApiResult<HashMap<&'st
         message_data.insert_opt("enc", notification.headers.encryption.as_ref());
         message_data.insert_opt("cryptokey", notification.headers.crypto_key.as_ref());
         message_data.insert_opt("enckey", notification.headers.encryption_key.as_ref());
+        // Report the data to the UA. How this value is reported back is still a work in progress.
+        trace!(
+            "🔍 Sending Reliability ID: {:?}",
+            notification.subscription.reliability_id
+        );
+        message_data.insert_opt("rid", notification.subscription.reliability_id.as_ref());
     }
 
     Ok(message_data)
@@ -239,7 +245,7 @@ pub mod tests {
                 user,
                 channel_id: channel_id(),
                 vapid: None,
-                tracking_id: None,
+                reliability_id: None,
             },
             headers: NotificationHeaders {
                 ttl: 0,

autoendpoint/src/routers/webpush.rs

@@ -179,7 +179,7 @@ impl WebPushRouter {
         node_id: &str,
     ) -> Result<Response, reqwest::Error> {
         let url = format!("{}/push/{}", node_id, notification.subscription.user.uaid);
-        let notification = notification.serialize_for_delivery();
+        let notification = notification.serialize_for_delivery().unwrap();
 
         self.http.put(&url).json(&notification).send().await
     }

autoendpoint/src/server.rs

@@ -48,7 +48,7 @@ pub struct AppState {
     pub apns_router: Arc<ApnsRouter>,
     #[cfg(feature = "stub")]
     pub stub_router: Arc<StubRouter>,
-    pub reliability: Arc<VapidTracker>,
+    pub vapid_tracker: Arc<VapidTracker>,
 }
 
 pub struct Server;
@@ -109,7 +109,7 @@ impl Server {
             )
             .await?,
         );
-        let reliability = Arc::new(VapidTracker(settings.tracking_keys()));
+        let vapid_tracker = Arc::new(VapidTracker(settings.tracking_keys()));
         #[cfg(feature = "stub")]
         let stub_router = Arc::new(StubRouter::new(settings.stub.clone())?);
         let app_state = AppState {
@@ -122,7 +122,7 @@ impl Server {
             apns_router,
             #[cfg(feature = "stub")]
             stub_router,
-            reliability,
+            vapid_tracker,
         };
 
         spawn_pool_periodic_reporter(

autoendpoint/src/settings.rs

@@ -169,9 +169,11 @@ impl Settings {
     // public key, but that may not always be true.
     pub fn tracking_keys(&self) -> Vec<String> {
         let keys = &self.tracking_keys.replace(['"', ' '], "");
-        Self::read_list_from_str(keys, "Invalid AUTOEND_TRACKING_KEYS")
-            .map(|v| v.to_owned())
-            .collect()
+        let result = Self::read_list_from_str(keys, "Invalid AUTOEND_TRACKING_KEYS")
+            .map(|v| v.to_owned().replace("=", ""))
+            .collect();
+        trace!("🔍 tracking_keys: {:?}", result);
+        result
     }
 
     /// Get the URL for this endpoint server
@@ -193,11 +195,20 @@ impl VapidTracker {
     pub fn is_trackable(&self, vapid: &VapidHeaderWithKey) -> bool {
         // ideally, [Settings.with_env_and_config_file()] does the work of pre-populating
         // the Settings.tracking_vapid_pubs cache, but we can't rely on that.
-        self.0.contains(&vapid.public_key)
+        let key = vapid.public_key.replace('=', "");
+        let result = self.0.contains(&key);
+        debug!("🔍 Checking {key} {}", {
+            if result {
+                "Match!"
+            } else {
+                "no match"
+            }
+        });
+        result
     }
 
     /// Extract the message Id from the headers (if present), otherwise just make one up.
-    pub fn get_tracking_id(&self, headers: &HeaderMap) -> String {
+    pub fn get_id(&self, headers: &HeaderMap) -> String {
         headers
             .get("X-MessageId")
             .and_then(|v|
@@ -304,7 +315,7 @@ mod tests {
     #[test]
     fn test_tracking_keys() -> ApiResult<()> {
         let settings = Settings{
-            tracking_keys: r#"["BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7vI"]"#.to_owned(),
+            tracking_keys: r#"["BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7"]"#.to_owned(),
             ..Default::default()
         };
 
@@ -314,7 +325,7 @@ mod tests {
                 token: "".to_owned(),
                 version_data: crate::headers::vapid::VapidVersionData::Version1,
             },
-            public_key: "BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7vI".to_owned()
+            public_key: "BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7==".to_owned()
         };
 
         let key_set = settings.tracking_keys();
@@ -327,20 +338,20 @@ mod tests {
     }
 
     #[test]
-    fn test_tracking_id() -> ApiResult<()> {
+    fn test_reliability_id() -> ApiResult<()> {
         let mut headers = HeaderMap::new();
         let keys = Vec::new();
         let reliability = VapidTracker(keys);
 
-        let key = reliability.get_tracking_id(&headers);
+        let key = reliability.get_id(&headers);
         assert!(!key.is_empty());
 
         headers.insert(
             HeaderName::from_lowercase(b"x-messageid").unwrap(),
             HeaderValue::from_static("123foobar456"),
         );
 
-        let key = reliability.get_tracking_id(&headers);
+        let key = reliability.get_id(&headers);
         assert_eq!(key, "123foobar456".to_owned());
 
         Ok(())

autopush-common/src/db/bigtable/bigtable_client/mod.rs

@@ -720,6 +720,7 @@ impl BigTableClientImpl {
             )
         })?;
 
+        // Create from the known, required fields.
         let mut notif = Notification {
             channel_id: range_key.channel_id,
             topic: range_key.topic,
@@ -730,6 +731,7 @@ impl BigTableClientImpl {
             ..Default::default()
         };
 
+        // Backfill the Optional fields
         if let Some(cell) = row.take_cell("data") {
             notif.data = Some(to_string(cell.value, "data")?);
         }
@@ -739,6 +741,10 @@ impl BigTableClientImpl {
                     .map_err(|e| DbError::Serialization(e.to_string()))?,
             );
         }
+        if let Some(cell) = row.take_cell("reliability_id") {
+            trace!("🚣  Is reliable");
+            notif.reliability_id = Some(to_string(cell.value, "reliability_id")?);
+        }
 
         trace!("🚣  Deserialized message row: {:?}", &notif);
         Ok(notif)
@@ -1182,6 +1188,15 @@ impl DbClient for BigTableClientImpl {
                 ..Default::default()
             });
         }
+
+        if let Some(reliability_id) = message.reliability_id {
+            cells.push(cell::Cell {
+                qualifier: "reliability_id".to_owned(),
+                value: reliability_id.into_bytes(),
+                timestamp: expiry,
+                ..Default::default()
+            });
+        }
         row.add_cells(family, cells);
         trace!("🉑 Adding row");
         self.write_row(row).await?;
@@ -1301,6 +1316,7 @@ impl DbClient for BigTableClientImpl {
         );
 
         let messages = self.rows_to_notifications(rows)?;
+
         // Note: Bigtable always returns a timestamp of None.
         // Under Bigtable `current_timestamp` is instead initially read
         // from [get_user].

autopush-common/src/db/mod.rs

@@ -245,6 +245,10 @@ pub struct NotificationRecord {
     /// value before sending it to storage or a connection node.
     #[serde(skip_serializing_if = "Option::is_none")]
     updateid: Option<String>,
+    /// Internal Push Reliability tracking id. (Applied only to subscription updates generated
+    /// by Mozilla owned and consumed messages, like SendTab updates.)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    reliability_id: Option<String>,
 }
 
 impl NotificationRecord {
@@ -333,6 +337,7 @@ impl NotificationRecord {
             data: self.data,
             headers: self.headers.map(|m| m.into()),
             sortkey_timestamp: key.sortkey_timestamp,
+            reliability_id: None,
         })
     }