Page MenuHomePhabricator
Differential D7539 Diff 25420 services/identity/src/database.rs

Changeset View

Standalone View

View Options

services/identity/src/database.rs

use std::collections::HashMap; use std::collections::HashMap;
use std::fmt::{Display, Formatter, Result as FmtResult}; use std::fmt::{Display, Formatter, Result as FmtResult};
use std::str::FromStr; use std::str::FromStr;
use std::sync::Arc; use std::sync::Arc;
use aws_config::SdkConfig; use aws_config::SdkConfig;
use aws_sdk_dynamodb::model::AttributeValue; use aws_sdk_dynamodb::model::AttributeValue;
use aws_sdk_dynamodb::output::{ use aws_sdk_dynamodb::output::{
DeleteItemOutput, GetItemOutput, PutItemOutput, QueryOutput, DeleteItemOutput, GetItemOutput, PutItemOutput, QueryOutput,
}; };
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, ServerRegistration}; 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::RegistrationState; use crate::client_service::{LoginState, 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_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;
use comm_opaque::Cipher;
#[derive(Serialize, Deserialize)] #[derive(Serialize, Deserialize)]
pub struct OlmKeys { pub struct OlmKeys {
pub curve25519: String, pub curve25519: String,
pub ed25519: String, pub ed25519: String,
} }
#[derive(Serialize, Deserialize)] #[derive(Serialize, Deserialize)]
▲ Show 20 Lines • Show All 165 Lines • ▼ Show 20 Lines self
.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(
&self,
login_state: LoginState,
) -> Result<(), Error> {
let device_info = HashMap::from([
(
USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(Device::Client.to_string()),
),
(
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(login_state.flattened_device_key_upload.key_payload),
),
(
USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME
.to_string(),
AttributeValue::S(
login_state
.flattened_device_key_upload
.key_payload_signature,
),
),
(
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(
login_state.flattened_device_key_upload.identity_prekey,
),
),
(
USERS_TABLE_DEVICES_MAP_IDENTITY_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(),
AttributeValue::S(
login_state
.flattened_device_key_upload
.identity_prekey_signature,
),
),
(
USERS_TABLE_DEVICES_MAP_IDENTITY_ONETIME_KEYS_ATTRIBUTE_NAME
.to_string(),
AttributeValue::L(
login_state
.flattened_device_key_upload
.identity_onetime_keys
.into_iter()
.map(AttributeValue::S)
.collect(),
),
),
(
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(),
AttributeValue::S(login_state.flattened_device_key_upload.notif_prekey),
),
(
USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME
.to_string(),
AttributeValue::S(
login_state
.flattened_device_key_upload
.notif_prekey_signature,
),
),
(
USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(),
AttributeValue::L(
login_state
.flattened_device_key_upload
.notif_onetime_keys
.into_iter()
.map(AttributeValue::S)
.collect(),
),
),
]);
let update_expression =
format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",);
let expression_attribute_names = HashMap::from([(
format!("#{}", "deviceID"),
login_state.flattened_device_key_upload.device_id_key,
)]);
let expression_attribute_values =
HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]);
self
.client
.update_item()
.table_name(USERS_TABLE)
.key(
USERS_TABLE_PARTITION_KEY,
AttributeValue::S(login_state.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(())
}
pub async fn delete_user( pub async fn delete_user(
&self, &self,
user_id: String, user_id: String,
) -> Result<DeleteItemOutput, Error> { ) -> Result<DeleteItemOutput, Error> {
debug!("Attempting to delete user: {}", user_id); debug!("Attempting to delete user: {}", user_id);
match self match self
.client .client
▲ Show 20 Lines • Show All 206 Lines • ▼ Show 20 Lines match self
) )
.map(Some) .map(Some)
.map_err(Error::Attribute), .map_err(Error::Attribute),
Ok(_) => Ok(None), Ok(_) => Ok(None),
Err(e) => Err(e), Err(e) => Err(e),
} }
} }
pub async fn get_password_file_from_username( pub async fn get_user_id_and_password_file_from_username(
&self, &self,
username: &str, username: &str,
) -> Result<Option<Vec<u8>>, Error> { ) -> Result<Option<(String, Vec<u8>)>, Error> {
match self match self
.get_user_from_user_info(username.to_string(), AuthType::Password) .get_user_from_user_info(username.to_string(), AuthType::Password)
.await .await
{ {
Ok(Some(mut user)) => parse_registration_data_attribute( Ok(Some(mut user)) => {
let user_id = parse_string_attribute(
USERS_TABLE_PARTITION_KEY,
user.remove(USERS_TABLE_PARTITION_KEY),
)?;
let password_file = parse_registration_data_attribute(
user.remove(USERS_TABLE_REGISTRATION_ATTRIBUTE), user.remove(USERS_TABLE_REGISTRATION_ATTRIBUTE),
) )?;
.map(Some)
.map_err(Error::Attribute), Ok(Some((user_id, password_file)))
}
Ok(_) => { Ok(_) => {
info!( info!(
"No item found for user {} in PAKE registration table", "No item found for user {} in PAKE registration table",
username username
); );
Ok(None) Ok(None)
} }
Err(e) => { Err(e) => {
▲ Show 20 Lines • Show All 332 Lines • Show Last 20 Lines