feat: Add ability to track messages based on known VAPID keys

autoendpoint/src/extractors/subscription.rs

@@ -35,6 +35,10 @@ pub struct Subscription {
     pub user: User,
     pub channel_id: Uuid,
     pub vapid: Option<VapidHeaderWithKey>,
+    /// Should this subscription update be tracked internally?
+    /// (This should ONLY be applied for messages that match known
+    /// Mozilla provided VAPID public keys.)
+    pub tracking_id: Option<String>,
 }
 
 impl FromRequest for Subscription {
@@ -69,6 +73,11 @@ impl FromRequest for Subscription {
                 .transpose()?;
 
             trace!("raw vapid: {:?}", &vapid);
+            let trackable = if let Some(vapid) = &vapid {
+                app_state.settings.is_trackable(vapid)
+            } else {
+                false
+            };
 
             // 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
@@ -127,6 +136,13 @@ impl FromRequest for Subscription {
                 user,
                 channel_id,
                 vapid,
+                // Eventually this will be replaced by the generated, opaque
+                // tracking ID, for now, just use a simple UUID4 string.
+                tracking_id: if trackable {
+                    Some(uuid::Uuid::new_v4().as_simple().to_string())
+                } else {
+                    None
+                },
             })
         }
         .boxed_local()

autoendpoint/src/routers/common.rs

