Skip to content

Commit

Permalink
RUST-1411 Move create_encrypted_collection helper to `ClientEncrypt…
Browse files Browse the repository at this point in the history 
…ion` (#858)
  • Loading branch information
@abr-egn
abr-egn authored Apr 19, 2023
1 parent d551dab commit 3c5ad68
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 84 deletions.
53 changes: 53 additions & 0 deletions src/client/csfle/client_encryption.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,14 @@ use crate::{
},
client::options::TlsOptions,
coll::options::CollectionOptions,
db::options::CreateCollectionOptions,
error::{Error, Result},
options::{ReadConcern, WriteConcern},
results::DeleteResult,
Client,
Collection,
Cursor,
Database,
Namespace,
};

Expand Down Expand Up @@ -410,6 +412,57 @@ impl ClientEncryption {
.ok_or_else(|| Error::internal("invalid decryption result"))?
.to_raw_bson())
}

/// Creates a new collection with encrypted fields, automatically creating new data encryption
/// keys when needed based on the configured [`CreateCollectionOptions::encrypted_fields`].
///
/// Returns the potentially updated `encrypted_fields` along with status, as keys may have been
/// created even when a failure occurs.
///
/// Does not affect any auto encryption settings on existing MongoClients that are already
/// configured with auto encryption.
pub async fn create_encrypted_collection(
&self,
db: &Database,
name: impl AsRef<str>,
master_key: MasterKey,
options: CreateCollectionOptions,
) -> (Document, Result<()>) {
let ef = match options.encrypted_fields.as_ref() {
Some(ef) => ef,
None => {
return (
doc! {},
Err(Error::invalid_argument(
"no encrypted_fields defined for collection",
)),
);
}
};
let mut ef_prime = ef.clone();
if let Ok(fields) = ef_prime.get_array_mut("fields") {
for f in fields {
let f_doc = if let Some(d) = f.as_document_mut() {
d
} else {
continue;
};
if f_doc.get("keyId") == Some(&Bson::Null) {
let d = match self.create_data_key(master_key.clone()).run().await {
Ok(v) => v,
Err(e) => return (ef_prime, Err(e)),
};
f_doc.insert("keyId", d);
}
}
}
let mut opts_prime = options.clone();
opts_prime.encrypted_fields = Some(ef_prime.clone());
(
ef_prime,
db.create_collection(name.as_ref(), opts_prime).await,
)
}
}

/// Options for creating a data key.
Expand Down
52 changes: 0 additions & 52 deletions src/db/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@ use std::{fmt::Debug, sync::Arc};
use bson::doc;
use futures_util::stream::TryStreamExt;

