Page MenuHomePhabricator
Files F3720849

D8127.id27571.diff
No OneTemporary
Actions

D8127.id27571.diff
View Options

diff --git a/services/commtest/tests/identity_access_tokens_tests.rs b/services/commtest/tests/identity_access_tokens_tests.rs
--- a/services/commtest/tests/identity_access_tokens_tests.rs
+++ b/services/commtest/tests/identity_access_tokens_tests.rs
@@ -5,7 +5,8 @@
use proto::identity_client_service_client::IdentityClientServiceClient;
use proto::{
DeviceKeyUpload, IdentityKeyInfo, PreKey, RegistrationFinishRequest,
- RegistrationStartRequest, VerifyUserAccessTokenRequest,
+ RegistrationStartRequest, UploadOneTimeKeysRequest,
+ VerifyUserAccessTokenRequest,
};
use rand::{distributions::Alphanumeric, Rng};
@@ -113,3 +114,27 @@
assert_eq!(response.into_inner().token_valid, true);
}
+
+#[tokio::test]
+async fn upload_one_time_keys() {
+ 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!["a".to_string(), "b".to_string()],
+ notif_one_time_pre_keys: vec!["c".to_string(), "d".to_string()],
+ };
+
+ // This send will fail if the one-time keys weren't successfully added
+ 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
@@ -35,7 +35,7 @@
use moka::future::Cache;
use rand::rngs::OsRng;
use tonic::Response;
-use tracing::error;
+use tracing::{debug, error};
use self::client_proto::ReservedRegistrationStartRequest;
@@ -749,9 +749,41 @@
async fn upload_one_time_keys(
&self,
- _request: tonic::Request<UploadOneTimeKeysRequest>,
+ request: tonic::Request<UploadOneTimeKeysRequest>,
) -> Result<tonic::Response<Empty>, tonic::Status> {
- unimplemented!();
+ let message = request.into_inner();
+
+ debug!("Validating token: {:?}", message);
+ let token_valid = self
+ .client
+ .verify_access_token(
+ message.user_id.clone(),
+ message.device_id.clone(),
+ message.access_token,
+ )
+ .await
+ .map_err(handle_db_error)?;
+
+ if !token_valid {
+ return Err(tonic::Status::unauthenticated("Invalid token"));
+ }
+
+ debug!(
+ "Attempting to update one time keys for user: {}",
+ message.user_id
+ );
+ self
+ .client
+ .append_one_time_prekeys(
+ message.user_id,
+ message.device_id,
+ message.content_one_time_pre_keys,
+ message.notif_one_time_pre_keys,
+ )
+ .await
+ .map_err(handle_db_error)?;
+
+ Ok(tonic::Response::new(Empty {}))
}
async fn refresh_user_pre_keys(
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
@@ -218,6 +218,51 @@
.await
}
+ 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>,
+ ) -> 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();
+
+ 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)),
+ ]);
+
+ 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))
+ .send()
+ .await
+ .map_err(|e| Error::AwsSdk(e.into()))?;
+
+ Ok(())
+ }
+
async fn add_device_to_users_table(
&self,
user_id: String,

File Metadata

Mime Type
text/plain
Expires
Thu, Jan 9, 12:41 PM (43 m, 7 s)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
2820028
Default Alt Text
D8127.id27571.diff (4 KB)

Event Timeline