Page MenuHomePhabricator
Files F3366360

D8721.id30029.diff
No OneTemporary
Actions

D8721.id30029.diff
View Options

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<String>,
+ pub content_prekey: PreKey,
+ pub notif_prekey: PreKey,
+ pub content_one_time_key: Option<String>,
+ pub notif_one_time_key: Option<String>,
+}
+
#[derive(Clone)]
pub struct DatabaseClient {
client: Arc<Client>,
@@ -247,6 +266,197 @@
.await
}
+ pub async fn get_keyserver_keys_for_user(
+ &self,
+ user_id: &str,
+ ) -> Result<Option<OutboundKeys>, 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<String> = self
+ .get_onetime_key(keyserver_id, OlmAccountType::Notification)
+ .await?;
+ let content_one_time_key: Option<String> = 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<Option<String>, 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<QueryOutput, Error> {
+ 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<AttributeValue>, DBItemError>;
+ fn to_string(&self, attr_name: &str) -> Result<&String, DBItemError>;
+ fn to_hashmap(
+ &self,
+ attr_name: &str,
+ ) -> Result<&HashMap<String, AttributeValue>, 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<AttributeValue>, 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<String, AttributeValue>, 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<String, AttributeValue>, DBItemError>;
+ fn get_vec(&self, key: &str) -> Result<&Vec<AttributeValue>, DBItemError>;
+}
+
+impl AttributeValueFromHashMap for HashMap<String, AttributeValue> {
+ 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<String, AttributeValue>, 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<AttributeValue>, 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<OutboundKeysForUserRequest>,
+ request: Request<OutboundKeysForUserRequest>,
) -> Result<Response<KeyserverKeysResponse>, 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);
}
}

File Metadata

Mime Type
text/plain
Expires
Tue, Nov 26, 11:03 AM (20 h, 39 m)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
2581964
Default Alt Text
D8721.id30029.diff (16 KB)

Event Timeline