#[cfg(feature = "in-use-encryption-unstable")]
use crate::client_encryption::{ClientEncryption, MasterKey};
use crate::{
bson::{Bson, Document},
change_stream::{
Expand Down Expand Up @@ -423,56 +421,6 @@ impl Database {
self.create_collection_common(name, options, session).await
}

/// Creates a new collection with encrypted fields, automatically creating new data encryption
/// keys when needed based on the configured [`CreateCollectionOptions::encrypted_fields`].
///
/// Returns the potentially updated `encrypted_fields` along with status, as keys may have been
/// created even when a failure occurs.
///
/// Does not affect any auto encryption settings on existing MongoClients that are already
/// configured with auto encryption.
#[cfg(feature = "in-use-encryption-unstable")]
pub async fn create_encrypted_collection(
&self,
ce: &ClientEncryption,
name: impl AsRef<str>,
options: impl Into<Option<CreateCollectionOptions>>,
master_key: MasterKey,
) -> (Document, Result<()>) {
let options: Option<CreateCollectionOptions> = options.into();
let ef = match options.as_ref().and_then(|o| o.encrypted_fields.as_ref()) {
Some(ef) => ef,
None => {
return (
doc! {},
Err(Error::invalid_argument(
"no encrypted_fields defined for collection",
)),
);
}
};
let mut ef_prime = ef.clone();
if let Ok(fields) = ef_prime.get_array_mut("fields") {
for f in fields {
let f_doc = if let Some(d) = f.as_document_mut() {
d
} else {
continue;
};
if f_doc.get("keyId") == Some(&Bson::Null) {
let d = match ce.create_data_key(master_key.clone()).run().await {
Ok(v) => v,
Err(e) => return (ef_prime, Err(e)),
};
f_doc.insert("keyId", d);
}
}
}
let mut opts_prime = options.unwrap().clone(); // safe unwrap: no options would be caught by the encrypted_fields check
opts_prime.encrypted_fields = Some(ef_prime.clone());
(ef_prime, self.create_collection(name, opts_prime).await)
}

pub(crate) async fn run_command_common(
&self,
command: Document,
Expand Down
22 changes: 0 additions & 22 deletions src/sync/db.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -219,28 +219,6 @@ impl Database {
))
}

/// Creates a new collection with encrypted fields, automatically creating new data encryption
/// keys when needed based on the configured [`CreateCollectionOptions::encrypted_fields`].
///
/// Returns the potentially updated `encrypted_fields` along with status, as keys may have been
/// created even when a failure occurs.
///
/// Does not affect any auto encryption settings on existing MongoClients that are already
/// configured with auto encryption.
#[cfg(feature = "in-use-encryption-unstable")]
pub fn create_encrypted_collection(
&self,
ce: &crate::client_encryption::ClientEncryption,
name: impl AsRef<str>,
options: impl Into<Option<CreateCollectionOptions>>,
master_key: crate::client_encryption::MasterKey,
) -> (Document, Result<()>) {
runtime::block_on(
self.async_database
.create_encrypted_collection(ce, name, options, master_key),
)
}

/// Runs a database-level command.
///
/// Note that no inspection is done on `doc`, so the command will not use the database's default
Expand Down
25 changes: 15 additions & 10 deletions src/test/csfle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2928,7 +2928,7 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
)?;

// Case 1: Simple Creation and Validation
let opts = CreateCollectionOptions::builder()
let options = CreateCollectionOptions::builder()
.encrypted_fields(doc! {
"fields": [{
"path": "ssn",
Expand All @@ -2937,7 +2937,7 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
}],
})
.build();
db.create_encrypted_collection(&ce, "case_1", opts, master_key.clone())
ce.create_encrypted_collection(&db, "case_1", master_key.clone(), options)
.await
.1?;
let coll = db.collection::<Document>("case_1");
Expand All @@ -2949,8 +2949,13 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
);

// Case 2: Missing encryptedFields
let result = db
.create_encrypted_collection(&ce, "case_2", None, master_key.clone())
let result = ce
.create_encrypted_collection(
&db,
"case_2",
master_key.clone(),
CreateCollectionOptions::default(),
)
.await
.1;
assert!(
Expand All @@ -2960,7 +2965,7 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
);

// Case 3: Invalid keyId
let opts = CreateCollectionOptions::builder()
let options = CreateCollectionOptions::builder()
.encrypted_fields(doc! {
"fields": [{
"path": "ssn",
Expand All @@ -2969,8 +2974,8 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
}],
})
.build();
let result = db
.create_encrypted_collection(&ce, "case_1", opts, master_key.clone())
let result = ce
.create_encrypted_collection(&db, "case_1", master_key.clone(), options)
.await
.1;
assert!(
Expand All @@ -2980,7 +2985,7 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
);

// Case 4: Insert encrypted value
let opts = CreateCollectionOptions::builder()
let options = CreateCollectionOptions::builder()
.encrypted_fields(doc! {
"fields": [{
"path": "ssn",
Expand All @@ -2989,8 +2994,8 @@ async fn auto_encryption_keys(master_key: MasterKey) -> Result<()> {
}],
})
.build();
let (ef, result) = db
.create_encrypted_collection(&ce, "case_4", opts, master_key.clone())
let (ef, result) = ce
.create_encrypted_collection(&db, "case_4", master_key.clone(), options)
.await;
result?;
let key = match ef.get_array("fields")?[0]
Expand Down

0 comments on commit 3c5ad68

Please sign in to comment.