Page MenuHomePhabricator
Differential D7540 Diff 25421 services/identity/src/database.rs

Changeset View

Standalone View

View Options

services/identity/src/database.rs

Show All 9 Lines
}; };
use aws_sdk_dynamodb::types::Blob; use aws_sdk_dynamodb::types::Blob;
use aws_sdk_dynamodb::{Client, Error as DynamoDBError}; use aws_sdk_dynamodb::{Client, Error as DynamoDBError};
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use opaque_ke::errors::ProtocolError; use opaque_ke::errors::ProtocolError;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tracing::{debug, error, info, warn}; use tracing::{debug, error, info, warn};
use crate::client_service::{LoginState, RegistrationState}; use crate::client_service::{FlattenedDeviceKeyUpload, RegistrationState};
use crate::config::CONFIG; use crate::config::CONFIG;
use crate::constants::{ use crate::constants::{
ACCESS_TOKEN_SORT_KEY, ACCESS_TOKEN_TABLE, ACCESS_TOKEN_SORT_KEY, ACCESS_TOKEN_TABLE,
ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE, ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE,
ACCESS_TOKEN_TABLE_PARTITION_KEY, ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE, ACCESS_TOKEN_TABLE_PARTITION_KEY, ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE,
ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, NONCE_TABLE, ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, NONCE_TABLE,
NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY, USERS_TABLE, NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY, USERS_TABLE,
USERS_TABLE_DEVICES_ATTRIBUTE, USERS_TABLE_DEVICES_ATTRIBUTE,
USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_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_ONETIME_KEYS_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME,
USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME,
USERS_TABLE_PARTITION_KEY, USERS_TABLE_REGISTRATION_ATTRIBUTE, USERS_TABLE_PARTITION_KEY, USERS_TABLE_REGISTRATION_ATTRIBUTE,
USERS_TABLE_USERNAME_ATTRIBUTE, USERS_TABLE_USERNAME_INDEX, USERS_TABLE_USERNAME_ATTRIBUTE, USERS_TABLE_USERNAME_INDEX,
USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE, USERS_TABLE_WALLET_ADDRESS_INDEX, USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE, USERS_TABLE_WALLET_ADDRESS_INDEX,
}; };
use crate::nonce::NonceData; use crate::nonce::NonceData;
use crate::token::{AccessTokenData, AuthType}; use crate::token::{AccessTokenData, AuthType};
use crate::utils::generate_uuid; use crate::utils::generate_uuid;
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Lines let client = match &CONFIG.localstack_endpoint {
None => Client::new(aws_config), None => Client::new(aws_config),
}; };
DatabaseClient { DatabaseClient {
client: Arc::new(client), client: Arc::new(client),
} }
} }
pub async fn add_user_to_users_table( pub async fn add_password_user_to_users_table(
&self, &self,
registration_state: RegistrationState, registration_state: RegistrationState,
password_file: Vec<u8>, password_file: Vec<u8>,
) -> Result<String, Error> { ) -> Result<String, Error> {
self
.add_user_to_users_table(
registration_state.flattened_device_key_upload,
Some((registration_state.username, Blob::new(password_file))),
None,
None,
)
.await
}
pub async fn add_wallet_user_to_users_table(
&self,
flattened_device_key_upload: FlattenedDeviceKeyUpload,
wallet_address: String,
social_proof: String,
) -> Result<String, Error> {
self
.add_user_to_users_table(
flattened_device_key_upload,
None,
Some(wallet_address),
Some(social_proof),
)
.await
}
async fn add_user_to_users_table(
&self,
flattened_device_key_upload: FlattenedDeviceKeyUpload,
username_and_password_file: Option<(String, Blob)>,
wallet_address: Option<String>,
social_proof: Option<String>,
) -> Result<String, Error> {
let user_id = generate_uuid(); let user_id = generate_uuid();
let device_info = HashMap::from([ let mut device_info = HashMap::from([
( (
USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(Device::Client.to_string()), AttributeValue::S(Device::Client.to_string()),
), ),
( (
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.key_payload),
registration_state.flattened_device_key_upload.key_payload,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.key_payload_signature),
registration_state
.flattened_device_key_upload
.key_payload_signature,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.identity_prekey),
registration_state
.flattened_device_key_upload
.identity_prekey,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(
registration_state flattened_device_key_upload.identity_prekey_signature,
.flattened_device_key_upload
.identity_prekey_signature,
), ),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::L( AttributeValue::L(
registration_state flattened_device_key_upload
.flattened_device_key_upload
.identity_onetime_keys .identity_onetime_keys
.into_iter() .into_iter()
.map(AttributeValue::S) .map(AttributeValue::S)
.collect(), .collect(),
), ),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.notif_prekey),
registration_state.flattened_device_key_upload.notif_prekey,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.notif_prekey_signature),
registration_state
.flattened_device_key_upload
.notif_prekey_signature,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(),
AttributeValue::L( AttributeValue::L(
registration_state flattened_device_key_upload
.flattened_device_key_upload
.notif_onetime_keys .notif_onetime_keys
.into_iter() .into_iter()
.map(AttributeValue::S) .map(AttributeValue::S)
.collect(), .collect(),
), ),
), ),
]); ]);
if let Some(social_proof) = social_proof {
device_info.insert(
USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(social_proof),
);
}
let devices = HashMap::from([( let devices = HashMap::from([(
registration_state.flattened_device_key_upload.device_id_key, flattened_device_key_upload.device_id_key,
AttributeValue::M(device_info), AttributeValue::M(device_info),
)]); )]);
let mut user = HashMap::from([
let user = HashMap::from([
( (
USERS_TABLE_PARTITION_KEY.to_string(), USERS_TABLE_PARTITION_KEY.to_string(),
AttributeValue::S(user_id.clone()), AttributeValue::S(user_id.clone()),
), ),
( (
USERS_TABLE_USERNAME_ATTRIBUTE.to_string(),
AttributeValue::S(registration_state.username),
),
(
USERS_TABLE_DEVICES_ATTRIBUTE.to_string(), USERS_TABLE_DEVICES_ATTRIBUTE.to_string(),
AttributeValue::M(devices), AttributeValue::M(devices),
), ),
(
USERS_TABLE_REGISTRATION_ATTRIBUTE.to_string(),
AttributeValue::B(Blob::new(password_file)),
),
]); ]);
if let Some((username, password_file)) = username_and_password_file {
user.insert(
USERS_TABLE_USERNAME_ATTRIBUTE.to_string(),
AttributeValue::S(username),
);
user.insert(
USERS_TABLE_REGISTRATION_ATTRIBUTE.to_string(),
AttributeValue::B(password_file),
);
}
if let Some(address) = wallet_address {
user.insert(
USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE.to_string(),
AttributeValue::S(address),
);
}
self self
.client .client
.put_item() .put_item()
.table_name(USERS_TABLE) .table_name(USERS_TABLE)
.set_item(Some(user)) .set_item(Some(user))
.send() .send()
.await .await
.map_err(|e| Error::AwsSdk(e.into()))?; .map_err(|e| Error::AwsSdk(e.into()))?;
Ok(user_id) Ok(user_id)
} }
pub async fn add_device_to_users_table( pub async fn add_password_user_device_to_users_table(
&self,
user_id: String,
flattened_device_key_upload: FlattenedDeviceKeyUpload,
) -> Result<(), Error> {
self
.add_device_to_users_table(user_id, flattened_device_key_upload, None)
.await
}
pub async fn add_wallet_user_device_to_users_table(
&self, &self,
login_state: LoginState, user_id: String,
flattened_device_key_upload: FlattenedDeviceKeyUpload,
social_proof: String,
) -> Result<(), Error> { ) -> Result<(), Error> {
let device_info = HashMap::from([ self
.add_device_to_users_table(
user_id,
flattened_device_key_upload,
Some(social_proof),
)
.await
}
async fn add_device_to_users_table(
&self,
user_id: String,
flattened_device_key_upload: FlattenedDeviceKeyUpload,
social_proof: Option<String>,
) -> Result<(), Error> {
let mut device_info = HashMap::from([
( (
USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(Device::Client.to_string()), AttributeValue::S(Device::Client.to_string()),
), ),
( (
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(login_state.flattened_device_key_upload.key_payload), AttributeValue::S(flattened_device_key_upload.key_payload),
), ),
( (
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.key_payload_signature),
login_state
.flattened_device_key_upload
.key_payload_signature,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.identity_prekey),
login_state.flattened_device_key_upload.identity_prekey,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(
login_state flattened_device_key_upload.identity_prekey_signature,
.flattened_device_key_upload
.identity_prekey_signature,
), ),
), ),
( (
USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::L( AttributeValue::L(
login_state flattened_device_key_upload
.flattened_device_key_upload
.identity_onetime_keys .identity_onetime_keys
.into_iter() .into_iter()
.map(AttributeValue::S) .map(AttributeValue::S)
.collect(), .collect(),
), ),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(login_state.flattened_device_key_upload.notif_prekey), AttributeValue::S(flattened_device_key_upload.notif_prekey),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(), .to_string(),
AttributeValue::S( AttributeValue::S(flattened_device_key_upload.notif_prekey_signature),
login_state
.flattened_device_key_upload
.notif_prekey_signature,
),
), ),
( (
USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(),
AttributeValue::L( AttributeValue::L(
login_state flattened_device_key_upload
.flattened_device_key_upload
.notif_onetime_keys .notif_onetime_keys
.into_iter() .into_iter()
.map(AttributeValue::S) .map(AttributeValue::S)
.collect(), .collect(),
), ),
), ),
]); ]);
if let Some(social_proof) = social_proof {
device_info.insert(
USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(social_proof),
);
}
let update_expression = let update_expression =
format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",); format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",);
let expression_attribute_names = HashMap::from([( let expression_attribute_names = HashMap::from([(
format!("#{}", "deviceID"), format!("#{}", "deviceID"),
login_state.flattened_device_key_upload.device_id_key, flattened_device_key_upload.device_id_key,
)]); )]);
let expression_attribute_values = let expression_attribute_values =
HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]); HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]);
self self
.client .client
.update_item() .update_item()
.table_name(USERS_TABLE) .table_name(USERS_TABLE)
.key( .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id))
USERS_TABLE_PARTITION_KEY,
AttributeValue::S(login_state.user_id),
)
.update_expression(update_expression) .update_expression(update_expression)
.set_expression_attribute_names(Some(expression_attribute_names)) .set_expression_attribute_names(Some(expression_attribute_names))
.set_expression_attribute_values(Some(expression_attribute_values)) .set_expression_attribute_values(Some(expression_attribute_values))
.send() .send()
.await .await
.map_err(|e| Error::AwsSdk(e.into()))?; .map_err(|e| Error::AwsSdk(e.into()))?;
Ok(()) Ok(())
▲ Show 20 Lines • Show All 583 Lines • Show Last 20 Lines