feat: Add ability to track messages based on known VAPID keys (#739)

_*Note*_: This introduces a new setting for autoendpoint:
`tracking_keys`
This is a JSON formatted list of the "raw"/x962 formatted VAPID public
keys that should be monitored for Push tracking and follows the same
format as other key data fields.

The newly added `script/convert_pem_to_x962.py` script aids by reading a
VAPID public key PEM file and outputing a x962 formatted string.

Closes: [SYNC-4349](https://mozilla-hub.atlassian.net/browse/SYNC-4349)

[SYNC-4349]:
https://mozilla-hub.atlassian.net/browse/SYNC-4349?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

autoendpoint/src/extractors/subscription.rs

@@ -35,6 +35,12 @@ pub struct Subscription {
     pub user: User,
     pub channel_id: Uuid,
     pub vapid: Option<VapidHeaderWithKey>,
+    /// A stored value here indicates that the subscription update
+    /// should 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 +75,11 @@ impl FromRequest for Subscription {
                 .transpose()?;
 
             trace!("raw vapid: {:?}", &vapid);
+            let trackable = if let Some(vapid) = &vapid {
+                app_state.reliability.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
@@ -123,10 +134,14 @@ 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,
             })
         }
         .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/server.rs

@@ -18,7 +18,6 @@ use autopush_common::{
     middleware::sentry::SentryWrapper,
 };
 
-use crate::error::{ApiError, ApiErrorKind, ApiResult};
 use crate::metrics;
 #[cfg(feature = "stub")]
 use crate::routers::stub::router::StubRouter;
@@ -32,6 +31,10 @@ use crate::routes::{
     webpush::{delete_notification_route, webpush_route},
 };
 use crate::settings::Settings;
+use crate::{
+    error::{ApiError, ApiErrorKind, ApiResult},
+    settings::VapidTracker,
+};
 
 #[derive(Clone)]
 pub struct AppState {
@@ -45,6 +48,7 @@ pub struct AppState {
     pub apns_router: Arc<ApnsRouter>,
     #[cfg(feature = "stub")]
     pub stub_router: Arc<StubRouter>,
+    pub reliability: Arc<VapidTracker>,
 }
 
 pub struct Server;
@@ -105,6 +109,7 @@ impl Server {
             )
             .await?,
         );
+        let reliability = Arc::new(VapidTracker(settings.tracking_keys()));
         #[cfg(feature = "stub")]
         let stub_router = Arc::new(StubRouter::new(settings.stub.clone())?);
         let app_state = AppState {
@@ -117,6 +122,7 @@ impl Server {
             apns_router,
             #[cfg(feature = "stub")]
             stub_router,
+            reliability,
         };
 
         spawn_pool_periodic_reporter(

autoendpoint/src/settings.rs

@@ -1,10 +1,12 @@
 //! Application settings
 
+use actix_http::header::HeaderMap;
 use config::{Config, ConfigError, Environment, File};
 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 +33,14 @@ 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,
+
     pub max_data_bytes: usize,
     pub crypto_keys: String,
     pub auth_keys: String,
@@ -71,6 +81,7 @@ impl Default for Settings {
             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 +111,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 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 +130,9 @@ impl Settings {
                     error
                 }
             }
-        })
+        })?;
+
+        Ok(built)
     }
 
     /// Convert a string like `[item1,item2]` into a iterator over `item1` and `item2`.
@@ -158,6 +169,17 @@ impl Settings {
             .collect()
     }
 
+    /// Get the list of tracking public keys
+    // TODO: this should return a Vec<[u8]> so that key formatting errors do not cause
+    // false rejections. This is not a problem now since we have access to the source
+    // 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()
+    }
+
     /// Get the URL for this endpoint server
     pub fn endpoint_url(&self) -> Url {
         let endpoint = if self.endpoint_url.is_empty() {
@@ -169,10 +191,42 @@ impl Settings {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct VapidTracker(pub Vec<String>);
+impl VapidTracker {
+    /// 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.0.contains(&vapid.public_key)
+    }
+
+    /// Extract the message Id from the headers (if present), otherwise just make one up.
+    pub fn get_tracking_id(&self, headers: &HeaderMap) -> String {
+        headers
+            .get("X-MessageId")
+            .and_then(|v|
+                // TODO: we should convert the public key string to a bitarray
+                // this would prevent any formatting errors from falsely rejecting
+                // the key. We're ok with comparing strings because we currently
+                // have access to the same public key value string that is being
+                // used, but that may not always be the case.
+                v.to_str().ok())
+            .map(|v| v.to_owned())
+            .unwrap_or_else(|| uuid::Uuid::new_v4().as_simple().to_string())
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use super::Settings;
-    use crate::error::ApiResult;
+    use actix_http::header::{HeaderMap, HeaderName, HeaderValue};
+
+    use super::{Settings, VapidTracker};
+    use crate::{
+        error::ApiResult,
+        headers::vapid::{VapidHeader, VapidHeaderWithKey},
+    };
 
     #[test]
     fn test_auth_keys() -> ApiResult<()> {
@@ -252,4 +306,49 @@ mod tests {
             env::remove_var(&timeout);
         }
     }
+
+    #[test]
+    fn test_tracking_keys() -> ApiResult<()> {
+        let 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 key_set = settings.tracking_keys();
+        assert!(!key_set.is_empty());
+
+        let reliability = VapidTracker(key_set);
+        assert!(reliability.is_trackable(&test_header));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_tracking_id() -> ApiResult<()> {
+        let mut headers = HeaderMap::new();
+        let keys = Vec::new();
+        let reliability = VapidTracker(keys);
+
+        let key = reliability.get_tracking_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);
+        assert_eq!(key, "123foobar456".to_owned());
+
+        Ok(())
+    }
 }

scripts/convert_pem_to_x962.py

@@ -0,0 +1,32 @@
+"""
+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
+from cryptography.hazmat.primitives import serialization
+
+try:
+    with open(sys.argv[1], "rb") as fp:
+        content = fp.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()}")