feat: Suspend Broken Tenants & Delete Broken Clients

migrations/1691518766_add-suspension.sql

@@ -0,0 +1,5 @@
+alter table public.tenants
+    add suspended bool not null default false;
+
+alter table public.tenants
+    add suspended_reason text;

migrations/new.sh

@@ -1,3 +1,2 @@
-#!/bin/bash
 DESCRIPTION=$1
 touch "./$(date +%s)_$DESCRIPTION.sql"

src/error.rs

@@ -171,6 +171,12 @@ pub enum Error {
 
     #[error("invalid apns creds")]
     BadApnsCredentials,
+
+    #[error("invalid device token")]
+    ClientDeleted,
+
+    #[error("invalid tenant configuration")]
+    TenantSuspended,
 }
 
 impl IntoResponse for Error {
@@ -525,6 +531,18 @@ impl IntoResponse for Error {
                     location: ErrorLocation::Path,
                 }
             ]),
+            Error::ClientDeleted => crate::handlers::Response::new_failure(StatusCode::ACCEPTED, vec![
+                ResponseError {
+                    name: "client_deleted".to_string(),
+                    message: "Request Accepted, client deleted due to invalid token".to_string(),
+                },
+            ], vec![]),
+            Error::TenantSuspended => crate::handlers::Response::new_failure(StatusCode::ACCEPTED, vec![
+                ResponseError {
+                    name: "tenant_suspended".to_string(),
+                    message: "Request Accepted, tenant suspended due to invalid configuration".to_string(),
+                },
+            ], vec![]),
             e => {
                 warn!("Error does not have response clause, {:?}", e);
 

src/handlers/get_tenant.rs

@@ -23,6 +23,8 @@ pub struct GetTenantResponse {
     enabled_providers: Vec<String>,
     apns_topic: Option<String>,
     apns_type: Option<ApnsType>,
+    suspended: bool,
+    suspended_reason: Option<String>,
 }
 
 pub async fn handler(
@@ -64,6 +66,8 @@ pub async fn handler(
         enabled_providers: tenant.providers().iter().map(Into::into).collect(),
         apns_topic: None,
         apns_type: None,
+        suspended: tenant.suspended,
+        suspended_reason: tenant.suspended_reason,
     };
 
     if providers.contains(&ProviderKind::Apns) {

src/handlers/push_message.rs

@@ -242,7 +242,7 @@
                     return Err((Error::MissmatchedTenantId, analytics));
                 }
 
                 return Err((Error::MissmatchedTenantId, None));
             }
         }
 
@@ -346,6 +346,10 @@
         "fetched tenant"
     );
 
+    if tenant.suspended {
+        return Err((Error::TenantSuspended, analytics.clone()));
+    }
+
     let mut provider = tenant
         .provider(&client.push_type)
         .map_err(|e| (e, analytics.clone()))?;
@@ -358,10 +362,29 @@
         "fetched provider"
     );
 
-    provider
-        .send_notification(client.token, body.payload)
-        .await
-        .map_err(|e| (e, analytics.clone()))?;
+    match provider.send_notification(client.token, body.payload).await {
+        Ok(_) => Ok(()),
+        Err(error) => match error {
+            Error::BadDeviceToken => {
+                state.client_store.delete_client(&tenant_id, &id);
+                Err(Error::ClientDeleted)
+            }
+            Error::BadApnsCredentials => {
+                state
+                    .tenant_store
+                    .suspend_tenant(&tenant_id, "Invalid APNS Credentials");
+                Err(Error::TenantSuspended)
+            }
+            Error::BadFcmApiKey => {
+                state
+                    .tenant_store
+                    .suspend_tenant(&tenant_id, "Invalid FCM Credentials");
+                Err(Error::TenantSuspended)
+            }
+            e => Err(e),
+        },
+    }
+    .map_err(|e| (e, analytics.clone()))?;
 
     info!(
         %request_id,
@@ -390,5 +413,5 @@
         return Ok(((StatusCode::ACCEPTED).into_response(), analytics));
     }
 
     Ok(((StatusCode::ACCEPTED).into_response(), None))
 }