@@ -239,6 +239,7 @@ pub mod tests {
                 user,
                 channel_id: channel_id(),
                 vapid: None,
+                tracking_id: None,
             },
             headers: NotificationHeaders {
                 ttl: 0,

autoendpoint/src/settings.rs

@@ -5,6 +5,7 @@ use fernet::{Fernet, MultiFernet};
 use serde::Deserialize;
 use url::Url;
 
+use crate::headers::vapid::VapidHeaderWithKey;
 use crate::routers::apns::settings::ApnsSettings;
 use crate::routers::fcm::settings::FcmSettings;
 #[cfg(feature = "stub")]
@@ -31,6 +32,17 @@ pub struct Settings {
 
     pub vapid_aud: Vec<String>,
 
+    /// A stringified JSON list of VAPID public keys which should be tracked internally.
+    /// This should ONLY include Mozilla generated and consumed messages (e.g. "SendToTab", etc.)
+    /// These keys should be specified in stripped, b64encoded, X962 format (e.g. a single line of
+    /// base64 encoded data without padding).
+    /// You can use `scripts/convert_pem_to_x962.py` to easily convert EC Public keys stored in
+    /// PEM format into appropriate x962 format.
+    pub tracking_keys: String,
+    /// Cached, parsed tracking keys.
+    //TODO: convert to decoded Vec<u8>?
+    pub tracking_vapid_pubs: Vec<String>,
+
     pub max_data_bytes: usize,
     pub crypto_keys: String,
     pub auth_keys: String,
@@ -64,13 +76,15 @@ impl Default for Settings {
                 "https://push.services.mozilla.org".to_string(),
                 "http://127.0.0.1:9160".to_string(),
             ],
+            tracking_vapid_pubs: vec![],
             // max data is a bit hard to figure out, due to encryption. Using something
             // like pywebpush, if you encode a block of 4096 bytes, you'll get a
             // 4216 byte data block. Since we're going to be receiving this, we have to
             // presume base64 encoding, so we can bump things up to 5630 bytes max.
             max_data_bytes: 5630,
             crypto_keys: format!("[{}]", Fernet::generate_key()),
             auth_keys: r#"["AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAB="]"#.to_string(),
+            tracking_keys: r#"[]"#.to_string(),
             human_logs: false,
             connection_timeout_millis: 1000,
             request_timeout_millis: 3000,
@@ -100,9 +114,7 @@ impl Settings {
         // down to the sub structures.
         config = config.add_source(Environment::with_prefix(ENV_PREFIX).separator("__"));
 
-        let built = config.build()?;
-
-        built.try_deserialize::<Self>().map_err(|error| {
+        let mut built: Self = config.build()?.try_deserialize::<Self>().map_err(|error| {
             match error {
                 // Configuration errors are not very sysop friendly, Try to make them
                 // a bit more 3AM useful.
@@ -121,7 +133,12 @@ impl Settings {
                     error
                 }
             }
-        })
+        })?;
+
+        // cache the tracking keys we've built.
+        built.tracking_vapid_pubs = built.tracking_keys();
+
+        Ok(built)
     }
 
     /// Convert a string like `[item1,item2]` into a iterator over `item1` and `item2`.
@@ -158,6 +175,26 @@ impl Settings {
             .collect()
     }
 
+    /// Get the list of tracking public keys
+    pub fn tracking_keys(&self) -> Vec<String> {
+        // return the cached version if present.
+        if !self.tracking_vapid_pubs.is_empty() {
+            return self.tracking_vapid_pubs.clone();
+        };
+        let keys = &self.tracking_keys.replace(['"', ' '], "");
+        Self::read_list_from_str(keys, "Invalid AUTOEND_TRACKING_KEYS")
+            .map(|v| v.to_owned())
+            .collect()
+    }
+
+    /// Very simple string check to see if the Public Key specified in the Vapid header
+    /// matches the set of trackable keys.
+    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.tracking_keys().contains(&vapid.public_key)
+    }
+
     /// Get the URL for this endpoint server
     pub fn endpoint_url(&self) -> Url {
         let endpoint = if self.endpoint_url.is_empty() {
@@ -172,7 +209,10 @@ impl Settings {
 #[cfg(test)]
 mod tests {
     use super::Settings;
-    use crate::error::ApiResult;
+    use crate::{
+        error::ApiResult,
+        headers::vapid::{VapidHeader, VapidHeaderWithKey},
+    };
 
     #[test]
     fn test_auth_keys() -> ApiResult<()> {
@@ -198,6 +238,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_tracking_keys() -> ApiResult<()> {
+        let mut settings = Settings{
+            tracking_keys: r#"["BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7vI"]"#.to_owned(),
+            ..Default::default()
+        };
+
+        let test_header = VapidHeaderWithKey {
+            vapid: VapidHeader {
+                scheme: "".to_owned(),
+                token: "".to_owned(),
+                version_data: crate::headers::vapid::VapidVersionData::Version1,
+            },
+            public_key: "BLMymkOqvT6OZ1o9etCqV4jGPkvOXNz5FdBjsAR9zR5oeCV1x5CBKuSLTlHon-H_boHTzMtMoNHsAGDlDB6X7vI".to_owned()
+        };
+
+        let result = settings.tracking_keys();
+        assert!(!result.is_empty());
+
+        // emulate Settings.with_env_and_config_file()
+        settings.tracking_vapid_pubs = result;
+
+        assert!(settings.is_trackable(&test_header));
+        Ok(())
+    }
+
     #[test]
     fn test_endpoint_url() -> ApiResult<()> {
         let example = "https://example.org/";

scripts/convert_pem_to_x962.py

@@ -0,0 +1,31 @@
+"""
+Convert a EC Public key in PEM format into an b64 x962 string.
+
+Autopush will try to scan for known VAPID public keys to track. These keys
+are specified in the header as x962 formatted strings. X962 is effectively
+"raw" format and contains the two longs that are the coordinates for the
+public key.
+
+"""
+import base64
+import sys
+
+from typing import cast
+from cryptography.hazmat.primitives.asymmetric import ec, utils as ec_utils
+from cryptography.hazmat.primitives import serialization
+
+try:
+    content = open(sys.argv[1], "rb").read()
+    pubkey = serialization.load_pem_public_key(content)
+except IndexError:
+    print ("Please specify a public key PEM file to convert.")
+    exit()
+
+pk_string = cast(ec.EllipticCurvePublicKey, pubkey).public_bytes(
+    serialization.Encoding.X962,
+    serialization.PublicFormat.UncompressedPoint
+)
+
+pk_string = base64.urlsafe_b64encode(pk_string).strip(b'=')
+
+print(f"{pk_string.decode()}")