Differential D8748 Diff 30122 services/identity/src/database.rs

Changeset View

Standalone View

services/identity/src/database.rs

use constant_time_eq::constant_time_eq;		use constant_time_eq::constant_time_eq;
use std::collections::{HashMap, HashSet};		use std::collections::{HashMap, HashSet};
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 crate::ddb_utils::{		use crate::ddb_utils::{
create_one_time_key_partition_key, into_one_time_put_requests, OlmAccountType,		create_one_time_key_partition_key, into_one_time_put_requests, OlmAccountType,
};		};
use crate::error::{DBItemAttributeError, DBItemError, Error};		use crate::error::{consume_error, DBItemAttributeError, DBItemError, Error};
use aws_config::SdkConfig;		use aws_config::SdkConfig;
use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest};		use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest};
use aws_sdk_dynamodb::output::{		use aws_sdk_dynamodb::output::{
DeleteItemOutput, GetItemOutput, PutItemOutput, QueryOutput,		DeleteItemOutput, GetItemOutput, PutItemOutput, QueryOutput,
};		};
use aws_sdk_dynamodb::{types::Blob, Client};		use aws_sdk_dynamodb::{types::Blob, Client};
use chrono::{DateTime, Utc};		use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};		use serde::{Deserialize, Serialize};
▲ Show 20 Lines • Show All 355 Lines • ▼ Show 20 Lines	impl DatabaseClient {
/// Will "mint" a single onetime key by attempting to successfully deleting		/// Will "mint" a single onetime key by attempting to successfully deleting
/// a key		/// a key
pub async fn get_onetime_key(		pub async fn get_onetime_key(
&self,		&self,
device_id: &str,		device_id: &str,
account_type: OlmAccountType,		account_type: OlmAccountType,
) -> Result<Option<String>, Error> {		) -> Result<Option<String>, Error> {
use crate::constants::one_time_keys_table as otk_table;		use crate::constants::one_time_keys_table as otk_table;
		use crate::constants::ONETIME_KEY_MINIMUM_THRESHOLD;

let query_result = self.get_onetime_keys(device_id, account_type).await?;		let query_result = self.get_onetime_keys(device_id, account_type).await?;
let items = query_result.items();		let items = query_result.items();

// If no onetime keys exists, return none early		// If no onetime keys exists, return none early
if items.is_none() {		let Some(item_vec) = items else {
debug!("Unable to find {:?} onetime-key", account_type);		debug!("Unable to find {:?} onetime-key", account_type);
return Ok(None);		return Ok(None);
		};

		if item_vec.len() < ONETIME_KEY_MINIMUM_THRESHOLD {
		// Avoid device_id being moved out-of-scope by "move"
		let device_id = device_id.to_string();
		tokio::spawn(async move {
		debug!("Attempting to request more keys for device: {}", &device_id);
		let result =
		crate::tunnelbroker::send_refresh_keys_request(&device_id).await;
		consume_error(result);
		});
}		}

let mut result = 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		// Attempt to delete the onetime keys individually, a successful delete
// mints the onetime key to the requester		// mints the onetime key to the requester
for item in items.unwrap() {		for item in item_vec {
let pk = item.get_string(otk_table::PARTITION_KEY)?;		let pk = item.get_string(otk_table::PARTITION_KEY)?;
let otk = item.get_string(otk_table::SORT_KEY)?;		let otk = item.get_string(otk_table::SORT_KEY)?;

let composite_key = HashMap::from([		let composite_key = HashMap::from([
(		(
otk_table::PARTITION_KEY.to_string(),		otk_table::PARTITION_KEY.to_string(),
AttributeValue::S(pk.to_string()),		AttributeValue::S(pk.to_string()),
),		),
▲ Show 20 Lines • Show All 1,052 Lines • Show Last 20 Lines