diff --git a/services/commtest/src/identity/device.rs b/services/commtest/src/identity/device.rs --- a/services/commtest/src/identity/device.rs +++ b/services/commtest/src/identity/device.rs @@ -3,10 +3,10 @@ mod proto { tonic::include_proto!("identity.client"); } +use proto as client; use proto::{ identity_client_service_client::IdentityClientServiceClient, DeviceKeyUpload, - DeviceType, IdentityKeyInfo, PreKey, RegistrationFinishRequest, - RegistrationStartRequest, + IdentityKeyInfo, PreKey, RegistrationFinishRequest, RegistrationStartRequest, }; pub struct DeviceInfo { @@ -51,7 +51,7 @@ }), onetime_content_prekeys: Vec::new(), onetime_notif_prekeys: Vec::new(), - device_type: DeviceType::Keyserver.into(), + device_type: client::DeviceType::Keyserver.into(), }), }; diff --git a/services/commtest/tests/identity_keyserver_tests.rs b/services/commtest/tests/identity_keyserver_tests.rs new file mode 100644 --- /dev/null +++ b/services/commtest/tests/identity_keyserver_tests.rs @@ -0,0 +1,85 @@ +mod proto { + tonic::include_proto!("identity.client"); +} +use proto as client; +mod auth_proto { + tonic::include_proto!("identity.authenticated"); +} +use auth_proto::identity_client_service_client::IdentityClientServiceClient as AuthClient; +use auth_proto::OutboundKeysForUserRequest; +use client::UploadOneTimeKeysRequest; +use commtest::identity::device::create_device; +use tonic::{transport::Endpoint, Request}; + +#[tokio::test] +async fn set_prekey() { + let device_info = create_device().await; + + let channel = Endpoint::from_static("http://[::1]:50054") + .connect() + .await + .unwrap(); + + let mut client = + AuthClient::with_interceptor(channel, |mut request: Request<()>| { + let metadata = request.metadata_mut(); + metadata.insert("user_id", device_info.user_id.parse().unwrap()); + metadata.insert("device_id", device_info.device_id.parse().unwrap()); + metadata + .insert("access_token", device_info.access_token.parse().unwrap()); + Ok(request) + }); + + let upload_request = UploadOneTimeKeysRequest { + user_id: device_info.user_id.to_string(), + device_id: device_info.device_id.to_string(), + access_token: device_info.access_token.to_string(), + content_one_time_pre_keys: vec!["content1".to_string()], + notif_one_time_pre_keys: vec!["notif1".to_string()], + }; + + let mut unauthenticated_client = + proto::identity_client_service_client::IdentityClientServiceClient::connect("http://127.0.0.1:50054") + .await + .expect("Couldn't connect to identitiy service"); + + unauthenticated_client + .upload_one_time_keys(upload_request) + .await + .expect("Failed to upload keys"); + + // Currently allowed to request your own outbound keys + let keyserver_request = OutboundKeysForUserRequest { + user_id: device_info.user_id.clone(), + }; + + println!("Getting keyserver info for user, {}", device_info.user_id); + let first_reponse = client + .get_keyserver_keys(keyserver_request.clone()) + .await + .expect("Second keyserver keys request failed") + .into_inner() + .keyserver_info + .unwrap(); + + assert_eq!( + first_reponse.onetime_content_prekey, + Some("content1".to_string()) + ); + assert_eq!( + first_reponse.onetime_notif_prekey, + Some("notif1".to_string()) + ); + + let second_reponse = client + .get_keyserver_keys(keyserver_request) + .await + .expect("Second keyserver keys request failed") + .into_inner() + .keyserver_info + .unwrap(); + + // The one time keys should be exhausted + assert_eq!(second_reponse.onetime_content_prekey, None); + assert_eq!(second_reponse.onetime_notif_prekey, None); +} 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,7 +4,9 @@ use std::str::FromStr; use std::sync::Arc; -use crate::ddb_utils::{into_one_time_put_requests, OlmAccountType}; +use crate::ddb_utils::{ + create_one_time_key_partition_key, into_one_time_put_requests, OlmAccountType, +}; use crate::error::{DBItemAttributeError, DBItemError, Error}; use aws_config::SdkConfig; use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest}; @@ -39,6 +41,7 @@ USERS_TABLE_USERNAME_ATTRIBUTE, USERS_TABLE_USERNAME_INDEX, USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE, USERS_TABLE_WALLET_ADDRESS_INDEX, }; +use crate::error::{AttributeValueFromHashMap, FromAttributeValue}; use crate::id::generate_uuid; use crate::nonce::NonceData; use crate::token::{AccessTokenData, AuthType}; @@ -102,6 +105,22 @@ } } +// This is very similar to the protobuf definitions, however, +// coupling the protobuf schema to the database API should be avoided. +pub struct PreKey { + pub prekey: String, + pub prekey_signature: String, +} +pub struct OutboundKeys { + pub key_payload: String, + pub key_payload_signature: String, + pub social_proof: Option, + pub content_prekey: PreKey, + pub notif_prekey: PreKey, + pub content_one_time_key: Option, + pub notif_one_time_key: Option, +} + #[derive(Clone)] pub struct DatabaseClient { client: Arc, @@ -247,6 +266,197 @@ .await } + pub async fn get_keyserver_keys_for_user( + &self, + user_id: &str, + ) -> Result, Error> { + // DynamoDB doesn't have a way to "pop" a value from a list, so we must + // first read in user info, then update one_time_keys with value we + // gave to requester + let user_info = self + .get_item_from_users_table(&user_id) + .await? + .item + .ok_or(Error::MissingItem)?; + + let devices = user_info + .get(USERS_TABLE_DEVICES_ATTRIBUTE) + .ok_or(Error::MissingItem)? + .to_hashmap(USERS_TABLE_DEVICES_ATTRIBUTE)?; + + let mut maybe_keyserver_id = None; + for (device_id, device_info) in devices { + let device_type = device_info + .to_hashmap("device_id")? + .get(USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME) + .ok_or(Error::MissingItem)? + .to_string(USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME)?; + + if device_type == "keyserver" { + maybe_keyserver_id = Some(device_id); + break; + } + } + + // Assert that the user has a keyserver, if they don't return None + let keyserver_id = match maybe_keyserver_id { + None => return Ok(None), + Some(id) => id, + }; + + let keyserver = devices.get_map(keyserver_id)?; + let notif_one_time_key: Option = self + .get_onetime_key(keyserver_id, OlmAccountType::Notification) + .await?; + let content_one_time_key: Option = self + .get_onetime_key(keyserver_id, OlmAccountType::Content) + .await?; + + debug!( + "Able to get notif key for keyserver {}: {}", + keyserver_id, + notif_one_time_key.is_some() + ); + debug!( + "Able to get content key for keyserver {}: {}", + keyserver_id, + content_one_time_key.is_some() + ); + + let content_prekey = keyserver + .get_string(USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME)?; + let content_prekey_signature = keyserver.get_string( + USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + )?; + let notif_prekey = keyserver + .get_string(USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME)?; + let notif_prekey_signature = keyserver.get_string( + USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + )?; + let key_payload = keyserver + .get_string(USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME)? + .to_string(); + let key_payload_signature = keyserver + .get_string(USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME)? + .to_string(); + let social_proof = keyserver + .get(USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME) + .map(|s| { + s.to_string(USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME) + .ok() + }) + .flatten() + .map(|s| s.to_owned()); + + let full_content_prekey = PreKey { + prekey: content_prekey.to_string(), + prekey_signature: content_prekey_signature.to_string(), + }; + + let full_notif_prekey = PreKey { + prekey: notif_prekey.to_string(), + prekey_signature: notif_prekey_signature.to_string(), + }; + + let outbound_payload = OutboundKeys { + key_payload, + key_payload_signature, + social_proof, + content_prekey: full_content_prekey, + notif_prekey: full_notif_prekey, + content_one_time_key, + notif_one_time_key, + }; + + return Ok(Some(outbound_payload)); + } + + /// Will "mint" a single onetime key by attempting to successfully deleting + /// a key + pub async fn get_onetime_key( + &self, + device_id: &str, + account_type: OlmAccountType, + ) -> Result, Error> { + use crate::constants::one_time_keys_table as otk_table; + + let query_result = self.get_onetime_keys(device_id, account_type).await?; + let items = query_result.items(); + + // If no onetime keys exists, return none early + if items.is_none() { + debug!("Unable to find {:?} onetime-key", account_type); + return Ok(None); + } + + let mut result = None; + + // "items" was checked to be None above, will be safe to unwrap here. + // Attempt to delete the onetime keys individually, a successful delete + // mints the onetime key to the requester + for item in items.unwrap() { + let pk = item.get_string(otk_table::PARTITION_KEY)?; + let otk = item.get_string(otk_table::SORT_KEY)?; + + let composite_key = HashMap::from([ + ( + otk_table::PARTITION_KEY.to_string(), + AttributeValue::S(pk.to_string()), + ), + ( + otk_table::SORT_KEY.to_string(), + AttributeValue::S(otk.to_string()), + ), + ]); + + debug!("Attempting to delete a {:?} onetime-key", account_type); + match self + .client + .delete_item() + .set_key(Some(composite_key)) + .table_name(otk_table::NAME) + .send() + .await + { + Ok(_) => { + result = Some(otk.to_string()); + break; + } + // This err should only happen if a delete occurred between the read + // above and this delete + Err(e) => { + debug!("Unable to delete key: {:?}", e); + continue; + } + } + } + + // Return deleted key + Ok(result) + } + + pub async fn get_onetime_keys( + &self, + device_id: &str, + account_type: OlmAccountType, + ) -> Result { + use crate::constants::one_time_keys_table::*; + + // Add related prefix to partition key to grab the correct result set + let partition_key = + create_one_time_key_partition_key(device_id, account_type); + + self + .client + .query() + .table_name(NAME) + .key_condition_expression(format!("{} = :pk", PARTITION_KEY)) + .expression_attribute_values(":pk", AttributeValue::S(partition_key)) + .send() + .await + .map_err(|e| Error::AwsSdk(e.into())) + } + pub async fn set_prekey( &self, user_id: String, diff --git a/services/identity/src/ddb_utils.rs b/services/identity/src/ddb_utils.rs --- a/services/identity/src/ddb_utils.rs +++ b/services/identity/src/ddb_utils.rs @@ -10,7 +10,7 @@ // Prefix the one time keys with the olm account variant. This allows for a single // DDB table to contain both notification and content keys for a device. -fn create_one_time_key_partition_key( +pub fn create_one_time_key_partition_key( device_id: &str, account_type: OlmAccountType, ) -> String { diff --git a/services/identity/src/error.rs b/services/identity/src/error.rs --- a/services/identity/src/error.rs +++ b/services/identity/src/error.rs @@ -1,4 +1,5 @@ use aws_sdk_dynamodb::{model::AttributeValue, Error as DynamoDBError}; +use std::collections::hash_map::HashMap; use std::fmt::{Display, Formatter, Result as FmtResult}; #[derive( @@ -53,3 +54,90 @@ #[display(...)] InvalidValue, } + +pub trait FromAttributeValue { + fn to_vec( + &self, + attr_name: &str, + ) -> Result<&Vec, DBItemError>; + fn to_string(&self, attr_name: &str) -> Result<&String, DBItemError>; + fn to_hashmap( + &self, + attr_name: &str, + ) -> Result<&HashMap, DBItemError>; +} + +fn handle_attr_failure(value: &AttributeValue, attr_name: &str) -> DBItemError { + DBItemError { + attribute_name: attr_name.to_string(), + attribute_value: Some(value.clone()), + attribute_error: DBItemAttributeError::IncorrectType, + } +} + +impl FromAttributeValue for AttributeValue { + fn to_vec( + &self, + attr_name: &str, + ) -> Result<&Vec, DBItemError> { + self.as_l().map_err(|e| handle_attr_failure(e, attr_name)) + } + + fn to_string(&self, attr_name: &str) -> Result<&String, DBItemError> { + self.as_s().map_err(|e| handle_attr_failure(e, attr_name)) + } + + fn to_hashmap( + &self, + attr_name: &str, + ) -> Result<&HashMap, DBItemError> { + self.as_m().map_err(|e| handle_attr_failure(e, attr_name)) + } +} + +pub trait AttributeValueFromHashMap { + fn get_string(&self, key: &str) -> Result<&String, DBItemError>; + fn get_map( + &self, + key: &str, + ) -> Result<&HashMap, DBItemError>; + fn get_vec(&self, key: &str) -> Result<&Vec, DBItemError>; +} + +impl AttributeValueFromHashMap for HashMap { + fn get_string(&self, key: &str) -> Result<&String, DBItemError> { + self + .get(key) + .ok_or(DBItemError { + attribute_name: key.to_string(), + attribute_value: None, + attribute_error: DBItemAttributeError::Missing, + })? + .to_string(key) + } + + fn get_map( + &self, + key: &str, + ) -> Result<&HashMap, DBItemError> { + self + .get(key) + .ok_or(DBItemError { + attribute_name: key.to_string(), + attribute_value: None, + attribute_error: DBItemAttributeError::Missing, + })? + .to_hashmap(key) + } + + fn get_vec(&self, key: &str) -> Result<&Vec, DBItemError> { + self + .get(key) + .ok_or(DBItemError { + attribute_name: key.to_string(), + attribute_value: None, + attribute_error: DBItemAttributeError::Missing, + })? + .to_vec(key) + } +} diff --git a/services/identity/src/grpc_services/authenticated.rs b/services/identity/src/grpc_services/authenticated.rs --- a/services/identity/src/grpc_services/authenticated.rs +++ b/services/identity/src/grpc_services/authenticated.rs @@ -12,9 +12,9 @@ } use auth_proto::{ identity_client_service_server::IdentityClientService, KeyserverKeysResponse, - OutboundKeysForUserRequest, RefreshUserPreKeysRequest, + OutboundKeyInfo, OutboundKeysForUserRequest, RefreshUserPreKeysRequest, }; -use client::Empty; +use client::{Empty, IdentityKeyInfo}; use tracing::debug; #[derive(derive_more::Constructor)] @@ -116,8 +116,37 @@ async fn get_keyserver_keys( &self, - _request: Request, + request: Request, ) -> Result, Status> { - unimplemented!(); + let message = request.into_inner(); + + let inner_response = self + .db_client + .get_keyserver_keys_for_user(&message.user_id) + .await + .map_err(handle_db_error)? + .map(|db_keys| OutboundKeyInfo { + identity_info: Some(IdentityKeyInfo { + payload: db_keys.key_payload, + payload_signature: db_keys.key_payload_signature, + social_proof: db_keys.social_proof, + }), + content_prekey: Some(client::PreKey { + pre_key: db_keys.content_prekey.prekey, + pre_key_signature: db_keys.content_prekey.prekey_signature, + }), + notif_prekey: Some(client::PreKey { + pre_key: db_keys.notif_prekey.prekey, + pre_key_signature: db_keys.notif_prekey.prekey_signature, + }), + onetime_content_prekey: db_keys.content_one_time_key, + onetime_notif_prekey: db_keys.notif_one_time_key, + }); + + let response = Response::new(KeyserverKeysResponse { + keyserver_info: inner_response, + }); + + return Ok(response); } }