src/providers/fcm.rs

@@ -1,11 +1,12 @@
 use {
     crate::{
         blob::DecryptedPayloadBlob,
+        error::Error,
         handlers::push_message::MessagePayload,
         providers::PushProvider,
     },
     async_trait::async_trait,
-    fcm::{MessageBuilder, NotificationBuilder},
+    fcm::{ErrorReason, FcmError, FcmResponse, MessageBuilder, NotificationBuilder},
     std::fmt::{Debug, Formatter},
     tracing::span,
 };
@@ -36,12 +37,12 @@ impl PushProvider for FcmProvider {
 
         let mut message_builder = MessageBuilder::new(self.api_key.as_str(), token.as_str());
 
-        if payload.is_encrypted() {
+        let result = if payload.is_encrypted() {
             message_builder.data(&payload)?;
 
             let fcm_message = message_builder.finalize();
 
-            let _ = self.client.send(fcm_message).await?;
+            self.client.send(fcm_message).await
         } else {
             let blob = DecryptedPayloadBlob::from_base64_encoded(payload.clone().blob)?;
 
@@ -55,10 +56,30 @@ impl PushProvider for FcmProvider {
 
             let fcm_message = message_builder.finalize();
 
-            let _ = self.client.send(fcm_message).await?;
+            self.client.send(fcm_message).await
+        };
+
+        match result {
+            Ok(val) => match val {
+                FcmResponse { error, .. } => {
+                    if let Some(error) = error {
+                        match error {
+                            ErrorReason::MissingRegistration
+                            | ErrorReason::InvalidRegistration
+                            | ErrorReason::NotRegistered => Err(Error::BadDeviceToken),
+                            ErrorReason::InvalidApnsCredential => Err(Error::BadApnsCredentials),
+                            _ => Ok(()),
+                        }
+                    } else {
+                        Ok(())
+                    }
+                }
+            },
+            Err(e) => match e {
+                FcmError::Unauthorized => Err(Error::BadFcmApiKey),
+                _ => Ok(()),
+            },
         }
-
-        Ok(())
     }
 }
 

src/stores/tenant.rs

@@ -95,6 +95,10 @@
     pub apns_key_id: Option<String>,
     pub apns_team_id: Option<String>,
 
+    // Suspension
+    pub suspended: bool,
+    pub suspended_reason: Option<String>,
+
     pub created_at: DateTime<Utc>,
     pub updated_at: DateTime<Utc>,
 }
@@ -261,6 +265,7 @@
         id: &str,
         params: TenantApnsUpdateAuth,
     ) -> Result<Tenant>;
+    async fn suspend_tenant(&self, id: &str, reason: &str) -> Result<()>;
 }
 
 #[async_trait]
@@ -373,6 +378,18 @@
 
         Ok(res)
     }
+
+    async fn suspend_tenant(&self, id: &str, reason: &str) -> Result<()> {
+        sqlx::query_as::<sqlx::postgres::Postgres, Tenant>(
+            "UPDATE public.tenants SET suspended = true, suspended_reason = $2::text WHERE id = \
+             $1 RETURNING *;",
+        )
+        .bind(id)
+        .bind(reason)
+        .await?;
+
+        Ok(())
+    }
 }
 
 #[cfg(not(feature = "multitenant"))]
@@ -381,7 +398,7 @@
 #[cfg(not(feature = "multitenant"))]
 impl DefaultTenantStore {
     pub fn new(config: Arc<Config>) -> Result<DefaultTenantStore> {
         Ok(DefaultTenantStore(Tenant {
             id: DEFAULT_TENANT_ID.to_string(),
             fcm_api_key: config.fcm_api_key.clone(),
             apns_type: config.apns_type,
@@ -435,4 +452,8 @@
     ) -> Result<Tenant> {
         panic!("Shouldn't have run in single tenant mode")
     }
+
+    async fn suspend_tenant(&self, id: &str, reason: &str) -> Result<()> {
+        panic!("Shouldn't have run in single tenant mode")
+    }
 }