Page MenuHomePhabricator
Files F3361601

D8676.diff
No OneTemporary
Actions

D8676.diff
View Options

diff --git a/services/commtest/tests/identity_onetime_key_tests.rs b/services/commtest/tests/identity_onetime_key_tests.rs
new file mode 100644
--- /dev/null
+++ b/services/commtest/tests/identity_onetime_key_tests.rs
@@ -0,0 +1,35 @@
+mod client {
+ tonic::include_proto!("identity.client");
+}
+mod auth_proto {
+ tonic::include_proto!("identity.authenticated");
+}
+use client::identity_client_service_client::IdentityClientServiceClient;
+use client::UploadOneTimeKeysRequest;
+use commtest::identity::device::create_device;
+
+#[tokio::test]
+async fn verify_access_token() {
+ let device_info = create_device().await;
+
+ let mut identity_client =
+ IdentityClientServiceClient::connect("http://127.0.0.1:50054")
+ .await
+ .expect("Couldn't connect to identitiy service");
+
+ let upload_request = UploadOneTimeKeysRequest {
+ user_id: device_info.user_id,
+ device_id: device_info.device_id,
+ access_token: device_info.access_token,
+ content_one_time_pre_keys: vec![
+ "content1".to_string(),
+ "content2".to_string(),
+ ],
+ notif_one_time_pre_keys: vec!["notif1".to_string(), "notif2".to_string()],
+ };
+
+ identity_client
+ .upload_one_time_keys(upload_request)
+ .await
+ .unwrap();
+}
diff --git a/services/identity/src/client_service.rs b/services/identity/src/client_service.rs
--- a/services/identity/src/client_service.rs
+++ b/services/identity/src/client_service.rs
@@ -791,7 +791,6 @@
self
.client
.append_one_time_prekeys(
- message.user_id,
message.device_id,
message.content_one_time_pre_keys,
message.notif_one_time_pre_keys,
diff --git a/services/identity/src/constants.rs b/services/identity/src/constants.rs
--- a/services/identity/src/constants.rs
+++ b/services/identity/src/constants.rs
@@ -88,7 +88,7 @@
// One time keys table, which need to exist in their own table to ensure
// atomicity of additions and removals
pub mod content_one_time_keys_table {
- pub const NAME: &'static str = "identity-content-one-time-keys";
+ pub const NAME: &'static str = "identity-content-onetime-keys";
pub const PARTITION_KEY: &'static str = "deviceID";
pub const DEVICE_ID: &'static str = PARTITION_KEY;
pub const SORT_KEY: &'static str = "oneTimeKey";
@@ -96,7 +96,7 @@
}
pub mod notif_one_time_keys_table {
- pub const NAME: &'static str = "identity-notif-one-time-keys";
+ pub const NAME: &'static str = "identity-notif-onetime-keys";
pub const PARTITION_KEY: &'static str = "deviceID";
pub const DEVICE_ID: &'static str = PARTITION_KEY;
pub const SORT_KEY: &'static str = "oneTimeKey";
diff --git a/services/identity/src/database.rs b/services/identity/src/database.rs
--- a/services/identity/src/database.rs
+++ b/services/identity/src/database.rs
@@ -4,6 +4,7 @@
use std::str::FromStr;
use std::sync::Arc;
+use crate::ddb_utils::into_onetime_put_requests;
use crate::error::{DBItemAttributeError, DBItemError, Error};
use aws_config::SdkConfig;
use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest};
@@ -25,13 +26,11 @@
NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY,
RESERVED_USERNAMES_TABLE, RESERVED_USERNAMES_TABLE_PARTITION_KEY,
USERS_TABLE, USERS_TABLE_DEVICES_ATTRIBUTE,
- USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME,
- USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME,
@@ -295,42 +294,24 @@
pub async fn append_one_time_prekeys(
&self,
- user_id: String,
device_id: String,
- content_one_time_keys: Vec<String>,
- notif_one_time_keys: Vec<String>,
+ content_onetime_keys: Vec<String>,
+ notif_onetime_keys: Vec<String>,
) -> Result<(), Error> {
- let notif_keys_av: Vec<AttributeValue> = notif_one_time_keys
- .into_iter()
- .map(AttributeValue::S)
- .collect();
- let content_keys_av: Vec<AttributeValue> = content_one_time_keys
- .into_iter()
- .map(AttributeValue::S)
- .collect();
+ use crate::constants::{
+ content_one_time_keys_table, notif_one_time_keys_table,
+ };
- let update_expression =
- format!("SET {0}.#{1}.{2} = list_append({0}.#{1}.{2}, :n), {0}.#{1}.{3} = list_append({0}.#{1}.{3}, :i)",
- USERS_TABLE_DEVICES_ATTRIBUTE,
- "deviceID",
- USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME,
- USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME
- );
- let expression_attribute_names =
- HashMap::from([(format!("#{}", "deviceID"), device_id)]);
- let expression_attribute_values = HashMap::from([
- (":n".to_string(), AttributeValue::L(notif_keys_av)),
- (":i".to_string(), AttributeValue::L(content_keys_av)),
- ]);
+ let content_otk_requests =
+ into_onetime_put_requests(device_id.clone(), content_onetime_keys);
+ let notif_otk_requests =
+ into_onetime_put_requests(device_id.clone(), notif_onetime_keys);
self
.client
- .update_item()
- .table_name(USERS_TABLE)
- .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id))
- .update_expression(update_expression)
- .set_expression_attribute_names(Some(expression_attribute_names))
- .set_expression_attribute_values(Some(expression_attribute_values))
+ .batch_write_item()
+ .request_items(content_one_time_keys_table::NAME, content_otk_requests)
+ .request_items(notif_one_time_keys_table::NAME, notif_otk_requests)
.send()
.await
.map_err(|e| Error::AwsSdk(e.into()))?;
@@ -344,14 +325,19 @@
flattened_device_key_upload: FlattenedDeviceKeyUpload,
social_proof: Option<String>,
) -> Result<(), Error> {
+ // Avoid borrowing from lifetime of flattened_device_key_upload
+ let device_id = flattened_device_key_upload.device_id_key.clone();
+ let content_onetime_keys =
+ flattened_device_key_upload.content_onetime_keys.clone();
+ let notif_onetime_keys =
+ flattened_device_key_upload.notif_onetime_keys.clone();
+
let device_info =
- create_device_info(flattened_device_key_upload.clone(), social_proof);
+ create_device_info(flattened_device_key_upload, social_proof);
let update_expression =
format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",);
- let expression_attribute_names = HashMap::from([(
- format!("#{}", "deviceID"),
- flattened_device_key_upload.device_id_key,
- )]);
+ let expression_attribute_names =
+ HashMap::from([(format!("#{}", "deviceID"), device_id.clone())]);
let expression_attribute_values =
HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]);
@@ -367,6 +353,26 @@
.await
.map_err(|e| Error::AwsSdk(e.into()))?;
+ let content_otk_requests =
+ into_onetime_put_requests(device_id.clone(), content_onetime_keys);
+ let notif_otk_requests =
+ into_onetime_put_requests(device_id.clone(), notif_onetime_keys);
+
+ self
+ .client
+ .batch_write_item()
+ .request_items(
+ crate::constants::content_one_time_keys_table::NAME,
+ content_otk_requests,
+ )
+ .request_items(
+ crate::constants::notif_one_time_keys_table::NAME,
+ notif_otk_requests,
+ )
+ .send()
+ .await
+ .map_err(|e| Error::AwsSdk(e.into()))?;
+
Ok(())
}
@@ -1042,16 +1048,6 @@
.to_string(),
AttributeValue::S(flattened_device_key_upload.content_prekey_signature),
),
- (
- USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(),
- AttributeValue::L(
- flattened_device_key_upload
- .content_onetime_keys
- .into_iter()
- .map(AttributeValue::S)
- .collect(),
- ),
- ),
(
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(flattened_device_key_upload.notif_prekey),
@@ -1060,16 +1056,6 @@
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(flattened_device_key_upload.notif_prekey_signature),
),
- (
- USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(),
- AttributeValue::L(
- flattened_device_key_upload
- .notif_onetime_keys
- .into_iter()
- .map(AttributeValue::S)
- .collect(),
- ),
- ),
]);
if let Some(social_proof) = social_proof {
diff --git a/services/identity/src/ddb_utils.rs b/services/identity/src/ddb_utils.rs
new file mode 100644
--- /dev/null
+++ b/services/identity/src/ddb_utils.rs
@@ -0,0 +1,38 @@
+use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest};
+use std::collections::HashMap;
+use std::iter::IntoIterator;
+
+fn create_onetime_key_put_request(
+ device_id: String,
+ onetime_key: String,
+) -> WriteRequest {
+ // The content and notif onetime key tables use the same partition and sort
+ // key
+ use crate::constants::content_one_time_keys_table::*;
+
+ let builder = PutRequest::builder();
+ let attrs = HashMap::from([
+ (PARTITION_KEY.to_string(), AttributeValue::S(device_id)),
+ (SORT_KEY.to_string(), AttributeValue::S(onetime_key)),
+ ]);
+
+ let put_request = builder.set_item(Some(attrs)).build();
+
+ WriteRequest::builder().put_request(put_request).build()
+}
+
+pub fn into_onetime_put_requests<T>(
+ device_id: String,
+ onetime_keys: T,
+) -> Vec<WriteRequest>
+where
+ T: IntoIterator,
+ <T as IntoIterator>::Item: ToString,
+{
+ onetime_keys
+ .into_iter()
+ .map(|otk| {
+ create_onetime_key_put_request(device_id.clone(), otk.to_string())
+ })
+ .collect()
+}
diff --git a/services/identity/src/main.rs b/services/identity/src/main.rs
--- a/services/identity/src/main.rs
+++ b/services/identity/src/main.rs
@@ -9,6 +9,7 @@
mod config;
pub mod constants;
mod database;
+pub mod ddb_utils;
pub mod error;
mod grpc_services;
mod id;

File Metadata

Mime Type
text/plain
Expires
Mon, Nov 25, 6:24 PM (20 h, 50 m)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
2580653
Default Alt Text
D8676.diff (10 KB)

Event Timeline