From d967211af62d1020a09793091e42dead95cace66 Mon Sep 17 00:00:00 2001
From: Varun Dhananjaya <varun.dhananjaya@gmail.com>
Date: Fri, 14 Jul 2023 15:53:15 -0400
Subject: [PATCH] [identity] helper functions for removing data from DDB

Summary:
on logout, we need to delete/remove the access token and keys associated with the logged out device.

the first helper function removes the device's keys from the users table with an update expression.

the second helper function deletes the device's access token from the access tokens table by calling delete_item().

Test Plan: created test data in local DDB and called the helper functions from main to confirm that the keys and token were removed

Reviewers: bartek, jon

Reviewed By: bartek, jon

Subscribers: ashoat, tomek

Differential Revision: https://phab.comm.dev/D8500
---
 services/identity/src/database.rs | 51 +++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/services/identity/src/database.rs b/services/identity/src/database.rs
index 6f73018468..c32820fd85 100644
--- a/services/identity/src/database.rs
+++ b/services/identity/src/database.rs
@@ -370,6 +370,33 @@ impl DatabaseClient {
     Ok(())
   }
 
+  pub async fn remove_device_from_users_table(
+    &self,
+    user_id: String,
+    device_id_key: String,
+  ) -> Result<(), Error> {
+    let update_expression =
+      format!("REMOVE {}.{}", USERS_TABLE_DEVICES_ATTRIBUTE, ":deviceID");
+
+    let expression_attribute_values = HashMap::from([(
+      ":deviceID".to_string(),
+      AttributeValue::S(device_id_key),
+    )]);
+
+    self
+      .client
+      .update_item()
+      .table_name(USERS_TABLE)
+      .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id))
+      .update_expression(update_expression)
+      .set_expression_attribute_values(Some(expression_attribute_values))
+      .send()
+      .await
+      .map_err(|e| Error::AwsSdk(e.into()))?;
+
+    Ok(())
+  }
+
   pub async fn update_user_password(
     &self,
     user_id: String,
@@ -553,6 +580,30 @@ impl DatabaseClient {
       .map_err(|e| Error::AwsSdk(e.into()))
   }
 
+  pub async fn delete_access_token_data(
+    &self,
+    user_id: String,
+    device_id_key: String,
+  ) -> Result<(), Error> {
+    self
+      .client
+      .delete_item()
+      .table_name(ACCESS_TOKEN_TABLE)
+      .key(
+        ACCESS_TOKEN_TABLE_PARTITION_KEY.to_string(),
+        AttributeValue::S(user_id),
+      )
+      .key(
+        ACCESS_TOKEN_SORT_KEY.to_string(),
+        AttributeValue::S(device_id_key),
+      )
+      .send()
+      .await
+      .map_err(|e| Error::AwsSdk(e.into()))?;
+
+    Ok(())
+  }
+
   pub async fn username_taken(&self, username: String) -> Result<bool, Error> {
     let result = self
       .get_user_id_from_user_info(username, AuthType::Password)