From 9d05d23958c18bcca38f78a2b651ea4d1ba57e76 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 25 Jul 2022 12:27:59 -0400 Subject: [PATCH 01/35] api outline --- src/client/csfle.rs | 1 + src/client/csfle/state_machine.rs | 11 +++++++++++ 2 files changed, 12 insertions(+) create mode 100644 src/client/csfle/state_machine.rs diff --git a/src/client/csfle.rs b/src/client/csfle.rs index 61cfe4641..f4f569629 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -1,4 +1,5 @@ pub mod options; +mod state_machine; use std::{ path::Path, diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs new file mode 100644 index 000000000..5826c6f08 --- /dev/null +++ b/src/client/csfle/state_machine.rs @@ -0,0 +1,11 @@ +use bson::RawDocumentBuf; +use mongocrypt::ctx::Ctx; + +use crate::Client; +use crate::error::Result; + +impl Client { + async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx) -> Result { + todo!() + } +} \ No newline at end of file From 5bd3324a09bf5a00e304f002d610508c3bb0e592 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 25 Jul 2022 13:01:17 -0400 Subject: [PATCH 02/35] wip --- src/client/csfle/state_machine.rs | 15 +++++-- src/operation/mod.rs | 1 + src/operation/raw_output.rs | 72 +++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 3 deletions(-) create mode 100644 src/operation/raw_output.rs diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 5826c6f08..b6f16d501 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,11 +1,20 @@ use bson::RawDocumentBuf; -use mongocrypt::ctx::Ctx; +use mongocrypt::ctx::{Ctx, State}; use crate::Client; use crate::error::Result; impl Client { - async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx) -> Result { - todo!() + pub(crate) async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx) -> Result { + let mut result = RawDocumentBuf::new(); + loop { + match ctx.state()? { + State::NeedMongoCollinfo => { + let filter = ctx.mongo_op()?; + } + _ => todo!(), + } + } + Ok(result) } } \ No newline at end of file diff --git a/src/operation/mod.rs b/src/operation/mod.rs index e59ce4bad..ff4297359 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -17,6 +17,7 @@ mod insert; mod list_collections; mod list_databases; mod list_indexes; +mod raw_output; mod run_command; mod update; diff --git a/src/operation/raw_output.rs b/src/operation/raw_output.rs new file mode 100644 index 000000000..b54a7f9d5 --- /dev/null +++ b/src/operation/raw_output.rs @@ -0,0 +1,72 @@ +use crate::cmap::{StreamDescription, RawCommandResponse, Command}; +use crate::error::Result; + +use super::Operation; + +pub(crate) struct RawOutput(Op); + +impl Operation for RawOutput { + type O = RawCommandResponse; + type Command = Op::Command; + const NAME: &'static str = Op::NAME; + + fn build(&mut self, description: &StreamDescription) -> Result> { + todo!() + } + + fn handle_response( + &self, + response: RawCommandResponse, + description: &StreamDescription, + ) -> Result { + todo!() + } + + fn serialize_command(&mut self, cmd: Command) -> Result> { + Ok(bson::to_vec(&cmd)?) + } + + fn extract_at_cluster_time(&self, _response: &bson::RawDocument) -> Result> { + Ok(None) + } + + fn handle_error(&self, error: crate::error::Error) -> Result { + Err(error) + } + + fn selection_criteria(&self) -> Option<&crate::selection_criteria::SelectionCriteria> { + None + } + + fn is_acknowledged(&self) -> bool { + self.write_concern() + .map(crate::options::WriteConcern::is_acknowledged) + .unwrap_or(true) + } + + fn write_concern(&self) -> Option<&crate::options::WriteConcern> { + None + } + + fn supports_read_concern(&self, _description: &StreamDescription) -> bool { + false + } + + fn supports_sessions(&self) -> bool { + true + } + + fn retryability(&self) -> super::Retryability { + super::Retryability::None + } + + fn update_for_retry(&mut self) {} + + fn pinned_connection(&self) -> Option<&crate::cmap::conn::PinnedConnectionHandle> { + None + } + + fn name(&self) -> &str { + Self::NAME + } +} \ No newline at end of file From 56c41600ecb9279d7827e0a1e7233f409dc4974a Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 25 Jul 2022 13:21:24 -0400 Subject: [PATCH 03/35] shift operation defaults to mirror trait --- src/operation/abort_transaction/mod.rs | 6 +- src/operation/aggregate/change_stream.rs | 4 +- src/operation/aggregate/mod.rs | 6 +- src/operation/commit_transaction/mod.rs | 4 +- src/operation/count/mod.rs | 4 +- src/operation/count_documents/mod.rs | 4 +- src/operation/create/mod.rs | 4 +- src/operation/create_indexes/mod.rs | 4 +- src/operation/delete/mod.rs | 4 +- src/operation/distinct/mod.rs | 4 +- src/operation/drop_collection/mod.rs | 4 +- src/operation/drop_database/mod.rs | 4 +- src/operation/drop_indexes/mod.rs | 4 +- src/operation/find/mod.rs | 4 +- src/operation/find_and_modify/mod.rs | 4 +- src/operation/get_more/mod.rs | 4 +- src/operation/insert/mod.rs | 4 +- src/operation/list_collections/mod.rs | 4 +- src/operation/list_databases/mod.rs | 4 +- src/operation/list_indexes/mod.rs | 4 +- src/operation/mod.rs | 186 ++++++++++++++++++----- src/operation/run_command/mod.rs | 4 +- src/operation/update/mod.rs | 4 +- 23 files changed, 196 insertions(+), 82 deletions(-) diff --git a/src/operation/abort_transaction/mod.rs b/src/operation/abort_transaction/mod.rs index 2e4f9b5ec..35a992993 100644 --- a/src/operation/abort_transaction/mod.rs +++ b/src/operation/abort_transaction/mod.rs @@ -5,12 +5,12 @@ use crate::{ client::session::TransactionPin, cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{Operation, Retryability}, + operation::{Retryability}, options::WriteConcern, selection_criteria::SelectionCriteria, }; -use super::WriteConcernOnlyBody; +use super::{WriteConcernOnlyBody, OperationWithDefaults}; pub(crate) struct AbortTransaction { write_concern: Option, @@ -26,7 +26,7 @@ impl AbortTransaction { } } -impl Operation for AbortTransaction { +impl OperationWithDefaults for AbortTransaction { type O = (); type Command = Document; diff --git a/src/operation/aggregate/change_stream.rs b/src/operation/aggregate/change_stream.rs index e6c1bb5ac..50a108cc3 100644 --- a/src/operation/aggregate/change_stream.rs +++ b/src/operation/aggregate/change_stream.rs @@ -4,7 +4,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, cursor::CursorSpecification, error::Result, - operation::{append_options, Operation, Retryability}, + operation::{append_options, OperationWithDefaults, Retryability}, options::{ChangeStreamOptions, SelectionCriteria, WriteConcern}, }; @@ -39,7 +39,7 @@ impl ChangeStreamAggregate { } } -impl Operation for ChangeStreamAggregate { +impl OperationWithDefaults for ChangeStreamAggregate { type O = (CursorSpecification, ChangeStreamData); type Command = Document; diff --git a/src/operation/aggregate/mod.rs b/src/operation/aggregate/mod.rs index b707bfc10..b2ba5c8ab 100644 --- a/src/operation/aggregate/mod.rs +++ b/src/operation/aggregate/mod.rs @@ -9,12 +9,12 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, cursor::CursorSpecification, error::Result, - operation::{append_options, remove_empty_write_concern, Operation, Retryability}, + operation::{append_options, remove_empty_write_concern, Retryability}, options::{AggregateOptions, SelectionCriteria, WriteConcern}, Namespace, }; -use super::{CursorBody, WriteConcernOnlyBody, SERVER_4_2_0_WIRE_VERSION}; +use super::{CursorBody, WriteConcernOnlyBody, SERVER_4_2_0_WIRE_VERSION, OperationWithDefaults}; pub(crate) use change_stream::ChangeStreamAggregate; @@ -46,7 +46,7 @@ impl Aggregate { // IMPORTANT: If new method implementations are added here, make sure `ChangeStreamAggregate` has // the equivalent delegations. -impl Operation for Aggregate { +impl OperationWithDefaults for Aggregate { type O = CursorSpecification; type Command = Document; diff --git a/src/operation/commit_transaction/mod.rs b/src/operation/commit_transaction/mod.rs index 1a3e4b532..3c3d455d6 100644 --- a/src/operation/commit_transaction/mod.rs +++ b/src/operation/commit_transaction/mod.rs @@ -5,7 +5,7 @@ use bson::{doc, Document}; use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, Operation, Retryability}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults, Retryability}, options::{Acknowledgment, TransactionOptions, WriteConcern}, }; @@ -21,7 +21,7 @@ impl CommitTransaction { } } -impl Operation for CommitTransaction { +impl OperationWithDefaults for CommitTransaction { type O = (); type Command = Document; diff --git a/src/operation/count/mod.rs b/src/operation/count/mod.rs index 50d4dc7e5..15c0f5fe5 100644 --- a/src/operation/count/mod.rs +++ b/src/operation/count/mod.rs @@ -9,7 +9,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, coll::{options::EstimatedDocumentCountOptions, Namespace}, error::{Error, Result}, - operation::{append_options, Operation, Retryability}, + operation::{append_options, OperationWithDefaults, Retryability}, selection_criteria::SelectionCriteria, }; @@ -35,7 +35,7 @@ impl Count { } } -impl Operation for Count { +impl OperationWithDefaults for Count { type O = u64; type Command = Document; diff --git a/src/operation/count_documents/mod.rs b/src/operation/count_documents/mod.rs index 3853b944c..413059326 100644 --- a/src/operation/count_documents/mod.rs +++ b/src/operation/count_documents/mod.rs @@ -5,7 +5,7 @@ use std::convert::TryInto; use serde::Deserialize; -use super::{Operation, Retryability, SingleCursorResult}; +use super::{OperationWithDefaults, Retryability, SingleCursorResult}; use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, error::{Error, ErrorKind, Result}, @@ -75,7 +75,7 @@ impl CountDocuments { } } -impl Operation for CountDocuments { +impl OperationWithDefaults for CountDocuments { type O = u64; type Command = Document; diff --git a/src/operation/create/mod.rs b/src/operation/create/mod.rs index bda286099..d70cbe2c8 100644 --- a/src/operation/create/mod.rs +++ b/src/operation/create/mod.rs @@ -7,7 +7,7 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, Operation, WriteConcernOnlyBody}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, options::{CreateCollectionOptions, WriteConcern}, Namespace, }; @@ -35,7 +35,7 @@ impl Create { } } -impl Operation for Create { +impl OperationWithDefaults for Create { type O = (); type Command = Document; diff --git a/src/operation/create_indexes/mod.rs b/src/operation/create_indexes/mod.rs index a0815aff2..247e6cb40 100644 --- a/src/operation/create_indexes/mod.rs +++ b/src/operation/create_indexes/mod.rs @@ -6,7 +6,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, error::{ErrorKind, Result}, index::IndexModel, - operation::{append_options, remove_empty_write_concern, Operation}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults}, options::{CreateIndexOptions, WriteConcern}, results::CreateIndexesResult, Namespace, @@ -47,7 +47,7 @@ impl CreateIndexes { } } -impl Operation for CreateIndexes { +impl OperationWithDefaults for CreateIndexes { type O = CreateIndexesResult; type Command = Document; const NAME: &'static str = "createIndexes"; diff --git a/src/operation/delete/mod.rs b/src/operation/delete/mod.rs index 7af98f73e..7458011d4 100644 --- a/src/operation/delete/mod.rs +++ b/src/operation/delete/mod.rs @@ -10,7 +10,7 @@ use crate::{ operation::{ append_options, remove_empty_write_concern, - Operation, + OperationWithDefaults, Retryability, WriteResponseBody, }, @@ -59,7 +59,7 @@ impl Delete { } } -impl Operation for Delete { +impl OperationWithDefaults for Delete { type O = DeleteResult; type Command = Document; diff --git a/src/operation/distinct/mod.rs b/src/operation/distinct/mod.rs index a48d3594b..e5f18cdca 100644 --- a/src/operation/distinct/mod.rs +++ b/src/operation/distinct/mod.rs @@ -9,7 +9,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, coll::{options::DistinctOptions, Namespace}, error::Result, - operation::{append_options, Operation, Retryability}, + operation::{append_options, OperationWithDefaults, Retryability}, selection_criteria::SelectionCriteria, }; @@ -49,7 +49,7 @@ impl Distinct { } } -impl Operation for Distinct { +impl OperationWithDefaults for Distinct { type O = Vec; type Command = Document; diff --git a/src/operation/drop_collection/mod.rs b/src/operation/drop_collection/mod.rs index 6d79b985b..6f9f29aff 100644 --- a/src/operation/drop_collection/mod.rs +++ b/src/operation/drop_collection/mod.rs @@ -7,7 +7,7 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::{Error, Result}, - operation::{append_options, remove_empty_write_concern, Operation, WriteConcernOnlyBody}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, options::{DropCollectionOptions, WriteConcern}, Namespace, }; @@ -35,7 +35,7 @@ impl DropCollection { } } -impl Operation for DropCollection { +impl OperationWithDefaults for DropCollection { type O = (); type Command = Document; diff --git a/src/operation/drop_database/mod.rs b/src/operation/drop_database/mod.rs index d4cfafbcc..53ba013d6 100644 --- a/src/operation/drop_database/mod.rs +++ b/src/operation/drop_database/mod.rs @@ -7,7 +7,7 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, Operation, WriteConcernOnlyBody}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, options::{DropDatabaseOptions, WriteConcern}, }; @@ -28,7 +28,7 @@ impl DropDatabase { } } -impl Operation for DropDatabase { +impl OperationWithDefaults for DropDatabase { type O = (); type Command = Document; diff --git a/src/operation/drop_indexes/mod.rs b/src/operation/drop_indexes/mod.rs index 70c236607..dd7c9fbd5 100644 --- a/src/operation/drop_indexes/mod.rs +++ b/src/operation/drop_indexes/mod.rs @@ -5,7 +5,7 @@ use crate::{ bson::{doc, Document}, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, Operation}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults}, options::{DropIndexOptions, WriteConcern}, Namespace, }; @@ -34,7 +34,7 @@ impl DropIndexes { } } -impl Operation for DropIndexes { +impl OperationWithDefaults for DropIndexes { type O = (); type Command = Document; const NAME: &'static str = "dropIndexes"; diff --git a/src/operation/find/mod.rs b/src/operation/find/mod.rs index ecb124c2a..f4f50141f 100644 --- a/src/operation/find/mod.rs +++ b/src/operation/find/mod.rs @@ -6,7 +6,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, cursor::CursorSpecification, error::{ErrorKind, Result}, - operation::{append_options, CursorBody, Operation, Retryability}, + operation::{append_options, CursorBody, OperationWithDefaults, Retryability}, options::{CursorType, FindOptions, SelectionCriteria}, Namespace, }; @@ -44,7 +44,7 @@ impl Find { } } -impl Operation for Find { +impl OperationWithDefaults for Find { type O = CursorSpecification; type Command = Document; const NAME: &'static str = "find"; diff --git a/src/operation/find_and_modify/mod.rs b/src/operation/find_and_modify/mod.rs index 1e9e3101f..bdc4815e0 100644 --- a/src/operation/find_and_modify/mod.rs +++ b/src/operation/find_and_modify/mod.rs @@ -21,7 +21,7 @@ use crate::{ Namespace, }, error::{ErrorKind, Result}, - operation::{append_options, remove_empty_write_concern, Operation, Retryability}, + operation::{append_options, remove_empty_write_concern, OperationWithDefaults, Retryability}, options::WriteConcern, }; @@ -95,7 +95,7 @@ where } } -impl Operation for FindAndModify +impl OperationWithDefaults for FindAndModify where T: DeserializeOwned, { diff --git a/src/operation/get_more/mod.rs b/src/operation/get_more/mod.rs index 4a9e64a27..c12f5abe7 100644 --- a/src/operation/get_more/mod.rs +++ b/src/operation/get_more/mod.rs @@ -12,7 +12,7 @@ use crate::{ cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription}, cursor::CursorInformation, error::{ErrorKind, Result}, - operation::Operation, + operation::OperationWithDefaults, options::SelectionCriteria, results::GetMoreResult, Namespace, @@ -44,7 +44,7 @@ impl<'conn> GetMore<'conn> { } } -impl<'conn> Operation for GetMore<'conn> { +impl<'conn> OperationWithDefaults for GetMore<'conn> { type O = GetMoreResult; type Command = Document; diff --git a/src/operation/insert/mod.rs b/src/operation/insert/mod.rs index 0cf129b28..031dfff41 100644 --- a/src/operation/insert/mod.rs +++ b/src/operation/insert/mod.rs @@ -11,7 +11,7 @@ use crate::{ bson_util, cmap::{Command, RawCommandResponse, StreamDescription}, error::{BulkWriteFailure, Error, ErrorKind, Result}, - operation::{remove_empty_write_concern, Operation, Retryability, WriteResponseBody}, + operation::{remove_empty_write_concern, OperationWithDefaults, Retryability, WriteResponseBody}, options::{InsertManyOptions, WriteConcern}, results::InsertManyResult, Namespace, @@ -49,7 +49,7 @@ impl<'a, T> Insert<'a, T> { } } -impl<'a, T: Serialize> Operation for Insert<'a, T> { +impl<'a, T: Serialize> OperationWithDefaults for Insert<'a, T> { type O = InsertManyResult; type Command = InsertCommand; diff --git a/src/operation/list_collections/mod.rs b/src/operation/list_collections/mod.rs index c49c78a68..951c165f8 100644 --- a/src/operation/list_collections/mod.rs +++ b/src/operation/list_collections/mod.rs @@ -6,7 +6,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, cursor::CursorSpecification, error::Result, - operation::{append_options, CursorBody, Operation, Retryability}, + operation::{append_options, CursorBody, OperationWithDefaults, Retryability}, options::{ListCollectionsOptions, ReadPreference, SelectionCriteria}, }; @@ -39,7 +39,7 @@ impl ListCollections { } } -impl Operation for ListCollections { +impl OperationWithDefaults for ListCollections { type O = CursorSpecification; type Command = Document; diff --git a/src/operation/list_databases/mod.rs b/src/operation/list_databases/mod.rs index c202ef4f0..8360fd78f 100644 --- a/src/operation/list_databases/mod.rs +++ b/src/operation/list_databases/mod.rs @@ -8,7 +8,7 @@ use crate::{ bson::{doc, Document}, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, Operation, Retryability}, + operation::{append_options, OperationWithDefaults, Retryability}, options::ListDatabasesOptions, selection_criteria::{ReadPreference, SelectionCriteria}, }; @@ -43,7 +43,7 @@ impl ListDatabases { } } -impl Operation for ListDatabases { +impl OperationWithDefaults for ListDatabases { type O = Vec; type Command = Document; diff --git a/src/operation/list_indexes/mod.rs b/src/operation/list_indexes/mod.rs index 9189b62af..ff9130017 100644 --- a/src/operation/list_indexes/mod.rs +++ b/src/operation/list_indexes/mod.rs @@ -3,7 +3,7 @@ use crate::{ cmap::{Command, RawCommandResponse, StreamDescription}, cursor::CursorSpecification, error::Result, - operation::{append_options, Operation}, + operation::{append_options, OperationWithDefaults}, options::ListIndexesOptions, selection_criteria::{ReadPreference, SelectionCriteria}, Namespace, @@ -36,7 +36,7 @@ impl ListIndexes { } } -impl Operation for ListIndexes { +impl OperationWithDefaults for ListIndexes { type O = CursorSpecification; type Command = Document; diff --git a/src/operation/mod.rs b/src/operation/mod.rs index ff4297359..62a03a2ee 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -74,6 +74,8 @@ pub(crate) use update::Update; const SERVER_4_2_0_WIRE_VERSION: i32 = 8; /// A trait modeling the behavior of a server side operation. +/// +/// No methods in this trait should have default behaviors to ensure that wrapper operations replicate all behavior. Default behavior is provided by the `OperationDefault` trait. pub(crate) trait Operation { /// The output type of this operation. type O; @@ -90,15 +92,11 @@ pub(crate) trait Operation { /// Perform custom serialization of the built command. /// By default, this will just call through to the `Serialize` implementation of the command. - fn serialize_command(&mut self, cmd: Command) -> Result> { - Ok(bson::to_vec(&cmd)?) - } + fn serialize_command(&mut self, cmd: Command) -> Result>; /// Parse the response for the atClusterTime field. /// Depending on the operation, this may be found in different locations. - fn extract_at_cluster_time(&self, _response: &RawDocument) -> Result> { - Ok(None) - } + fn extract_at_cluster_time(&self, _response: &RawDocument) -> Result>; /// Interprets the server response to the command. fn handle_response( @@ -109,52 +107,32 @@ pub(crate) trait Operation { /// Interpret an error encountered while sending the built command to the server, potentially /// recovering. - fn handle_error(&self, error: Error) -> Result { - Err(error) - } + fn handle_error(&self, error: Error) -> Result; /// Criteria to use for selecting the server that this operation will be executed on. - fn selection_criteria(&self) -> Option<&SelectionCriteria> { - None - } + fn selection_criteria(&self) -> Option<&SelectionCriteria>; /// Whether or not this operation will request acknowledgment from the server. - fn is_acknowledged(&self) -> bool { - self.write_concern() - .map(WriteConcern::is_acknowledged) - .unwrap_or(true) - } + fn is_acknowledged(&self) -> bool; /// The write concern to use for this operation, if any. - fn write_concern(&self) -> Option<&WriteConcern> { - None - } + fn write_concern(&self) -> Option<&WriteConcern>; /// Returns whether or not this command supports the `readConcern` field. - fn supports_read_concern(&self, _description: &StreamDescription) -> bool { - false - } + fn supports_read_concern(&self, _description: &StreamDescription) -> bool; /// Whether this operation supports sessions or not. - fn supports_sessions(&self) -> bool { - true - } + fn supports_sessions(&self) -> bool; /// The level of retryability the operation supports. - fn retryability(&self) -> Retryability { - Retryability::None - } + fn retryability(&self) -> Retryability; /// Updates this operation as needed for a retry. - fn update_for_retry(&mut self) {} + fn update_for_retry(&mut self); - fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> { - None - } + fn pinned_connection(&self) -> Option<&PinnedConnectionHandle>; - fn name(&self) -> &str { - Self::NAME - } + fn name(&self) -> &str; } pub(crate) trait CommandBody: Serialize { @@ -399,3 +377,139 @@ macro_rules! remove_empty_write_concern { } pub(crate) use remove_empty_write_concern; + +// A mirror of the `Operation` trait, with default behavior where appropriate. Should only be implemented by leaf operation types. +pub(crate) trait OperationWithDefaults { + /// The output type of this operation. + type O; + + /// The format of the command body constructed in `build`. + type Command: CommandBody; + + /// The name of the server side command associated with this operation. + const NAME: &'static str; + + /// Returns the command that should be sent to the server as part of this operation. + /// The operation may store some additional state that is required for handling the response. + fn build(&mut self, description: &StreamDescription) -> Result>; + + /// Perform custom serialization of the built command. + /// By default, this will just call through to the `Serialize` implementation of the command. + fn serialize_command(&mut self, cmd: Command) -> Result> { + Ok(bson::to_vec(&cmd)?) + } + + /// Parse the response for the atClusterTime field. + /// Depending on the operation, this may be found in different locations. + fn extract_at_cluster_time(&self, _response: &RawDocument) -> Result> { + Ok(None) + } + + /// Interprets the server response to the command. + fn handle_response( + &self, + response: RawCommandResponse, + description: &StreamDescription, + ) -> Result; + + /// Interpret an error encountered while sending the built command to the server, potentially + /// recovering. + fn handle_error(&self, error: Error) -> Result { + Err(error) + } + + /// Criteria to use for selecting the server that this operation will be executed on. + fn selection_criteria(&self) -> Option<&SelectionCriteria> { + None + } + + /// Whether or not this operation will request acknowledgment from the server. + fn is_acknowledged(&self) -> bool { + self.write_concern() + .map(WriteConcern::is_acknowledged) + .unwrap_or(true) + } + + /// The write concern to use for this operation, if any. + fn write_concern(&self) -> Option<&WriteConcern> { + None + } + + /// Returns whether or not this command supports the `readConcern` field. + fn supports_read_concern(&self, _description: &StreamDescription) -> bool { + false + } + + /// Whether this operation supports sessions or not. + fn supports_sessions(&self) -> bool { + true + } + + /// The level of retryability the operation supports. + fn retryability(&self) -> Retryability { + Retryability::None + } + + /// Updates this operation as needed for a retry. + fn update_for_retry(&mut self) {} + + fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> { + None + } + + fn name(&self) -> &str { + Self::NAME + } +} + +impl Operation for T { + type O = T::O; + type Command = T::Command; + const NAME: &'static str = T::NAME; + fn build(&mut self, description: &StreamDescription) -> Result> { + self.build(description) + } + fn serialize_command(&mut self, cmd: Command) -> Result> { + self.serialize_command(cmd) + } + fn extract_at_cluster_time(&self, response: &RawDocument) -> Result> { + self.extract_at_cluster_time(response) + } + fn handle_response( + &self, + response: RawCommandResponse, + description: &StreamDescription, + ) -> Result { + self.handle_response(response, description) + } + fn handle_error(&self, error: Error) -> Result { + self.handle_error(error) + } + fn selection_criteria(&self) -> Option<&SelectionCriteria> { + self.selection_criteria() + } + fn is_acknowledged(&self) -> bool { + self.is_acknowledged() + } + fn write_concern(&self) -> Option<&WriteConcern> { + self.write_concern() + } + fn supports_read_concern(&self, description: &StreamDescription) -> bool { + self.supports_read_concern(description) + } + fn supports_sessions(&self) -> bool { + self.supports_sessions() + } + fn retryability(&self) -> Retryability { + self.retryability() + } + fn update_for_retry(&mut self) { + self.update_for_retry() + } + fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> { + self.pinned_connection() + } + fn name(&self) -> &str { + self.name() + } +} \ No newline at end of file diff --git a/src/operation/run_command/mod.rs b/src/operation/run_command/mod.rs index ff862e1a0..8460a643f 100644 --- a/src/operation/run_command/mod.rs +++ b/src/operation/run_command/mod.rs @@ -5,7 +5,7 @@ use std::convert::TryInto; use bson::RawBsonRef; -use super::{CursorBody, Operation}; +use super::{CursorBody, Operation, OperationWithDefaults}; use crate::{ bson::Document, client::SESSIONS_UNSUPPORTED_COMMANDS, @@ -50,7 +50,7 @@ impl<'conn> RunCommand<'conn> { } } -impl<'conn> Operation for RunCommand<'conn> { +impl<'conn> OperationWithDefaults for RunCommand<'conn> { type O = Document; type Command = Document; diff --git a/src/operation/update/mod.rs b/src/operation/update/mod.rs index a1b4fcda9..782a5c5d0 100644 --- a/src/operation/update/mod.rs +++ b/src/operation/update/mod.rs @@ -8,7 +8,7 @@ use crate::{ bson_util, cmap::{Command, RawCommandResponse, StreamDescription}, error::{convert_bulk_errors, Result}, - operation::{Operation, Retryability, WriteResponseBody}, + operation::{OperationWithDefaults, Retryability, WriteResponseBody}, options::{UpdateModifications, UpdateOptions, WriteConcern}, results::UpdateResult, Namespace, @@ -55,7 +55,7 @@ impl Update { } } -impl Operation for Update { +impl OperationWithDefaults for Update { type O = UpdateResult; type Command = Document; From 1349cffeb4af7cf29b304b07233a7c773b91f9c4 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 25 Jul 2022 14:53:30 -0400 Subject: [PATCH 04/35] needmongocollinfo --- src/client/csfle/state_machine.rs | 22 ++++++++++---- src/operation/mod.rs | 1 + src/operation/raw_output.rs | 48 +++++++++++++++---------------- src/operation/run_command/mod.rs | 2 +- 4 files changed, 43 insertions(+), 30 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index b6f16d501..c07c847b9 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,16 +1,28 @@ -use bson::RawDocumentBuf; +use std::convert::TryInto; + +use bson::{RawDocumentBuf, Document, RawDocument}; use mongocrypt::ctx::{Ctx, State}; -use crate::Client; -use crate::error::Result; +use crate::{Client, Database}; +use crate::error::{Error, Result}; +use crate::operation::{ListCollections, RawOutput}; + +fn raw_to_doc(raw: &RawDocument) -> Result { + raw.try_into().map_err(|e| Error::internal("???")) +} impl Client { - pub(crate) async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx) -> Result { + pub(crate) async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx, db: Option) -> Result { let mut result = RawDocumentBuf::new(); loop { match ctx.state()? { State::NeedMongoCollinfo => { - let filter = ctx.mongo_op()?; + let filter: Document = raw_to_doc(ctx.mongo_op()?)?; + let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?; + let mut cursor = db.list_collections(filter, None).await?; + if cursor.advance().await? { + ctx.mongo_feed(cursor.current())?; + } } _ => todo!(), } diff --git a/src/operation/mod.rs b/src/operation/mod.rs index 62a03a2ee..ab8b94e3f 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -68,6 +68,7 @@ pub(crate) use insert::Insert; pub(crate) use list_collections::ListCollections; pub(crate) use list_databases::ListDatabases; pub(crate) use list_indexes::ListIndexes; +pub(crate) use raw_output::RawOutput; pub(crate) use run_command::RunCommand; pub(crate) use update::Update; diff --git a/src/operation/raw_output.rs b/src/operation/raw_output.rs index b54a7f9d5..0bc311e8b 100644 --- a/src/operation/raw_output.rs +++ b/src/operation/raw_output.rs @@ -3,7 +3,7 @@ use crate::error::Result; use super::Operation; -pub(crate) struct RawOutput(Op); +pub(crate) struct RawOutput(pub(crate) Op); impl Operation for RawOutput { type O = RawCommandResponse; @@ -11,23 +11,23 @@ impl Operation for RawOutput { const NAME: &'static str = Op::NAME; fn build(&mut self, description: &StreamDescription) -> Result> { - todo!() + self.0.build(description) } - fn handle_response( - &self, - response: RawCommandResponse, - description: &StreamDescription, - ) -> Result { - todo!() + fn serialize_command(&mut self, cmd: Command) -> Result> { + self.0.serialize_command(cmd) } - fn serialize_command(&mut self, cmd: Command) -> Result> { - Ok(bson::to_vec(&cmd)?) + fn extract_at_cluster_time(&self, response: &bson::RawDocument) -> Result> { + self.0.extract_at_cluster_time(response) } - fn extract_at_cluster_time(&self, _response: &bson::RawDocument) -> Result> { - Ok(None) + fn handle_response( + &self, + response: RawCommandResponse, + _description: &StreamDescription, + ) -> Result { + Ok(response) } fn handle_error(&self, error: crate::error::Error) -> Result { @@ -35,38 +35,38 @@ impl Operation for RawOutput { } fn selection_criteria(&self) -> Option<&crate::selection_criteria::SelectionCriteria> { - None + self.0.selection_criteria() } fn is_acknowledged(&self) -> bool { - self.write_concern() - .map(crate::options::WriteConcern::is_acknowledged) - .unwrap_or(true) + self.0.is_acknowledged() } fn write_concern(&self) -> Option<&crate::options::WriteConcern> { - None + self.0.write_concern() } - fn supports_read_concern(&self, _description: &StreamDescription) -> bool { - false + fn supports_read_concern(&self, description: &StreamDescription) -> bool { + self.0.supports_read_concern(description) } fn supports_sessions(&self) -> bool { - true + self.0.supports_sessions() } fn retryability(&self) -> super::Retryability { - super::Retryability::None + self.0.retryability() } - fn update_for_retry(&mut self) {} + fn update_for_retry(&mut self) { + self.0.update_for_retry() + } fn pinned_connection(&self) -> Option<&crate::cmap::conn::PinnedConnectionHandle> { - None + self.0.pinned_connection() } fn name(&self) -> &str { - Self::NAME + self.0.name() } } \ No newline at end of file diff --git a/src/operation/run_command/mod.rs b/src/operation/run_command/mod.rs index 8460a643f..6cf77cc88 100644 --- a/src/operation/run_command/mod.rs +++ b/src/operation/run_command/mod.rs @@ -5,7 +5,7 @@ use std::convert::TryInto; use bson::RawBsonRef; -use super::{CursorBody, Operation, OperationWithDefaults}; +use super::{CursorBody, OperationWithDefaults}; use crate::{ bson::Document, client::SESSIONS_UNSUPPORTED_COMMANDS, From d1ccd799345db59b043c1b2925332b0a86ec7f50 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Tue, 26 Jul 2022 10:07:56 -0400 Subject: [PATCH 05/35] wip --- src/client/csfle/state_machine.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index c07c847b9..337586cd3 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -17,13 +17,20 @@ impl Client { loop { match ctx.state()? { State::NeedMongoCollinfo => { - let filter: Document = raw_to_doc(ctx.mongo_op()?)?; + let filter = raw_to_doc(ctx.mongo_op()?)?; let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?; let mut cursor = db.list_collections(filter, None).await?; if cursor.advance().await? { ctx.mongo_feed(cursor.current())?; } } + State::NeedMongoMarkings => { + let command = raw_to_doc(ctx.mongo_op()?)?; + let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; + let result = db.run_command(command, None).await?; + } + State::Ready => result = ctx.finalize()?.to_owned(), + State::Done => break, _ => todo!(), } } From 1709fd38550c432224c8700a75cf3c69aee9f7ff Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Tue, 26 Jul 2022 15:17:30 -0400 Subject: [PATCH 06/35] raw runcommand --- src/client/csfle/state_machine.rs | 14 +++++++++++--- src/operation/mod.rs | 11 +++++++++++ src/operation/run_command/mod.rs | 31 ++++++++++++++++++++++++++----- src/operation/run_command/test.rs | 5 ++++- 4 files changed, 52 insertions(+), 9 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 337586cd3..e224f3212 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -5,7 +5,7 @@ use mongocrypt::ctx::{Ctx, State}; use crate::{Client, Database}; use crate::error::{Error, Result}; -use crate::operation::{ListCollections, RawOutput}; +use crate::operation::{ListCollections, RawOutput, RunCommand}; fn raw_to_doc(raw: &RawDocument) -> Result { raw.try_into().map_err(|e| Error::internal("???")) @@ -23,11 +23,19 @@ impl Client { if cursor.advance().await? { ctx.mongo_feed(cursor.current())?; } + ctx.mongo_done()?; } State::NeedMongoMarkings => { - let command = raw_to_doc(ctx.mongo_op()?)?; let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; - let result = db.run_command(command, None).await?; + let op = RawOutput(RunCommand::new_raw( + db.name().to_string(), + ctx.mongo_op()?.to_raw_document_buf(), + None, + None, + )?); + let result = self.execute_operation(op, None).await?; + ctx.mongo_feed(result.raw_body())?; + ctx.mongo_done()?; } State::Ready => result = ctx.finalize()?.to_owned(), State::Done => break, diff --git a/src/operation/mod.rs b/src/operation/mod.rs index ab8b94e3f..f7d441448 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -153,6 +153,17 @@ impl CommandBody for Document { } } +impl CommandBody for RawDocumentBuf { + fn should_redact(&self) -> bool { + if let Some(Ok((command_name, _))) = self.into_iter().next() { + HELLO_COMMAND_NAMES.contains(command_name.to_lowercase().as_str()) + && self.get("speculativeAuthenticate").ok().flatten().is_some() + } else { + false + } + } +} + impl Command { pub(crate) fn should_redact(&self) -> bool { let name = self.name.to_lowercase(); diff --git a/src/operation/run_command/mod.rs b/src/operation/run_command/mod.rs index 6cf77cc88..fa33888fb 100644 --- a/src/operation/run_command/mod.rs +++ b/src/operation/run_command/mod.rs @@ -3,7 +3,7 @@ mod test; use std::convert::TryInto; -use bson::RawBsonRef; +use bson::{RawBsonRef, RawDocumentBuf}; use super::{CursorBody, OperationWithDefaults}; use crate::{ @@ -18,7 +18,7 @@ use crate::{ #[derive(Debug)] pub(crate) struct RunCommand<'conn> { db: String, - command: Document, + command: RawDocumentBuf, selection_criteria: Option, write_concern: Option, pinned_connection: Option<&'conn PinnedConnectionHandle>, @@ -36,6 +36,27 @@ impl<'conn> RunCommand<'conn> { .map(|doc| bson::from_bson::(doc.clone())) .transpose()?; + Ok(Self { + db, + command: RawDocumentBuf::from_document(&command)?, + selection_criteria, + write_concern, + pinned_connection, + }) + } + + pub(crate) fn new_raw( + db: String, + command: RawDocumentBuf, + selection_criteria: Option, + pinned_connection: Option<&'conn PinnedConnectionHandle>, + ) -> Result { + let write_concern = command + .get("writeConcern")? + .and_then(|b| b.as_document()) + .map(|doc| bson::from_slice::(doc.as_bytes())) + .transpose()?; + Ok(Self { db, command, @@ -46,19 +67,19 @@ impl<'conn> RunCommand<'conn> { } fn command_name(&self) -> Option<&str> { - self.command.keys().next().map(String::as_str) + self.command.into_iter().next().and_then(|r| r.ok()).map(|(k, _)| k) } } impl<'conn> OperationWithDefaults for RunCommand<'conn> { type O = Document; - type Command = Document; + type Command = RawDocumentBuf; // Since we can't actually specify a string statically here, we just put a descriptive string // that should fail loudly if accidentally passed to the server. const NAME: &'static str = "$genericRunCommand"; - fn build(&mut self, _description: &StreamDescription) -> Result { + fn build(&mut self, _description: &StreamDescription) -> Result> { let command_name = self .command_name() .ok_or_else(|| ErrorKind::InvalidArgument { diff --git a/src/operation/run_command/test.rs b/src/operation/run_command/test.rs index 901aad90d..ad2b89a44 100644 --- a/src/operation/run_command/test.rs +++ b/src/operation/run_command/test.rs @@ -1,3 +1,5 @@ +use std::convert::TryInto; + use bson::Timestamp; use super::RunCommand; @@ -20,7 +22,8 @@ fn build() { command .body .get("hello") - .and_then(crate::bson_util::get_int), + .unwrap() + .and_then(|raw| crate::bson_util::get_int(&raw.try_into().unwrap())), Some(1) ); } From c7d02d96c84e1ff07c977942eaba2a27e932892a Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 12:02:11 -0400 Subject: [PATCH 07/35] local boxed run_mongocrypt_ctx --- src/client/csfle.rs | 4 ++ src/client/csfle/state_machine.rs | 63 +++++++++++++++++-------------- src/client/executor.rs | 19 +++++++++- 3 files changed, 57 insertions(+), 29 deletions(-) diff --git a/src/client/csfle.rs b/src/client/csfle.rs index f4f569629..b09439f65 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -66,6 +66,10 @@ impl ClientState { &self.opts } + pub(crate) fn crypt(&self) -> &Crypt { + &self.crypt + } + fn make_crypt(opts: &AutoEncryptionOptions) -> Result { let mut builder = Crypt::builder(); if Some(true) != opts.bypass_auto_encryption { diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index e224f3212..f87e5480e 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,6 +1,7 @@ use std::convert::TryInto; use bson::{RawDocumentBuf, Document, RawDocument}; +use futures_core::future::{BoxFuture, LocalBoxFuture}; use mongocrypt::ctx::{Ctx, State}; use crate::{Client, Database}; @@ -12,36 +13,42 @@ fn raw_to_doc(raw: &RawDocument) -> Result { } impl Client { - pub(crate) async fn run_mongocrypt_ctx(&self, ctx: &mut Ctx, db: Option) -> Result { - let mut result = RawDocumentBuf::new(); - loop { - match ctx.state()? { - State::NeedMongoCollinfo => { - let filter = raw_to_doc(ctx.mongo_op()?)?; - let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?; - let mut cursor = db.list_collections(filter, None).await?; - if cursor.advance().await? { - ctx.mongo_feed(cursor.current())?; + pub(crate) fn run_mongocrypt_ctx<'a>(&'a self, mut ctx: Ctx, db: Option<&'a str>) -> LocalBoxFuture<'a, Result> { + Box::pin(async move { + let mut result = None; + loop { + let state = ctx.state()?; + match state { + State::NeedMongoCollinfo => { + let filter = raw_to_doc(ctx.mongo_op()?)?; + let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); + let mut cursor = db.list_collections(filter, None).await?; + if cursor.advance().await? { + ctx.mongo_feed(cursor.current())?; + } + ctx.mongo_done()?; } - ctx.mongo_done()?; - } - State::NeedMongoMarkings => { - let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; - let op = RawOutput(RunCommand::new_raw( - db.name().to_string(), - ctx.mongo_op()?.to_raw_document_buf(), - None, - None, - )?); - let result = self.execute_operation(op, None).await?; - ctx.mongo_feed(result.raw_body())?; - ctx.mongo_done()?; + State::NeedMongoMarkings => { + let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; + let op = RawOutput(RunCommand::new_raw( + db.to_string(), + ctx.mongo_op()?.to_raw_document_buf(), + None, + None, + )?); + let guard = self.inner.csfle.read().await; + let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; + let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; + let result = mongocryptd_client.execute_operation(op, None).await?; + ctx.mongo_feed(result.raw_body())?; + ctx.mongo_done()?; + } + State::Ready => result = Some(ctx.finalize()?.to_owned()), + State::Done => break, + _ => todo!(), } - State::Ready => result = ctx.finalize()?.to_owned(), - State::Done => break, - _ => todo!(), } - } - Ok(result) + result.ok_or_else(|| Error::internal("libmongocrypt terminated without output")) + }) } } \ No newline at end of file diff --git a/src/client/executor.rs b/src/client/executor.rs index 731240fa5..21373c1d8 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,5 +1,6 @@ use bson::{doc, RawBsonRef, RawDocument, Timestamp}; use lazy_static::lazy_static; +use mongocrypt::ctx::{CtxBuilder, Ctx}; use serde::de::DeserializeOwned; use std::{collections::HashSet, sync::Arc, time::Instant}; @@ -597,7 +598,23 @@ impl Client { let cmd_name = cmd.name.clone(); let target_db = cmd.target_db.clone(); - let serialized = op.serialize_command(cmd)?; + let mut serialized = op.serialize_command(cmd)?; + #[cfg(feature = "csfle")] + { + let guard = self.inner.csfle.read().await; + if let Some(csfle) = guard.as_ref() { + if csfle.opts().bypass_auto_encryption != Some(true) { + /* + let ctx = csfle + .crypt() + .ctx_builder() + .build_encrypt(&target_db, RawDocument::from_bytes(&serialized)?)?; + let new_bytes = self.run_mongocrypt_ctx(ctx, Some(target_db.clone())).await?; + serialized = new_bytes.into_bytes(); + */ + } + } + } let raw_cmd = RawCommand { name: cmd_name.clone(), target_db, From 99f2c187675751bec9354161258c17d28760b7bf Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 14:04:50 -0400 Subject: [PATCH 08/35] dedicated ctx thread wip --- src/client/csfle.rs | 6 +- src/client/csfle/state_machine.rs | 97 ++++++++++++++++++++++++++++++- src/client/executor.rs | 41 ++++++++++++- 3 files changed, 137 insertions(+), 7 deletions(-) diff --git a/src/client/csfle.rs b/src/client/csfle.rs index b09439f65..e12275efc 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -3,7 +3,7 @@ mod state_machine; use std::{ path::Path, - process::{Command, Stdio}, + process::{Command, Stdio}, sync::Arc, }; use derivative::Derivative; @@ -32,7 +32,7 @@ use super::WeakClient; pub(super) struct ClientState { #[derivative(Debug = "ignore")] #[allow(dead_code)] - crypt: Crypt, + crypt: Arc, mongocryptd_client: Option, aux_clients: AuxClients, opts: AutoEncryptionOptions, @@ -55,7 +55,7 @@ impl ClientState { let aux_clients = Self::make_aux_clients(client, &opts)?; Ok(Self { - crypt, + crypt: Arc::new(crypt), mongocryptd_client, aux_clients, opts, diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index f87e5480e..1f259b1aa 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,8 +1,12 @@ use std::convert::TryInto; +use std::sync::Arc; +use std::sync::mpsc as sync_mpsc; +use std::thread; use bson::{RawDocumentBuf, Document, RawDocument}; use futures_core::future::{BoxFuture, LocalBoxFuture}; -use mongocrypt::ctx::{Ctx, State}; +use mongocrypt::ctx::{Ctx, State, CtxBuilder}; +use tokio::sync::mpsc::UnboundedSender; use crate::{Client, Database}; use crate::error::{Error, Result}; @@ -13,7 +17,42 @@ fn raw_to_doc(raw: &RawDocument) -> Result { } impl Client { - pub(crate) fn run_mongocrypt_ctx<'a>(&'a self, mut ctx: Ctx, db: Option<&'a str>) -> LocalBoxFuture<'a, Result> { + pub(crate) async fn run_mongocrypt_ctx(&self, build_ctx: impl FnOnce(CtxBuilder) -> Ctx + Send + 'static, db: Option<&str>) -> Result { + let guard = self.inner.csfle.read().await; + let crypt = match guard.as_ref() { + Some(csfle) => Arc::clone(&csfle.crypt), + None => return Err(Error::internal("no csfle state for mongocrypt ctx")), + }; + let (send, mut recv) = tokio::sync::mpsc::unbounded_channel(); + thread::spawn(move || { + let builder = crypt.ctx_builder(); + let ctx = build_ctx(builder); + match ctx_loop(ctx, send.clone()) { + Ok(Some(doc)) => send.send(CtxRequest::Done(doc)), + Ok(None) => send.send(CtxRequest::Err(Error::internal("libmongocrypt terminated without output"))), + Err(e) => send.send(CtxRequest::Err(e)), + } + }); + fn thread_err() -> Error { + Error::internal("ctx thread unexpectedly terminated") + } + loop { + let request = recv.recv().await.ok_or_else(thread_err)?; + match request { + CtxRequest::NeedMongoCollinfo { filter, reply } => { + let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); + let mut cursor = db.list_collections(filter, None).await?; + if cursor.advance().await? { + reply.send(Some(cursor.current().to_raw_document_buf())); + } else { + reply.send(None); + } + } + CtxRequest::Done(doc) => return Ok(doc), + CtxRequest::Err(e) => return Err(e), + } + } + /* Box::pin(async move { let mut result = None; loop { @@ -50,5 +89,59 @@ impl Client { } result.ok_or_else(|| Error::internal("libmongocrypt terminated without output")) }) + */ } +} + +fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result> { + let mut result = None; + loop { + match ctx.state()? { + State::NeedMongoCollinfo => { + let filter = raw_to_doc(ctx.mongo_op()?)?; + let (reply, reply_recv) = sync_oneshot(); + // TODO: terminate loop if send fails + send.send(CtxRequest::NeedMongoCollinfo { filter, reply }); + match reply_recv.recv() { + Ok(Some(v)) => ctx.mongo_feed(&v)?, + Ok(None) => (), + Err(_) => return Ok(None), // waiting async task terminated, no need to keep running + } + ctx.mongo_done()?; + } + State::Done => break, + _ => todo!(), + } + } + Ok(result) +} + +enum CtxRequest { + NeedMongoCollinfo { + filter: Document, + reply: SyncOneshotSender>, + }, + Done(RawDocumentBuf), + Err(Error), +} + +struct SyncOneshotSender(sync_mpsc::SyncSender); + +impl SyncOneshotSender { + fn send(self, value: T) -> std::result::Result<(), T> { + self.0.send(value).map_err(|sync_mpsc::SendError(value)| value) + } +} + +struct SyncOneshotReceiver(sync_mpsc::Receiver); + +impl SyncOneshotReceiver { + fn recv(self) -> std::result::Result { + self.0.recv() + } +} + +fn sync_oneshot() -> (SyncOneshotSender, SyncOneshotReceiver) { + let (sender, receiver) = sync_mpsc::sync_channel(1); + (SyncOneshotSender(sender), SyncOneshotReceiver(receiver)) } \ No newline at end of file diff --git a/src/client/executor.rs b/src/client/executor.rs index 21373c1d8..9752df02f 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,11 +1,12 @@ -use bson::{doc, RawBsonRef, RawDocument, Timestamp}; +use bson::{doc, RawBsonRef, RawDocument, Timestamp, RawDocumentBuf}; use lazy_static::lazy_static; use mongocrypt::ctx::{CtxBuilder, Ctx}; use serde::de::DeserializeOwned; +use tokio::task::LocalSet; use std::{collections::HashSet, sync::Arc, time::Instant}; -use super::{session::TransactionState, Client, ClientSession}; +use super::{session::TransactionState, Client, ClientSession, csfle::ClientState}; use crate::{ bson::Document, change_stream::{ @@ -601,9 +602,30 @@ impl Client { let mut serialized = op.serialize_command(cmd)?; #[cfg(feature = "csfle")] { + if self.inner.csfle.read().await.is_some() { + let client = self.clone(); + /* + local.spawn_local(async move { + let guard = match client.inner.csfle.read().await.as_ref() { + Some(g) => g, + None => return, + }; + }); + */ + } let guard = self.inner.csfle.read().await; if let Some(csfle) = guard.as_ref() { if csfle.opts().bypass_auto_encryption != Some(true) { + let local = LocalSet::new(); + let target_db = &target_db; + let serialized = &serialized; + let client = self.clone(); + /* + local.spawn_local(async move { + let new_message = client.auto_encrypt(serialized, todo!(), target_db).await; + }); + */ + //local.await; /* let ctx = csfle .crypt() @@ -782,6 +804,21 @@ impl Client { } } + /* + async fn auto_encrypt( + &self, + serialized: &[u8], + csfle: &ClientState, + target_db: &str, + ) -> Result { + let ctx = csfle + .crypt() + .ctx_builder() + .build_encrypt(&target_db, RawDocument::from_bytes(&serialized)?)?; + self.run_mongocrypt_ctx(ctx, Some(target_db)).await + } + */ + /// Start an implicit session if the operation and write concern are compatible with sessions. async fn start_implicit_session(&self, op: &T) -> Result> { match self.get_session_support_status().await? { From 21f0189ab8f5d81d9b4d368d0579df235bbcdba1 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 14:18:22 -0400 Subject: [PATCH 09/35] dedicated ctx thread scaffolding --- src/client/csfle/state_machine.rs | 12 +++++-- src/client/executor.rs | 59 ++++++++----------------------- 2 files changed, 23 insertions(+), 48 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 1f259b1aa..b380751a3 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -17,7 +17,7 @@ fn raw_to_doc(raw: &RawDocument) -> Result { } impl Client { - pub(crate) async fn run_mongocrypt_ctx(&self, build_ctx: impl FnOnce(CtxBuilder) -> Ctx + Send + 'static, db: Option<&str>) -> Result { + pub(crate) async fn run_mongocrypt_ctx(&self, build_ctx: impl FnOnce(CtxBuilder) -> Result + Send + 'static, db: Option<&str>) -> Result { let guard = self.inner.csfle.read().await; let crypt = match guard.as_ref() { Some(csfle) => Arc::clone(&csfle.crypt), @@ -26,12 +26,18 @@ impl Client { let (send, mut recv) = tokio::sync::mpsc::unbounded_channel(); thread::spawn(move || { let builder = crypt.ctx_builder(); - let ctx = build_ctx(builder); + let ctx = match build_ctx(builder) { + Ok(c) => c, + Err(e) => { + send.send(CtxRequest::Err(e)); + return; + } + }; match ctx_loop(ctx, send.clone()) { Ok(Some(doc)) => send.send(CtxRequest::Done(doc)), Ok(None) => send.send(CtxRequest::Err(Error::internal("libmongocrypt terminated without output"))), Err(e) => send.send(CtxRequest::Err(e)), - } + }; }); fn thread_err() -> Error { Error::internal("ctx thread unexpectedly terminated") diff --git a/src/client/executor.rs b/src/client/executor.rs index 9752df02f..d0c6be96c 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,4 +1,5 @@ use bson::{doc, RawBsonRef, RawDocument, Timestamp, RawDocumentBuf}; +use futures_core::future::BoxFuture; use lazy_static::lazy_static; use mongocrypt::ctx::{CtxBuilder, Ctx}; use serde::de::DeserializeOwned; @@ -603,38 +604,7 @@ impl Client { #[cfg(feature = "csfle")] { if self.inner.csfle.read().await.is_some() { - let client = self.clone(); - /* - local.spawn_local(async move { - let guard = match client.inner.csfle.read().await.as_ref() { - Some(g) => g, - None => return, - }; - }); - */ - } - let guard = self.inner.csfle.read().await; - if let Some(csfle) = guard.as_ref() { - if csfle.opts().bypass_auto_encryption != Some(true) { - let local = LocalSet::new(); - let target_db = &target_db; - let serialized = &serialized; - let client = self.clone(); - /* - local.spawn_local(async move { - let new_message = client.auto_encrypt(serialized, todo!(), target_db).await; - }); - */ - //local.await; - /* - let ctx = csfle - .crypt() - .ctx_builder() - .build_encrypt(&target_db, RawDocument::from_bytes(&serialized)?)?; - let new_bytes = self.run_mongocrypt_ctx(ctx, Some(target_db.clone())).await?; - serialized = new_bytes.into_bytes(); - */ - } + serialized = self.auto_encrypt(&serialized, &target_db).await?.into_bytes(); } } let raw_cmd = RawCommand { @@ -804,20 +774,19 @@ impl Client { } } - /* - async fn auto_encrypt( - &self, - serialized: &[u8], - csfle: &ClientState, - target_db: &str, - ) -> Result { - let ctx = csfle - .crypt() - .ctx_builder() - .build_encrypt(&target_db, RawDocument::from_bytes(&serialized)?)?; - self.run_mongocrypt_ctx(ctx, Some(target_db)).await + fn auto_encrypt<'a>( + &'a self, + serialized: &'a [u8], + target_db: &'a str, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let serialized = serialized.to_vec(); + let db = target_db.to_string(); + self.run_mongocrypt_ctx(move |builder| { + Ok(builder.build_encrypt(&db, RawDocument::from_bytes(&serialized)?)?) + }, Some(target_db)).await + }) } - */ /// Start an implicit session if the operation and write concern are compatible with sessions. async fn start_implicit_session(&self, op: &T) -> Result> { From e118f52014049276e054adf640c04f6709c31ac8 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 14:24:48 -0400 Subject: [PATCH 10/35] cleanup --- src/client/csfle.rs | 4 ---- src/client/csfle/state_machine.rs | 21 ++++++++++----------- src/client/executor.rs | 4 +--- 3 files changed, 11 insertions(+), 18 deletions(-) diff --git a/src/client/csfle.rs b/src/client/csfle.rs index e12275efc..68f31f609 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -66,10 +66,6 @@ impl ClientState { &self.opts } - pub(crate) fn crypt(&self) -> &Crypt { - &self.crypt - } - fn make_crypt(opts: &AutoEncryptionOptions) -> Result { let mut builder = Crypt::builder(); if Some(true) != opts.bypass_auto_encryption { diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index b380751a3..68e345c0a 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -4,7 +4,6 @@ use std::sync::mpsc as sync_mpsc; use std::thread; use bson::{RawDocumentBuf, Document, RawDocument}; -use futures_core::future::{BoxFuture, LocalBoxFuture}; use mongocrypt::ctx::{Ctx, State, CtxBuilder}; use tokio::sync::mpsc::UnboundedSender; @@ -29,19 +28,17 @@ impl Client { let ctx = match build_ctx(builder) { Ok(c) => c, Err(e) => { - send.send(CtxRequest::Err(e)); + let _ = send.send(CtxRequest::Err(e)); return; } }; - match ctx_loop(ctx, send.clone()) { + let _ = match ctx_loop(ctx, send.clone()) { Ok(Some(doc)) => send.send(CtxRequest::Done(doc)), Ok(None) => send.send(CtxRequest::Err(Error::internal("libmongocrypt terminated without output"))), Err(e) => send.send(CtxRequest::Err(e)), }; }); - fn thread_err() -> Error { - Error::internal("ctx thread unexpectedly terminated") - } + let thread_err = || Error::internal("ctx thread unexpectedly terminated"); loop { let request = recv.recv().await.ok_or_else(thread_err)?; match request { @@ -49,10 +46,10 @@ impl Client { let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); let mut cursor = db.list_collections(filter, None).await?; if cursor.advance().await? { - reply.send(Some(cursor.current().to_raw_document_buf())); + reply.send(Some(cursor.current().to_raw_document_buf())) } else { - reply.send(None); - } + reply.send(None) + }.map_err(|_| thread_err())?; } CtxRequest::Done(doc) => return Ok(doc), CtxRequest::Err(e) => return Err(e), @@ -106,8 +103,9 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result { let filter = raw_to_doc(ctx.mongo_op()?)?; let (reply, reply_recv) = sync_oneshot(); - // TODO: terminate loop if send fails - send.send(CtxRequest::NeedMongoCollinfo { filter, reply }); + if let Err(_) = send.send(CtxRequest::NeedMongoCollinfo { filter, reply }) { + return Ok(None); + } match reply_recv.recv() { Ok(Some(v)) => ctx.mongo_feed(&v)?, Ok(None) => (), @@ -115,6 +113,7 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result result = Some(ctx.finalize()?.to_owned()), State::Done => break, _ => todo!(), } diff --git a/src/client/executor.rs b/src/client/executor.rs index d0c6be96c..d5bdd8b29 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,13 +1,11 @@ use bson::{doc, RawBsonRef, RawDocument, Timestamp, RawDocumentBuf}; use futures_core::future::BoxFuture; use lazy_static::lazy_static; -use mongocrypt::ctx::{CtxBuilder, Ctx}; use serde::de::DeserializeOwned; -use tokio::task::LocalSet; use std::{collections::HashSet, sync::Arc, time::Instant}; -use super::{session::TransactionState, Client, ClientSession, csfle::ClientState}; +use super::{session::TransactionState, Client, ClientSession}; use crate::{ bson::Document, change_stream::{ From 36c75db4a7f1269aa6189da5fc97ce2d0ee6102b Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 14:34:16 -0400 Subject: [PATCH 11/35] catchup --- src/client/csfle/state_machine.rs | 74 ++++++++++++++----------------- 1 file changed, 33 insertions(+), 41 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 68e345c0a..54435f4ac 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -7,12 +7,12 @@ use bson::{RawDocumentBuf, Document, RawDocument}; use mongocrypt::ctx::{Ctx, State, CtxBuilder}; use tokio::sync::mpsc::UnboundedSender; -use crate::{Client, Database}; +use crate::{Client}; use crate::error::{Error, Result}; -use crate::operation::{ListCollections, RawOutput, RunCommand}; +use crate::operation::{RawOutput, RunCommand}; fn raw_to_doc(raw: &RawDocument) -> Result { - raw.try_into().map_err(|e| Error::internal("???")) + raw.try_into().map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) } impl Client { @@ -51,48 +51,24 @@ impl Client { reply.send(None) }.map_err(|_| thread_err())?; } + CtxRequest::NeedMongoMarkings { command, reply } => { + let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; + let op = RawOutput(RunCommand::new_raw( + db.to_string(), + command, + None, + None, + )?); + let guard = self.inner.csfle.read().await; + let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; + let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; + let result = mongocryptd_client.execute_operation(op, None).await?; + reply.send(result.into_raw_document_buf()).map_err(|_| thread_err())?; + } CtxRequest::Done(doc) => return Ok(doc), CtxRequest::Err(e) => return Err(e), } } - /* - Box::pin(async move { - let mut result = None; - loop { - let state = ctx.state()?; - match state { - State::NeedMongoCollinfo => { - let filter = raw_to_doc(ctx.mongo_op()?)?; - let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); - let mut cursor = db.list_collections(filter, None).await?; - if cursor.advance().await? { - ctx.mongo_feed(cursor.current())?; - } - ctx.mongo_done()?; - } - State::NeedMongoMarkings => { - let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; - let op = RawOutput(RunCommand::new_raw( - db.to_string(), - ctx.mongo_op()?.to_raw_document_buf(), - None, - None, - )?); - let guard = self.inner.csfle.read().await; - let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; - let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; - let result = mongocryptd_client.execute_operation(op, None).await?; - ctx.mongo_feed(result.raw_body())?; - ctx.mongo_done()?; - } - State::Ready => result = Some(ctx.finalize()?.to_owned()), - State::Done => break, - _ => todo!(), - } - } - result.ok_or_else(|| Error::internal("libmongocrypt terminated without output")) - }) - */ } } @@ -113,6 +89,18 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result { + let command = ctx.mongo_op()?.to_raw_document_buf(); + let (reply, reply_recv) = sync_oneshot(); + if let Err(_) = send.send(CtxRequest::NeedMongoMarkings { command, reply }) { + return Ok(None); + } + match reply_recv.recv() { + Ok(v) => ctx.mongo_feed(&v)?, + Err(_) => return Ok(None), + } + ctx.mongo_done()?; + } State::Ready => result = Some(ctx.finalize()?.to_owned()), State::Done => break, _ => todo!(), @@ -126,6 +114,10 @@ enum CtxRequest { filter: Document, reply: SyncOneshotSender>, }, + NeedMongoMarkings { + command: RawDocumentBuf, + reply: SyncOneshotSender, + }, Done(RawDocumentBuf), Err(Error), } From f586c4481e84df27d988ecc7fe950854a4d4b828 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 27 Jul 2022 14:41:17 -0400 Subject: [PATCH 12/35] tidy error handling --- src/client/csfle/state_machine.rs | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 54435f4ac..9c5f37684 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -38,6 +38,7 @@ impl Client { Err(e) => send.send(CtxRequest::Err(e)), }; }); + let thread_err = || Error::internal("ctx thread unexpectedly terminated"); loop { let request = recv.recv().await.ok_or_else(thread_err)?; @@ -72,6 +73,10 @@ impl Client { } } +fn ctx_err(_: T) -> Error { + Error::internal("ctx thread could not communicate with async task") +} + fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result> { let mut result = None; loop { @@ -79,26 +84,17 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result { let filter = raw_to_doc(ctx.mongo_op()?)?; let (reply, reply_recv) = sync_oneshot(); - if let Err(_) = send.send(CtxRequest::NeedMongoCollinfo { filter, reply }) { - return Ok(None); - } - match reply_recv.recv() { - Ok(Some(v)) => ctx.mongo_feed(&v)?, - Ok(None) => (), - Err(_) => return Ok(None), // waiting async task terminated, no need to keep running + send.send(CtxRequest::NeedMongoCollinfo { filter, reply }).map_err(ctx_err)?; + if let Some(v) = reply_recv.recv().map_err(ctx_err)? { + ctx.mongo_feed(&v)?; } ctx.mongo_done()?; } State::NeedMongoMarkings => { let command = ctx.mongo_op()?.to_raw_document_buf(); let (reply, reply_recv) = sync_oneshot(); - if let Err(_) = send.send(CtxRequest::NeedMongoMarkings { command, reply }) { - return Ok(None); - } - match reply_recv.recv() { - Ok(v) => ctx.mongo_feed(&v)?, - Err(_) => return Ok(None), - } + send.send(CtxRequest::NeedMongoMarkings { command, reply }).map_err(ctx_err)?; + ctx.mongo_feed(&reply_recv.recv().map_err(ctx_err)?)?; ctx.mongo_done()?; } State::Ready => result = Some(ctx.finalize()?.to_owned()), From 72440044988cba78a527dac3a327513b356725ea Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 28 Jul 2022 11:19:21 -0400 Subject: [PATCH 13/35] needmongokeys --- src/client/csfle/state_machine.rs | 71 +++++++++++++++++++++---------- 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 9c5f37684..c7063ca5a 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -4,6 +4,7 @@ use std::sync::mpsc as sync_mpsc; use std::thread; use bson::{RawDocumentBuf, Document, RawDocument}; +use futures_util::TryStreamExt; use mongocrypt::ctx::{Ctx, State, CtxBuilder}; use tokio::sync::mpsc::UnboundedSender; @@ -39,18 +40,17 @@ impl Client { }; }); - let thread_err = || Error::internal("ctx thread unexpectedly terminated"); loop { let request = recv.recv().await.ok_or_else(thread_err)?; match request { CtxRequest::NeedMongoCollinfo { filter, reply } => { let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); let mut cursor = db.list_collections(filter, None).await?; - if cursor.advance().await? { - reply.send(Some(cursor.current().to_raw_document_buf())) + reply.send(if cursor.advance().await? { + Some(cursor.current().to_raw_document_buf()) } else { - reply.send(None) - }.map_err(|_| thread_err())?; + None + })?; } CtxRequest::NeedMongoMarkings { command, reply } => { let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; @@ -64,7 +64,16 @@ impl Client { let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; let result = mongocryptd_client.execute_operation(op, None).await?; - reply.send(result.into_raw_document_buf()).map_err(|_| thread_err())?; + reply.send(result.into_raw_document_buf())?; + } + CtxRequest::NeedMongoKeys { filter, reply } => { + let guard = self.inner.csfle.read().await; + let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; + let kv_ns = &csfle.opts.key_vault_namespace; + let kv_client = csfle.aux_clients.key_vault_client.upgrade().ok_or_else(|| Error::internal("key vault client dropped"))?; + let kv_coll = kv_client.database(&kv_ns.db).collection::(&kv_ns.coll); + let results: Vec<_> = kv_coll.find(filter, None).await?.try_collect().await?; + reply.send(results)?; } CtxRequest::Done(doc) => return Ok(doc), CtxRequest::Err(e) => return Err(e), @@ -73,6 +82,10 @@ impl Client { } } +fn thread_err() -> Error { + Error::internal("ctx thread unexpectedly terminated") +} + fn ctx_err(_: T) -> Error { Error::internal("ctx thread could not communicate with async task") } @@ -83,18 +96,27 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result { let filter = raw_to_doc(ctx.mongo_op()?)?; - let (reply, reply_recv) = sync_oneshot(); + let (reply, reply_recv) = reply_oneshot(); send.send(CtxRequest::NeedMongoCollinfo { filter, reply }).map_err(ctx_err)?; - if let Some(v) = reply_recv.recv().map_err(ctx_err)? { + if let Some(v) = reply_recv.recv()? { ctx.mongo_feed(&v)?; } ctx.mongo_done()?; } State::NeedMongoMarkings => { let command = ctx.mongo_op()?.to_raw_document_buf(); - let (reply, reply_recv) = sync_oneshot(); + let (reply, reply_recv) = reply_oneshot(); send.send(CtxRequest::NeedMongoMarkings { command, reply }).map_err(ctx_err)?; - ctx.mongo_feed(&reply_recv.recv().map_err(ctx_err)?)?; + ctx.mongo_feed(&reply_recv.recv()?)?; + ctx.mongo_done()?; + } + State::NeedMongoKeys => { + let filter = raw_to_doc(ctx.mongo_op()?)?; + let (reply, reply_recv) = reply_oneshot(); + send.send(CtxRequest::NeedMongoKeys { filter, reply }).map_err(ctx_err)?; + for v in reply_recv.recv()? { + ctx.mongo_feed(&v)?; + } ctx.mongo_done()?; } State::Ready => result = Some(ctx.finalize()?.to_owned()), @@ -108,33 +130,38 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result>, + reply: ReplySender>, }, NeedMongoMarkings { command: RawDocumentBuf, - reply: SyncOneshotSender, + reply: ReplySender, + }, + NeedMongoKeys { + filter: Document, + reply: ReplySender>, }, Done(RawDocumentBuf), Err(Error), } -struct SyncOneshotSender(sync_mpsc::SyncSender); +struct ReplySender(sync_mpsc::SyncSender); -impl SyncOneshotSender { - fn send(self, value: T) -> std::result::Result<(), T> { - self.0.send(value).map_err(|sync_mpsc::SendError(value)| value) +impl ReplySender { + fn send(self, value: T) -> Result<()> { + self.0.send(value).map_err(|_| thread_err()) } } -struct SyncOneshotReceiver(sync_mpsc::Receiver); +struct ReplyReceiver(sync_mpsc::Receiver); -impl SyncOneshotReceiver { - fn recv(self) -> std::result::Result { - self.0.recv() +impl ReplyReceiver { + fn recv(self) -> Result { + self.0.recv().map_err(ctx_err) } } -fn sync_oneshot() -> (SyncOneshotSender, SyncOneshotReceiver) { +/// This is a sync version of `tokio::sync::oneshot::channel`; sending will never block (and so can be used in async code), receiving will block until a value is sent. +fn reply_oneshot() -> (ReplySender, ReplyReceiver) { let (sender, receiver) = sync_mpsc::sync_channel(1); - (SyncOneshotSender(sender), SyncOneshotReceiver(receiver)) + (ReplySender(sender), ReplyReceiver(receiver)) } \ No newline at end of file From 2646bedfe96e104efdd034a7184249eb95ab08d9 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Fri, 29 Jul 2022 15:42:34 -0400 Subject: [PATCH 14/35] start to rejigger to new ctx --- Cargo.toml | 3 +- src/client/csfle.rs | 6 +-- src/client/csfle/state_machine.rs | 69 +++++++++++++++++++++++-------- src/client/executor.rs | 15 +++---- 4 files changed, 64 insertions(+), 29 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 073b6a59f..e3975b12e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,8 @@ hex = "0.4.0" hmac = "0.12.1" lazy_static = "1.4.0" md-5 = "0.10.1" -mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } +# mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } +mongocrypt = { path = "../libmongocrypt-rust/mongocrypt", optional = true } openssl = { version = "0.10.38", optional = true } openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } diff --git a/src/client/csfle.rs b/src/client/csfle.rs index 68f31f609..d6c902793 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -3,7 +3,7 @@ mod state_machine; use std::{ path::Path, - process::{Command, Stdio}, sync::Arc, + process::{Command, Stdio}, }; use derivative::Derivative; @@ -32,7 +32,7 @@ use super::WeakClient; pub(super) struct ClientState { #[derivative(Debug = "ignore")] #[allow(dead_code)] - crypt: Arc, + pub(crate) crypt: Crypt, mongocryptd_client: Option, aux_clients: AuxClients, opts: AutoEncryptionOptions, @@ -55,7 +55,7 @@ impl ClientState { let aux_clients = Self::make_aux_clients(client, &opts)?; Ok(Self { - crypt: Arc::new(crypt), + crypt, mongocryptd_client, aux_clients, opts, diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index c7063ca5a..7f3252252 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,12 +1,9 @@ use std::convert::TryInto; use std::sync::Arc; -use std::sync::mpsc as sync_mpsc; -use std::thread; use bson::{RawDocumentBuf, Document, RawDocument}; use futures_util::TryStreamExt; -use mongocrypt::ctx::{Ctx, State, CtxBuilder}; -use tokio::sync::mpsc::UnboundedSender; +use mongocrypt::ctx::{Ctx, State}; use crate::{Client}; use crate::error::{Error, Result}; @@ -17,13 +14,26 @@ fn raw_to_doc(raw: &RawDocument) -> Result { } impl Client { - pub(crate) async fn run_mongocrypt_ctx(&self, build_ctx: impl FnOnce(CtxBuilder) -> Result + Send + 'static, db: Option<&str>) -> Result { + pub(crate) async fn run_mongocrypt_ctx(&self, mut ctx: Ctx, db: Option<&str>) -> Result { let guard = self.inner.csfle.read().await; let crypt = match guard.as_ref() { - Some(csfle) => Arc::clone(&csfle.crypt), + Some(csfle) => &csfle.crypt, None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; - let (send, mut recv) = tokio::sync::mpsc::unbounded_channel(); + let mut result = None; + loop { + match ctx.state()? { + State::Ready => result = Some(ctx.finalize()?.to_owned()), + State::Done => break, + _ => todo!(), + } + } + match result { + Some(doc) => Ok(doc), + None => Err(Error::internal("libmongocrypt terminated without output")), + } + /* + let (send, mut recv) = tokio_mpsc::unbounded_channel(); thread::spawn(move || { let builder = crypt.ctx_builder(); let ctx = match build_ctx(builder) { @@ -74,23 +84,20 @@ impl Client { let kv_coll = kv_client.database(&kv_ns.db).collection::(&kv_ns.coll); let results: Vec<_> = kv_coll.find(filter, None).await?.try_collect().await?; reply.send(results)?; + } + CtxRequest::NeedKms { messages } => { + } CtxRequest::Done(doc) => return Ok(doc), CtxRequest::Err(e) => return Err(e), } } + */ } } -fn thread_err() -> Error { - Error::internal("ctx thread unexpectedly terminated") -} - -fn ctx_err(_: T) -> Error { - Error::internal("ctx thread could not communicate with async task") -} - -fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result> { +/* +fn ctx_loop(mut ctx: Ctx, send: tokio_mpsc::UnboundedSender) -> Result> { let mut result = None; loop { match ctx.state()? { @@ -119,6 +126,17 @@ fn ctx_loop(mut ctx: Ctx, send: UnboundedSender) -> Result { + let scope = ctx.kms_scope(); + let mut messages = vec![]; + while let Some(kms_ctx) = scope.next_kms_ctx() { + let endpoint = kms_ctx.endpoint()?.to_string(); + let message = kms_ctx.message()?.to_vec(); + let (rsp_sender, responses_needed) = tokio_mpsc::unbounded_channel(); + messages.push(KmsMessage { endpoint, message, responses_needed }); + } + send.send(CtxRequest::NeedKms { messages }).map_err(ctx_err)?; + } State::Ready => result = Some(ctx.finalize()?.to_owned()), State::Done => break, _ => todo!(), @@ -140,10 +158,24 @@ enum CtxRequest { filter: Document, reply: ReplySender>, }, + NeedKms { + messages: Vec, + }, Done(RawDocumentBuf), Err(Error), } +struct KmsMessage { + endpoint: String, + message: Vec, + responses_needed: tokio_mpsc::UnboundedReceiver, +} + +struct ResponseNeeded { + bytes_needed: u8, + response: ReplySender>, +} + struct ReplySender(sync_mpsc::SyncSender); impl ReplySender { @@ -160,8 +192,9 @@ impl ReplyReceiver { } } -/// This is a sync version of `tokio::sync::oneshot::channel`; sending will never block (and so can be used in async code), receiving will block until a value is sent. +/// This is a sync version of `tokio::sync::oneshot::channel`; sending will never block (and so can be used in async code), receiving will synchronously block until a value is sent. fn reply_oneshot() -> (ReplySender, ReplyReceiver) { let (sender, receiver) = sync_mpsc::sync_channel(1); (ReplySender(sender), ReplyReceiver(receiver)) -} \ No newline at end of file +} +*/ \ No newline at end of file diff --git a/src/client/executor.rs b/src/client/executor.rs index d5bdd8b29..e1c89feac 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -5,7 +5,7 @@ use serde::de::DeserializeOwned; use std::{collections::HashSet, sync::Arc, time::Instant}; -use super::{session::TransactionState, Client, ClientSession}; +use super::{session::TransactionState, Client, ClientSession, csfle::ClientState}; use crate::{ bson::Document, change_stream::{ @@ -601,8 +601,9 @@ impl Client { let mut serialized = op.serialize_command(cmd)?; #[cfg(feature = "csfle")] { - if self.inner.csfle.read().await.is_some() { - serialized = self.auto_encrypt(&serialized, &target_db).await?.into_bytes(); + let guard = self.inner.csfle.read().await; + if let Some(ref csfle) = *guard { + serialized = self.auto_encrypt(csfle, &serialized, &target_db).await?.into_bytes(); } } let raw_cmd = RawCommand { @@ -774,15 +775,15 @@ impl Client { fn auto_encrypt<'a>( &'a self, + csfle: &'a ClientState, serialized: &'a [u8], target_db: &'a str, ) -> BoxFuture<'a, Result> { Box::pin(async move { - let serialized = serialized.to_vec(); + let doc = RawDocument::from_bytes(serialized)?.to_raw_document_buf(); let db = target_db.to_string(); - self.run_mongocrypt_ctx(move |builder| { - Ok(builder.build_encrypt(&db, RawDocument::from_bytes(&serialized)?)?) - }, Some(target_db)).await + let ctx = csfle.crypt.build_ctx(move |builder| builder.build_encrypt(&db, &doc))?; + self.run_mongocrypt_ctx(ctx, Some(target_db)).await }) } From edeaab409679ab5c731e283988179ca9510230f3 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Fri, 29 Jul 2022 16:23:28 -0400 Subject: [PATCH 15/35] progress --- src/client/csfle/state_machine.rs | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 7f3252252..018c3a6ae 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -2,7 +2,7 @@ use std::convert::TryInto; use std::sync::Arc; use bson::{RawDocumentBuf, Document, RawDocument}; -use futures_util::TryStreamExt; +use futures_util::{StreamExt, TryStreamExt, stream}; use mongocrypt::ctx::{Ctx, State}; use crate::{Client}; @@ -22,7 +22,28 @@ impl Client { }; let mut result = None; loop { - match ctx.state()? { + let state = ctx.state()?; + match state { + State::NeedMongoCollinfo => { + let filter = raw_to_doc(ctx.mongo_op()?)?; + let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); + let mut cursor = db.list_collections(filter, None).await?; + if cursor.advance().await? { + ctx.mongo_feed(cursor.current())?; + } + ctx.mongo_done()?; + } + State::NeedKms => { + let scope = ctx.kms_scope(); + let mut kms_ctxen: Vec> = vec![]; + while let Some(kms_ctx) = scope.next_kms_ctx()? { + kms_ctxen.push(Ok(kms_ctx)); + } + let f = stream::iter(kms_ctxen).try_for_each_concurrent(None, |kms_ctx| async move { + Ok(()) + }); + f.await; + } State::Ready => result = Some(ctx.finalize()?.to_owned()), State::Done => break, _ => todo!(), From 04a2faceba6b516f461640e5e1d0473bfc42c275 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 1 Aug 2022 14:22:54 -0400 Subject: [PATCH 16/35] select runtimes for ctx spawning --- Cargo.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index e3975b12e..e922d7328 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,12 +34,14 @@ tokio-runtime = [ "tokio/rt", "tokio/time", "serde_bytes", + "mongocrypt/tokio", ] async-std-runtime = [ "async-std", "async-std/attributes", "async-std-resolver", "tokio-util/compat", + "mongocrypt/async-std", ] sync = ["async-std-runtime"] tokio-sync = ["tokio-runtime"] From d0f974bb4b922fd68c5fd0966f9f5185bd14be36 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Tue, 2 Aug 2022 14:53:35 -0400 Subject: [PATCH 17/35] needkms --- src/client/csfle/state_machine.rs | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 018c3a6ae..a52c9a817 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -4,7 +4,11 @@ use std::sync::Arc; use bson::{RawDocumentBuf, Document, RawDocument}; use futures_util::{StreamExt, TryStreamExt, stream}; use mongocrypt::ctx::{Ctx, State}; +use tokio::io::{AsyncWriteExt, AsyncReadExt}; +use crate::client::options::{ServerAddress, TlsOptions}; +use crate::cmap::options::StreamOptions; +use crate::runtime::AsyncStream; use crate::{Client}; use crate::error::{Error, Result}; use crate::operation::{RawOutput, RunCommand}; @@ -39,10 +43,24 @@ impl Client { while let Some(kms_ctx) = scope.next_kms_ctx()? { kms_ctxen.push(Ok(kms_ctx)); } - let f = stream::iter(kms_ctxen).try_for_each_concurrent(None, |kms_ctx| async move { + stream::iter(kms_ctxen).try_for_each_concurrent(None, |mut kms_ctx| async move { + let endpoint = kms_ctx.endpoint()?; + let addr = ServerAddress::parse(endpoint)?; + let mut stream = AsyncStream::connect(StreamOptions::builder() + .address(addr) + .tls_options(TlsOptions::default()) + .build() + ).await?; + stream.write_all(kms_ctx.message()?).await?; + let mut buf = vec![]; + while kms_ctx.bytes_needed() > 0 { + let buf_size = kms_ctx.bytes_needed().try_into().map_err(|e| Error::internal(format!("buffer size overflow: {}", e)))?; + buf.resize(buf_size, 0); + let count = stream.read(&mut buf).await?; + kms_ctx.feed(&buf[0..count])?; + } Ok(()) - }); - f.await; + }).await?; } State::Ready => result = Some(ctx.finalize()?.to_owned()), State::Done => break, From ed668ccbe4d1fca35134cdd9ac4f2c515adfbc02 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 3 Aug 2022 12:57:47 -0400 Subject: [PATCH 18/35] fill back in previous states --- src/client/csfle/state_machine.rs | 198 +++++------------------------- src/client/executor.rs | 5 +- 2 files changed, 32 insertions(+), 171 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index a52c9a817..f3c1d8f06 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,8 +1,7 @@ use std::convert::TryInto; -use std::sync::Arc; use bson::{RawDocumentBuf, Document, RawDocument}; -use futures_util::{StreamExt, TryStreamExt, stream}; +use futures_util::{TryStreamExt, stream}; use mongocrypt::ctx::{Ctx, State}; use tokio::io::{AsyncWriteExt, AsyncReadExt}; @@ -20,8 +19,8 @@ fn raw_to_doc(raw: &RawDocument) -> Result { impl Client { pub(crate) async fn run_mongocrypt_ctx(&self, mut ctx: Ctx, db: Option<&str>) -> Result { let guard = self.inner.csfle.read().await; - let crypt = match guard.as_ref() { - Some(csfle) => &csfle.crypt, + let csfle = match guard.as_ref() { + Some(csfle) => csfle, None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; let mut result = None; @@ -37,6 +36,31 @@ impl Client { } ctx.mongo_done()?; } + State::NeedMongoMarkings => { + let command = ctx.mongo_op()?.to_raw_document_buf(); + let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; + let op = RawOutput(RunCommand::new_raw( + db.to_string(), + command, + None, + None, + )?); + let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; + let response = mongocryptd_client.execute_operation(op, None).await?; + ctx.mongo_feed(response.raw_body())?; + ctx.mongo_done()?; + } + State::NeedMongoKeys => { + let filter = raw_to_doc(ctx.mongo_op()?)?; + let kv_ns = &csfle.opts.key_vault_namespace; + let kv_client = csfle.aux_clients.key_vault_client.upgrade().ok_or_else(|| Error::internal("key vault client dropped"))?; + let kv_coll = kv_client.database(&kv_ns.db).collection::(&kv_ns.coll); + let mut cursor = kv_coll.find(filter, None).await?; + while let Some(result) = cursor.try_next().await? { + ctx.mongo_feed(&result)? + } + ctx.mongo_done()?; + } State::NeedKms => { let scope = ctx.kms_scope(); let mut kms_ctxen: Vec> = vec![]; @@ -71,169 +95,5 @@ impl Client { Some(doc) => Ok(doc), None => Err(Error::internal("libmongocrypt terminated without output")), } - /* - let (send, mut recv) = tokio_mpsc::unbounded_channel(); - thread::spawn(move || { - let builder = crypt.ctx_builder(); - let ctx = match build_ctx(builder) { - Ok(c) => c, - Err(e) => { - let _ = send.send(CtxRequest::Err(e)); - return; - } - }; - let _ = match ctx_loop(ctx, send.clone()) { - Ok(Some(doc)) => send.send(CtxRequest::Done(doc)), - Ok(None) => send.send(CtxRequest::Err(Error::internal("libmongocrypt terminated without output"))), - Err(e) => send.send(CtxRequest::Err(e)), - }; - }); - - loop { - let request = recv.recv().await.ok_or_else(thread_err)?; - match request { - CtxRequest::NeedMongoCollinfo { filter, reply } => { - let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); - let mut cursor = db.list_collections(filter, None).await?; - reply.send(if cursor.advance().await? { - Some(cursor.current().to_raw_document_buf()) - } else { - None - })?; - } - CtxRequest::NeedMongoMarkings { command, reply } => { - let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; - let op = RawOutput(RunCommand::new_raw( - db.to_string(), - command, - None, - None, - )?); - let guard = self.inner.csfle.read().await; - let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; - let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; - let result = mongocryptd_client.execute_operation(op, None).await?; - reply.send(result.into_raw_document_buf())?; - } - CtxRequest::NeedMongoKeys { filter, reply } => { - let guard = self.inner.csfle.read().await; - let csfle = guard.as_ref().ok_or_else(|| Error::internal("csfle state not found"))?; - let kv_ns = &csfle.opts.key_vault_namespace; - let kv_client = csfle.aux_clients.key_vault_client.upgrade().ok_or_else(|| Error::internal("key vault client dropped"))?; - let kv_coll = kv_client.database(&kv_ns.db).collection::(&kv_ns.coll); - let results: Vec<_> = kv_coll.find(filter, None).await?.try_collect().await?; - reply.send(results)?; - } - CtxRequest::NeedKms { messages } => { - - } - CtxRequest::Done(doc) => return Ok(doc), - CtxRequest::Err(e) => return Err(e), - } - } - */ } -} - -/* -fn ctx_loop(mut ctx: Ctx, send: tokio_mpsc::UnboundedSender) -> Result> { - let mut result = None; - loop { - match ctx.state()? { - State::NeedMongoCollinfo => { - let filter = raw_to_doc(ctx.mongo_op()?)?; - let (reply, reply_recv) = reply_oneshot(); - send.send(CtxRequest::NeedMongoCollinfo { filter, reply }).map_err(ctx_err)?; - if let Some(v) = reply_recv.recv()? { - ctx.mongo_feed(&v)?; - } - ctx.mongo_done()?; - } - State::NeedMongoMarkings => { - let command = ctx.mongo_op()?.to_raw_document_buf(); - let (reply, reply_recv) = reply_oneshot(); - send.send(CtxRequest::NeedMongoMarkings { command, reply }).map_err(ctx_err)?; - ctx.mongo_feed(&reply_recv.recv()?)?; - ctx.mongo_done()?; - } - State::NeedMongoKeys => { - let filter = raw_to_doc(ctx.mongo_op()?)?; - let (reply, reply_recv) = reply_oneshot(); - send.send(CtxRequest::NeedMongoKeys { filter, reply }).map_err(ctx_err)?; - for v in reply_recv.recv()? { - ctx.mongo_feed(&v)?; - } - ctx.mongo_done()?; - } - State::NeedKms => { - let scope = ctx.kms_scope(); - let mut messages = vec![]; - while let Some(kms_ctx) = scope.next_kms_ctx() { - let endpoint = kms_ctx.endpoint()?.to_string(); - let message = kms_ctx.message()?.to_vec(); - let (rsp_sender, responses_needed) = tokio_mpsc::unbounded_channel(); - messages.push(KmsMessage { endpoint, message, responses_needed }); - } - send.send(CtxRequest::NeedKms { messages }).map_err(ctx_err)?; - } - State::Ready => result = Some(ctx.finalize()?.to_owned()), - State::Done => break, - _ => todo!(), - } - } - Ok(result) -} - -enum CtxRequest { - NeedMongoCollinfo { - filter: Document, - reply: ReplySender>, - }, - NeedMongoMarkings { - command: RawDocumentBuf, - reply: ReplySender, - }, - NeedMongoKeys { - filter: Document, - reply: ReplySender>, - }, - NeedKms { - messages: Vec, - }, - Done(RawDocumentBuf), - Err(Error), -} - -struct KmsMessage { - endpoint: String, - message: Vec, - responses_needed: tokio_mpsc::UnboundedReceiver, -} - -struct ResponseNeeded { - bytes_needed: u8, - response: ReplySender>, -} - -struct ReplySender(sync_mpsc::SyncSender); - -impl ReplySender { - fn send(self, value: T) -> Result<()> { - self.0.send(value).map_err(|_| thread_err()) - } -} - -struct ReplyReceiver(sync_mpsc::Receiver); - -impl ReplyReceiver { - fn recv(self) -> Result { - self.0.recv().map_err(ctx_err) - } -} - -/// This is a sync version of `tokio::sync::oneshot::channel`; sending will never block (and so can be used in async code), receiving will synchronously block until a value is sent. -fn reply_oneshot() -> (ReplySender, ReplyReceiver) { - let (sender, receiver) = sync_mpsc::sync_channel(1); - (ReplySender(sender), ReplyReceiver(receiver)) -} -*/ \ No newline at end of file +} \ No newline at end of file diff --git a/src/client/executor.rs b/src/client/executor.rs index e1c89feac..5c8e3730f 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -5,7 +5,7 @@ use serde::de::DeserializeOwned; use std::{collections::HashSet, sync::Arc, time::Instant}; -use super::{session::TransactionState, Client, ClientSession, csfle::ClientState}; +use super::{session::TransactionState, Client, ClientSession}; use crate::{ bson::Document, change_stream::{ @@ -773,9 +773,10 @@ impl Client { } } + #[cfg(feature = "csfle")] fn auto_encrypt<'a>( &'a self, - csfle: &'a ClientState, + csfle: &'a super::csfle::ClientState, serialized: &'a [u8], target_db: &'a str, ) -> BoxFuture<'a, Result> { From 9c98ecbb5c3b3af812cc3b0dc198710d1c105b7f Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 3 Aug 2022 13:17:38 -0400 Subject: [PATCH 19/35] placeholder for needkmscredentials --- src/client/csfle/state_machine.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index f3c1d8f06..d102dd883 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -12,10 +12,6 @@ use crate::{Client}; use crate::error::{Error, Result}; use crate::operation::{RawOutput, RunCommand}; -fn raw_to_doc(raw: &RawDocument) -> Result { - raw.try_into().map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) -} - impl Client { pub(crate) async fn run_mongocrypt_ctx(&self, mut ctx: Ctx, db: Option<&str>) -> Result { let guard = self.inner.csfle.read().await; @@ -86,9 +82,10 @@ impl Client { Ok(()) }).await?; } + State::NeedKmsCredentials => todo!("RUST-1314"), State::Ready => result = Some(ctx.finalize()?.to_owned()), State::Done => break, - _ => todo!(), + s => return Err(Error::internal(format!("unhandled state {:?}", s))), } } match result { @@ -96,4 +93,8 @@ impl Client { None => Err(Error::internal("libmongocrypt terminated without output")), } } +} + +fn raw_to_doc(raw: &RawDocument) -> Result { + raw.try_into().map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) } \ No newline at end of file From 8a0cd31a1361e12c55ade7a9b282d00e338219a8 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 3 Aug 2022 13:29:27 -0400 Subject: [PATCH 20/35] auto decrypt --- src/client/executor.rs | 24 +++++++++++++++++++++++- src/cmap/conn/command.rs | 9 +++++---- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/client/executor.rs b/src/client/executor.rs index 5c8e3730f..1980cc947 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -737,7 +737,7 @@ impl Client { err.add_labels_and_update_pin(Some(connection), session, Some(retryability))?; op.handle_error(err) } - Ok(response) => { + Ok(mut response) => { self.emit_command_event(|handler| { let reply = if should_redact { Document::new() @@ -758,6 +758,15 @@ impl Client { handler.handle_command_succeeded_event(command_succeeded_event); }); + #[cfg(feature = "csfle")] + { + let guard = self.inner.csfle.read().await; + if let Some(ref csfle) = *guard { + let new_body = self.auto_decrypt(csfle, response.raw_body()).await?; + response = RawCommandResponse::new_raw(response.source, new_body); + } + } + match op.handle_response(response, connection.stream_description()?) { Ok(response) => Ok(response), Err(mut err) => { @@ -788,6 +797,19 @@ impl Client { }) } + #[cfg(feature = "csfle")] + fn auto_decrypt<'a>( + &'a self, + csfle: &'a super::csfle::ClientState, + doc: &'a RawDocument, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let doc = doc.to_raw_document_buf(); + let ctx = csfle.crypt.build_ctx(move |builder| builder.build_decrypt(&doc))?; + self.run_mongocrypt_ctx(ctx, None).await + }) + } + /// Start an implicit session if the operation and write concern are compatible with sessions. async fn start_implicit_session(&self, op: &T) -> Result> { match self.get_session_support_status().await? { diff --git a/src/cmap/conn/command.rs b/src/cmap/conn/command.rs index 265453feb..a6d6ddece 100644 --- a/src/cmap/conn/command.rs +++ b/src/cmap/conn/command.rs @@ -203,10 +203,11 @@ impl RawCommandResponse { pub(crate) fn new(source: ServerAddress, message: Message) -> Result { let raw = message.single_document_response()?; - Ok(Self { - source, - raw: RawDocumentBuf::from_bytes(raw)?, - }) + Ok(Self::new_raw(source, RawDocumentBuf::from_bytes(raw)?)) + } + + pub(crate) fn new_raw(source: ServerAddress, raw: RawDocumentBuf) -> Self { + Self { source, raw } } pub(crate) fn body<'a, T: Deserialize<'a>>(&'a self) -> Result { From 54362a663063bda8b30ac8d517a2ae206defa3eb Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 3 Aug 2022 13:41:13 -0400 Subject: [PATCH 21/35] rustfmt --- src/client/csfle/state_machine.rs | 105 ++++++++++++++---------- src/client/executor.rs | 29 ++++--- src/operation/abort_transaction/mod.rs | 4 +- src/operation/aggregate/mod.rs | 2 +- src/operation/create/mod.rs | 7 +- src/operation/drop_collection/mod.rs | 7 +- src/operation/drop_database/mod.rs | 7 +- src/operation/insert/mod.rs | 7 +- src/operation/mod.rs | 108 +++++++++++++------------ src/operation/raw_output.rs | 13 ++- src/operation/run_command/mod.rs | 6 +- 11 files changed, 179 insertions(+), 116 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index d102dd883..e7a99e461 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -1,19 +1,25 @@ use std::convert::TryInto; -use bson::{RawDocumentBuf, Document, RawDocument}; -use futures_util::{TryStreamExt, stream}; +use bson::{Document, RawDocument, RawDocumentBuf}; +use futures_util::{stream, TryStreamExt}; use mongocrypt::ctx::{Ctx, State}; -use tokio::io::{AsyncWriteExt, AsyncReadExt}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use crate::client::options::{ServerAddress, TlsOptions}; -use crate::cmap::options::StreamOptions; -use crate::runtime::AsyncStream; -use crate::{Client}; -use crate::error::{Error, Result}; -use crate::operation::{RawOutput, RunCommand}; +use crate::{ + client::options::{ServerAddress, TlsOptions}, + cmap::options::StreamOptions, + error::{Error, Result}, + operation::{RawOutput, RunCommand}, + runtime::AsyncStream, + Client, +}; impl Client { - pub(crate) async fn run_mongocrypt_ctx(&self, mut ctx: Ctx, db: Option<&str>) -> Result { + pub(crate) async fn run_mongocrypt_ctx( + &self, + mut ctx: Ctx, + db: Option<&str>, + ) -> Result { let guard = self.inner.csfle.read().await; let csfle = match guard.as_ref() { Some(csfle) => csfle, @@ -25,7 +31,9 @@ impl Client { match state { State::NeedMongoCollinfo => { let filter = raw_to_doc(ctx.mongo_op()?)?; - let db = self.database(db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoCollinfo state"))?); + let db = self.database(db.as_ref().ok_or_else(|| { + Error::internal("db required for NeedMongoCollinfo state") + })?); let mut cursor = db.list_collections(filter, None).await?; if cursor.advance().await? { ctx.mongo_feed(cursor.current())?; @@ -34,14 +42,14 @@ impl Client { } State::NeedMongoMarkings => { let command = ctx.mongo_op()?.to_raw_document_buf(); - let db = db.as_ref().ok_or_else(|| Error::internal("db required for NeedMongoMarkings state"))?; - let op = RawOutput(RunCommand::new_raw( - db.to_string(), - command, - None, - None, - )?); - let mongocryptd_client = csfle.mongocryptd_client.as_ref().ok_or_else(|| Error::internal("mongocryptd client not found"))?; + let db = db.as_ref().ok_or_else(|| { + Error::internal("db required for NeedMongoMarkings state") + })?; + let op = RawOutput(RunCommand::new_raw(db.to_string(), command, None, None)?); + let mongocryptd_client = csfle + .mongocryptd_client + .as_ref() + .ok_or_else(|| Error::internal("mongocryptd client not found"))?; let response = mongocryptd_client.execute_operation(op, None).await?; ctx.mongo_feed(response.raw_body())?; ctx.mongo_done()?; @@ -49,38 +57,50 @@ impl Client { State::NeedMongoKeys => { let filter = raw_to_doc(ctx.mongo_op()?)?; let kv_ns = &csfle.opts.key_vault_namespace; - let kv_client = csfle.aux_clients.key_vault_client.upgrade().ok_or_else(|| Error::internal("key vault client dropped"))?; - let kv_coll = kv_client.database(&kv_ns.db).collection::(&kv_ns.coll); + let kv_client = csfle + .aux_clients + .key_vault_client + .upgrade() + .ok_or_else(|| Error::internal("key vault client dropped"))?; + let kv_coll = kv_client + .database(&kv_ns.db) + .collection::(&kv_ns.coll); let mut cursor = kv_coll.find(filter, None).await?; while let Some(result) = cursor.try_next().await? { ctx.mongo_feed(&result)? } ctx.mongo_done()?; - } + } State::NeedKms => { let scope = ctx.kms_scope(); let mut kms_ctxen: Vec> = vec![]; while let Some(kms_ctx) = scope.next_kms_ctx()? { kms_ctxen.push(Ok(kms_ctx)); } - stream::iter(kms_ctxen).try_for_each_concurrent(None, |mut kms_ctx| async move { - let endpoint = kms_ctx.endpoint()?; - let addr = ServerAddress::parse(endpoint)?; - let mut stream = AsyncStream::connect(StreamOptions::builder() - .address(addr) - .tls_options(TlsOptions::default()) - .build() - ).await?; - stream.write_all(kms_ctx.message()?).await?; - let mut buf = vec![]; - while kms_ctx.bytes_needed() > 0 { - let buf_size = kms_ctx.bytes_needed().try_into().map_err(|e| Error::internal(format!("buffer size overflow: {}", e)))?; - buf.resize(buf_size, 0); - let count = stream.read(&mut buf).await?; - kms_ctx.feed(&buf[0..count])?; - } - Ok(()) - }).await?; + stream::iter(kms_ctxen) + .try_for_each_concurrent(None, |mut kms_ctx| async move { + let endpoint = kms_ctx.endpoint()?; + let addr = ServerAddress::parse(endpoint)?; + let mut stream = AsyncStream::connect( + StreamOptions::builder() + .address(addr) + .tls_options(TlsOptions::default()) + .build(), + ) + .await?; + stream.write_all(kms_ctx.message()?).await?; + let mut buf = vec![]; + while kms_ctx.bytes_needed() > 0 { + let buf_size = kms_ctx.bytes_needed().try_into().map_err(|e| { + Error::internal(format!("buffer size overflow: {}", e)) + })?; + buf.resize(buf_size, 0); + let count = stream.read(&mut buf).await?; + kms_ctx.feed(&buf[0..count])?; + } + Ok(()) + }) + .await?; } State::NeedKmsCredentials => todo!("RUST-1314"), State::Ready => result = Some(ctx.finalize()?.to_owned()), @@ -96,5 +116,6 @@ impl Client { } fn raw_to_doc(raw: &RawDocument) -> Result { - raw.try_into().map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) -} \ No newline at end of file + raw.try_into() + .map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) +} diff --git a/src/client/executor.rs b/src/client/executor.rs index 1980cc947..c22de17c6 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,4 +1,4 @@ -use bson::{doc, RawBsonRef, RawDocument, Timestamp, RawDocumentBuf}; +use bson::{doc, RawBsonRef, RawDocument, RawDocumentBuf, Timestamp}; use futures_core::future::BoxFuture; use lazy_static::lazy_static; use serde::de::DeserializeOwned; @@ -603,7 +603,12 @@ impl Client { { let guard = self.inner.csfle.read().await; if let Some(ref csfle) = *guard { - serialized = self.auto_encrypt(csfle, &serialized, &target_db).await?.into_bytes(); + if csfle.opts().bypass_auto_encryption != Some(true) { + serialized = self + .auto_encrypt(csfle, &serialized, &target_db) + .await? + .into_bytes(); + } } } let raw_cmd = RawCommand { @@ -761,9 +766,11 @@ impl Client { #[cfg(feature = "csfle")] { let guard = self.inner.csfle.read().await; - if let Some(ref csfle) = *guard { - let new_body = self.auto_decrypt(csfle, response.raw_body()).await?; - response = RawCommandResponse::new_raw(response.source, new_body); + if let Some(ref csfle) = *guard { + if csfle.opts().bypass_auto_encryption != Some(true) { + let new_body = self.auto_decrypt(csfle, response.raw_body()).await?; + response = RawCommandResponse::new_raw(response.source, new_body); + } } } @@ -784,7 +791,7 @@ impl Client { #[cfg(feature = "csfle")] fn auto_encrypt<'a>( - &'a self, + &'a self, csfle: &'a super::csfle::ClientState, serialized: &'a [u8], target_db: &'a str, @@ -792,20 +799,24 @@ impl Client { Box::pin(async move { let doc = RawDocument::from_bytes(serialized)?.to_raw_document_buf(); let db = target_db.to_string(); - let ctx = csfle.crypt.build_ctx(move |builder| builder.build_encrypt(&db, &doc))?; + let ctx = csfle + .crypt + .build_ctx(move |builder| builder.build_encrypt(&db, &doc))?; self.run_mongocrypt_ctx(ctx, Some(target_db)).await }) } #[cfg(feature = "csfle")] fn auto_decrypt<'a>( - &'a self, + &'a self, csfle: &'a super::csfle::ClientState, doc: &'a RawDocument, ) -> BoxFuture<'a, Result> { Box::pin(async move { let doc = doc.to_raw_document_buf(); - let ctx = csfle.crypt.build_ctx(move |builder| builder.build_decrypt(&doc))?; + let ctx = csfle + .crypt + .build_ctx(move |builder| builder.build_decrypt(&doc))?; self.run_mongocrypt_ctx(ctx, None).await }) } diff --git a/src/operation/abort_transaction/mod.rs b/src/operation/abort_transaction/mod.rs index 35a992993..7a08861dd 100644 --- a/src/operation/abort_transaction/mod.rs +++ b/src/operation/abort_transaction/mod.rs @@ -5,12 +5,12 @@ use crate::{ client::session::TransactionPin, cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{Retryability}, + operation::Retryability, options::WriteConcern, selection_criteria::SelectionCriteria, }; -use super::{WriteConcernOnlyBody, OperationWithDefaults}; +use super::{OperationWithDefaults, WriteConcernOnlyBody}; pub(crate) struct AbortTransaction { write_concern: Option, diff --git a/src/operation/aggregate/mod.rs b/src/operation/aggregate/mod.rs index b2ba5c8ab..1bb5c9ad0 100644 --- a/src/operation/aggregate/mod.rs +++ b/src/operation/aggregate/mod.rs @@ -14,7 +14,7 @@ use crate::{ Namespace, }; -use super::{CursorBody, WriteConcernOnlyBody, SERVER_4_2_0_WIRE_VERSION, OperationWithDefaults}; +use super::{CursorBody, OperationWithDefaults, WriteConcernOnlyBody, SERVER_4_2_0_WIRE_VERSION}; pub(crate) use change_stream::ChangeStreamAggregate; diff --git a/src/operation/create/mod.rs b/src/operation/create/mod.rs index d70cbe2c8..d531e8ca1 100644 --- a/src/operation/create/mod.rs +++ b/src/operation/create/mod.rs @@ -7,7 +7,12 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, + operation::{ + append_options, + remove_empty_write_concern, + OperationWithDefaults, + WriteConcernOnlyBody, + }, options::{CreateCollectionOptions, WriteConcern}, Namespace, }; diff --git a/src/operation/drop_collection/mod.rs b/src/operation/drop_collection/mod.rs index 6f9f29aff..7d270ef4e 100644 --- a/src/operation/drop_collection/mod.rs +++ b/src/operation/drop_collection/mod.rs @@ -7,7 +7,12 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::{Error, Result}, - operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, + operation::{ + append_options, + remove_empty_write_concern, + OperationWithDefaults, + WriteConcernOnlyBody, + }, options::{DropCollectionOptions, WriteConcern}, Namespace, }; diff --git a/src/operation/drop_database/mod.rs b/src/operation/drop_database/mod.rs index 53ba013d6..2cd2b8cf7 100644 --- a/src/operation/drop_database/mod.rs +++ b/src/operation/drop_database/mod.rs @@ -7,7 +7,12 @@ use crate::{ bson::doc, cmap::{Command, RawCommandResponse, StreamDescription}, error::Result, - operation::{append_options, remove_empty_write_concern, OperationWithDefaults, WriteConcernOnlyBody}, + operation::{ + append_options, + remove_empty_write_concern, + OperationWithDefaults, + WriteConcernOnlyBody, + }, options::{DropDatabaseOptions, WriteConcern}, }; diff --git a/src/operation/insert/mod.rs b/src/operation/insert/mod.rs index 031dfff41..7f6a48ea4 100644 --- a/src/operation/insert/mod.rs +++ b/src/operation/insert/mod.rs @@ -11,7 +11,12 @@ use crate::{ bson_util, cmap::{Command, RawCommandResponse, StreamDescription}, error::{BulkWriteFailure, Error, ErrorKind, Result}, - operation::{remove_empty_write_concern, OperationWithDefaults, Retryability, WriteResponseBody}, + operation::{ + remove_empty_write_concern, + OperationWithDefaults, + Retryability, + WriteResponseBody, + }, options::{InsertManyOptions, WriteConcern}, results::InsertManyResult, Namespace, diff --git a/src/operation/mod.rs b/src/operation/mod.rs index f7d441448..dae5e014f 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -75,8 +75,9 @@ pub(crate) use update::Update; const SERVER_4_2_0_WIRE_VERSION: i32 = 8; /// A trait modeling the behavior of a server side operation. -/// -/// No methods in this trait should have default behaviors to ensure that wrapper operations replicate all behavior. Default behavior is provided by the `OperationDefault` trait. +/// +/// No methods in this trait should have default behaviors to ensure that wrapper operations +/// replicate all behavior. Default behavior is provided by the `OperationDefault` trait. pub(crate) trait Operation { /// The output type of this operation. type O; @@ -390,7 +391,8 @@ macro_rules! remove_empty_write_concern { pub(crate) use remove_empty_write_concern; -// A mirror of the `Operation` trait, with default behavior where appropriate. Should only be implemented by leaf operation types. +// A mirror of the `Operation` trait, with default behavior where appropriate. Should only be +// implemented by leaf operation types. pub(crate) trait OperationWithDefaults { /// The output type of this operation. type O; @@ -475,53 +477,53 @@ pub(crate) trait OperationWithDefaults { } impl Operation for T { - type O = T::O; - type Command = T::Command; - const NAME: &'static str = T::NAME; - fn build(&mut self, description: &StreamDescription) -> Result> { - self.build(description) - } - fn serialize_command(&mut self, cmd: Command) -> Result> { - self.serialize_command(cmd) - } - fn extract_at_cluster_time(&self, response: &RawDocument) -> Result> { - self.extract_at_cluster_time(response) - } - fn handle_response( - &self, - response: RawCommandResponse, - description: &StreamDescription, - ) -> Result { - self.handle_response(response, description) - } - fn handle_error(&self, error: Error) -> Result { - self.handle_error(error) - } - fn selection_criteria(&self) -> Option<&SelectionCriteria> { - self.selection_criteria() - } - fn is_acknowledged(&self) -> bool { - self.is_acknowledged() - } - fn write_concern(&self) -> Option<&WriteConcern> { - self.write_concern() - } - fn supports_read_concern(&self, description: &StreamDescription) -> bool { - self.supports_read_concern(description) - } - fn supports_sessions(&self) -> bool { - self.supports_sessions() - } - fn retryability(&self) -> Retryability { - self.retryability() - } - fn update_for_retry(&mut self) { - self.update_for_retry() - } - fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> { - self.pinned_connection() - } - fn name(&self) -> &str { - self.name() - } -} \ No newline at end of file + type O = T::O; + type Command = T::Command; + const NAME: &'static str = T::NAME; + fn build(&mut self, description: &StreamDescription) -> Result> { + self.build(description) + } + fn serialize_command(&mut self, cmd: Command) -> Result> { + self.serialize_command(cmd) + } + fn extract_at_cluster_time(&self, response: &RawDocument) -> Result> { + self.extract_at_cluster_time(response) + } + fn handle_response( + &self, + response: RawCommandResponse, + description: &StreamDescription, + ) -> Result { + self.handle_response(response, description) + } + fn handle_error(&self, error: Error) -> Result { + self.handle_error(error) + } + fn selection_criteria(&self) -> Option<&SelectionCriteria> { + self.selection_criteria() + } + fn is_acknowledged(&self) -> bool { + self.is_acknowledged() + } + fn write_concern(&self) -> Option<&WriteConcern> { + self.write_concern() + } + fn supports_read_concern(&self, description: &StreamDescription) -> bool { + self.supports_read_concern(description) + } + fn supports_sessions(&self) -> bool { + self.supports_sessions() + } + fn retryability(&self) -> Retryability { + self.retryability() + } + fn update_for_retry(&mut self) { + self.update_for_retry() + } + fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> { + self.pinned_connection() + } + fn name(&self) -> &str { + self.name() + } +} diff --git a/src/operation/raw_output.rs b/src/operation/raw_output.rs index 0bc311e8b..020878fa0 100644 --- a/src/operation/raw_output.rs +++ b/src/operation/raw_output.rs @@ -1,5 +1,7 @@ -use crate::cmap::{StreamDescription, RawCommandResponse, Command}; -use crate::error::Result; +use crate::{ + cmap::{Command, RawCommandResponse, StreamDescription}, + error::Result, +}; use super::Operation; @@ -18,7 +20,10 @@ impl Operation for RawOutput { self.0.serialize_command(cmd) } - fn extract_at_cluster_time(&self, response: &bson::RawDocument) -> Result> { + fn extract_at_cluster_time( + &self, + response: &bson::RawDocument, + ) -> Result> { self.0.extract_at_cluster_time(response) } @@ -69,4 +74,4 @@ impl Operation for RawOutput { fn name(&self) -> &str { self.0.name() } -} \ No newline at end of file +} diff --git a/src/operation/run_command/mod.rs b/src/operation/run_command/mod.rs index fa33888fb..9906c9a51 100644 --- a/src/operation/run_command/mod.rs +++ b/src/operation/run_command/mod.rs @@ -67,7 +67,11 @@ impl<'conn> RunCommand<'conn> { } fn command_name(&self) -> Option<&str> { - self.command.into_iter().next().and_then(|r| r.ok()).map(|(k, _)| k) + self.command + .into_iter() + .next() + .and_then(|r| r.ok()) + .map(|(k, _)| k) } } From edf0924ecbd8a4889c551bb58c0e5f2c19fd56d3 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 4 Aug 2022 11:04:51 -0400 Subject: [PATCH 22/35] put crypto finalization on a thread pool --- Cargo.toml | 5 ++--- src/client/csfle/state_machine.rs | 37 ++++++++++++++++++++++++++----- src/client/executor.rs | 9 ++++---- 3 files changed, 39 insertions(+), 12 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e922d7328..a33be4d68 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,14 +34,12 @@ tokio-runtime = [ "tokio/rt", "tokio/time", "serde_bytes", - "mongocrypt/tokio", ] async-std-runtime = [ "async-std", "async-std/attributes", "async-std-resolver", "tokio-util/compat", - "mongocrypt/async-std", ] sync = ["async-std-runtime"] tokio-sync = ["tokio-runtime"] @@ -68,7 +66,7 @@ zlib-compression = ["flate2"] snappy-compression = ["snap"] # DO NOT USE; see https://jira.mongodb.org/browse/RUST-569 for the status of CSFLE support in the Rust driver. -csfle = ["mongocrypt", "which"] +csfle = ["mongocrypt", "which", "rayon"] [dependencies] async-trait = "0.1.42" @@ -92,6 +90,7 @@ openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } percent-encoding = "2.0.0" rand = { version = "0.8.3", features = ["small_rng"] } +rayon = { version = "1.5.3", optional = true } rustc_version_runtime = "0.2.1" rustls-pemfile = "0.3.0" serde_with = "1.3.1" diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index e7a99e461..c0ec74c75 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -3,7 +3,7 @@ use std::convert::TryInto; use bson::{Document, RawDocument, RawDocumentBuf}; use futures_util::{stream, TryStreamExt}; use mongocrypt::ctx::{Ctx, State}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::{io::{AsyncReadExt, AsyncWriteExt}, sync::oneshot}; use crate::{ client::options::{ServerAddress, TlsOptions}, @@ -17,7 +17,7 @@ use crate::{ impl Client { pub(crate) async fn run_mongocrypt_ctx( &self, - mut ctx: Ctx, + ctx: Ctx, db: Option<&str>, ) -> Result { let guard = self.inner.csfle.read().await; @@ -25,11 +25,17 @@ impl Client { Some(csfle) => csfle, None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; + // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing thread for crypto finalization. An `Option` would also work here, but `Result` means we can return a helpful error if things get into a broken state rather than panicing. let mut result = None; + let num_cpus = std::thread::available_parallelism()?.get(); + let crypt_threads = rayon::ThreadPoolBuilder::new().num_threads(num_cpus).build() + .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?; + let mut ctx = Ok(ctx); loop { - let state = ctx.state()?; + let state = result_ref(&ctx)?.state()?; match state { State::NeedMongoCollinfo => { + let ctx = result_mut(&mut ctx)?; let filter = raw_to_doc(ctx.mongo_op()?)?; let db = self.database(db.as_ref().ok_or_else(|| { Error::internal("db required for NeedMongoCollinfo state") @@ -41,6 +47,7 @@ impl Client { ctx.mongo_done()?; } State::NeedMongoMarkings => { + let ctx = result_mut(&mut ctx)?; let command = ctx.mongo_op()?.to_raw_document_buf(); let db = db.as_ref().ok_or_else(|| { Error::internal("db required for NeedMongoMarkings state") @@ -55,6 +62,7 @@ impl Client { ctx.mongo_done()?; } State::NeedMongoKeys => { + let ctx = result_mut(&mut ctx)?; let filter = raw_to_doc(ctx.mongo_op()?)?; let kv_ns = &csfle.opts.key_vault_namespace; let kv_client = csfle @@ -72,9 +80,10 @@ impl Client { ctx.mongo_done()?; } State::NeedKms => { + let ctx = result_mut(&mut ctx)?; let scope = ctx.kms_scope(); let mut kms_ctxen: Vec> = vec![]; - while let Some(kms_ctx) = scope.next_kms_ctx()? { + while let Some(kms_ctx) = scope.next_kms_ctx() { kms_ctxen.push(Ok(kms_ctx)); } stream::iter(kms_ctxen) @@ -103,7 +112,17 @@ impl Client { .await?; } State::NeedKmsCredentials => todo!("RUST-1314"), - State::Ready => result = Some(ctx.finalize()?.to_owned()), + State::Ready => { + let (tx, rx) = oneshot::channel(); + let mut thread_ctx = std::mem::replace(&mut ctx, Err(Error::internal("crypto context not present")))?; + crypt_threads.spawn(move || { + let result = thread_ctx.finalize().map(|doc| doc.to_owned()); + let _ = tx.send((thread_ctx, result)); + }); + let (ctx_again, output) = rx.await.map_err(|_| Error::internal("crypto thread dropped"))?; + ctx = Ok(ctx_again); + result = Some(output?); + } State::Done => break, s => return Err(Error::internal(format!("unhandled state {:?}", s))), } @@ -115,6 +134,14 @@ impl Client { } } +fn result_ref(r: &Result) -> Result<&T> { + r.as_ref().map_err(Error::clone) +} + +fn result_mut(r: &mut Result) -> Result<&mut T> { + r.as_mut().map_err(|e| e.clone()) +} + fn raw_to_doc(raw: &RawDocument) -> Result { raw.try_into() .map_err(|e| Error::internal(format!("could not parse raw document: {}", e))) diff --git a/src/client/executor.rs b/src/client/executor.rs index c22de17c6..bfbb772bc 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -797,11 +797,12 @@ impl Client { target_db: &'a str, ) -> BoxFuture<'a, Result> { Box::pin(async move { - let doc = RawDocument::from_bytes(serialized)?.to_raw_document_buf(); + let doc = RawDocument::from_bytes(serialized)?; let db = target_db.to_string(); let ctx = csfle .crypt - .build_ctx(move |builder| builder.build_encrypt(&db, &doc))?; + .ctx_builder() + .build_encrypt(&db, doc)?; self.run_mongocrypt_ctx(ctx, Some(target_db)).await }) } @@ -813,10 +814,10 @@ impl Client { doc: &'a RawDocument, ) -> BoxFuture<'a, Result> { Box::pin(async move { - let doc = doc.to_raw_document_buf(); let ctx = csfle .crypt - .build_ctx(move |builder| builder.build_decrypt(&doc))?; + .ctx_builder() + .build_decrypt(doc)?; self.run_mongocrypt_ctx(ctx, None).await }) } From bf78b11e0125d4089a620f01d82268e00bd11c59 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 4 Aug 2022 11:46:24 -0400 Subject: [PATCH 23/35] format and clippy --- Cargo.toml | 2 +- src/client/csfle/state_machine.rs | 22 +++++++++++---- src/client/executor.rs | 45 +++++++++++++++++-------------- src/operation/mod.rs | 1 + src/operation/run_command/mod.rs | 1 + 5 files changed, 45 insertions(+), 26 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index a33be4d68..2d086fb2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,7 +84,7 @@ hmac = "0.12.1" lazy_static = "1.4.0" md-5 = "0.10.1" # mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } -mongocrypt = { path = "../libmongocrypt-rust/mongocrypt", optional = true } +mongocrypt = { git = "https://github.com/abr-egn/libmongocrypt-rust.git", branch = "RUST-1384/send-ctx", optional = true } openssl = { version = "0.10.38", optional = true } openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index c0ec74c75..7497d369d 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -3,7 +3,10 @@ use std::convert::TryInto; use bson::{Document, RawDocument, RawDocumentBuf}; use futures_util::{stream, TryStreamExt}; use mongocrypt::ctx::{Ctx, State}; -use tokio::{io::{AsyncReadExt, AsyncWriteExt}, sync::oneshot}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::oneshot, +}; use crate::{ client::options::{ServerAddress, TlsOptions}, @@ -25,10 +28,14 @@ impl Client { Some(csfle) => csfle, None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; - // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing thread for crypto finalization. An `Option` would also work here, but `Result` means we can return a helpful error if things get into a broken state rather than panicing. + // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing + // thread for crypto finalization. An `Option` would also work here, but `Result` means we + // can return a helpful error if things get into a broken state rather than panicing. let mut result = None; let num_cpus = std::thread::available_parallelism()?.get(); - let crypt_threads = rayon::ThreadPoolBuilder::new().num_threads(num_cpus).build() + let crypt_threads = rayon::ThreadPoolBuilder::new() + .num_threads(num_cpus) + .build() .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?; let mut ctx = Ok(ctx); loop { @@ -114,12 +121,17 @@ impl Client { State::NeedKmsCredentials => todo!("RUST-1314"), State::Ready => { let (tx, rx) = oneshot::channel(); - let mut thread_ctx = std::mem::replace(&mut ctx, Err(Error::internal("crypto context not present")))?; + let mut thread_ctx = std::mem::replace( + &mut ctx, + Err(Error::internal("crypto context not present")), + )?; crypt_threads.spawn(move || { let result = thread_ctx.finalize().map(|doc| doc.to_owned()); let _ = tx.send((thread_ctx, result)); }); - let (ctx_again, output) = rx.await.map_err(|_| Error::internal("crypto thread dropped"))?; + let (ctx_again, output) = rx + .await + .map_err(|_| Error::internal("crypto thread dropped"))?; ctx = Ok(ctx_again); result = Some(output?); } diff --git a/src/client/executor.rs b/src/client/executor.rs index bfbb772bc..6f08ac9ae 100644 --- a/src/client/executor.rs +++ b/src/client/executor.rs @@ -1,4 +1,7 @@ -use bson::{doc, RawBsonRef, RawDocument, RawDocumentBuf, Timestamp}; +#[cfg(feature = "csfle")] +use bson::RawDocumentBuf; +use bson::{doc, RawBsonRef, RawDocument, Timestamp}; +#[cfg(feature = "csfle")] use futures_core::future::BoxFuture; use lazy_static::lazy_static; use serde::de::DeserializeOwned; @@ -598,19 +601,22 @@ impl Client { let cmd_name = cmd.name.clone(); let target_db = cmd.target_db.clone(); - let mut serialized = op.serialize_command(cmd)?; + let serialized = op.serialize_command(cmd)?; #[cfg(feature = "csfle")] - { + let serialized = { let guard = self.inner.csfle.read().await; if let Some(ref csfle) = *guard { if csfle.opts().bypass_auto_encryption != Some(true) { - serialized = self - .auto_encrypt(csfle, &serialized, &target_db) + self.auto_encrypt(csfle, RawDocument::from_bytes(&serialized)?, &target_db) .await? - .into_bytes(); + .into_bytes() + } else { + serialized } + } else { + serialized } - } + }; let raw_cmd = RawCommand { name: cmd_name.clone(), target_db, @@ -742,7 +748,7 @@ impl Client { err.add_labels_and_update_pin(Some(connection), session, Some(retryability))?; op.handle_error(err) } - Ok(mut response) => { + Ok(response) => { self.emit_command_event(|handler| { let reply = if should_redact { Document::new() @@ -764,15 +770,19 @@ impl Client { }); #[cfg(feature = "csfle")] - { + let response = { let guard = self.inner.csfle.read().await; if let Some(ref csfle) = *guard { if csfle.opts().bypass_auto_encryption != Some(true) { let new_body = self.auto_decrypt(csfle, response.raw_body()).await?; - response = RawCommandResponse::new_raw(response.source, new_body); + RawCommandResponse::new_raw(response.source, new_body) + } else { + response } + } else { + response } - } + }; match op.handle_response(response, connection.stream_description()?) { Ok(response) => Ok(response), @@ -793,16 +803,14 @@ impl Client { fn auto_encrypt<'a>( &'a self, csfle: &'a super::csfle::ClientState, - serialized: &'a [u8], + command: &'a RawDocument, target_db: &'a str, ) -> BoxFuture<'a, Result> { Box::pin(async move { - let doc = RawDocument::from_bytes(serialized)?; - let db = target_db.to_string(); let ctx = csfle .crypt .ctx_builder() - .build_encrypt(&db, doc)?; + .build_encrypt(target_db, command)?; self.run_mongocrypt_ctx(ctx, Some(target_db)).await }) } @@ -811,13 +819,10 @@ impl Client { fn auto_decrypt<'a>( &'a self, csfle: &'a super::csfle::ClientState, - doc: &'a RawDocument, + response: &'a RawDocument, ) -> BoxFuture<'a, Result> { Box::pin(async move { - let ctx = csfle - .crypt - .ctx_builder() - .build_decrypt(doc)?; + let ctx = csfle.crypt.ctx_builder().build_decrypt(response)?; self.run_mongocrypt_ctx(ctx, None).await }) } diff --git a/src/operation/mod.rs b/src/operation/mod.rs index dae5e014f..f3b599f6e 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -68,6 +68,7 @@ pub(crate) use insert::Insert; pub(crate) use list_collections::ListCollections; pub(crate) use list_databases::ListDatabases; pub(crate) use list_indexes::ListIndexes; +#[cfg(feature = "csfle")] pub(crate) use raw_output::RawOutput; pub(crate) use run_command::RunCommand; pub(crate) use update::Update; diff --git a/src/operation/run_command/mod.rs b/src/operation/run_command/mod.rs index 9906c9a51..c8b784a85 100644 --- a/src/operation/run_command/mod.rs +++ b/src/operation/run_command/mod.rs @@ -45,6 +45,7 @@ impl<'conn> RunCommand<'conn> { }) } + #[cfg(feature = "csfle")] pub(crate) fn new_raw( db: String, command: RawDocumentBuf, From 15006f1a0f7411c064cdfc64c01a0c015ced65ee Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 4 Aug 2022 11:48:00 -0400 Subject: [PATCH 24/35] fix comment --- src/client/csfle/state_machine.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index 7497d369d..ee3a320f2 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -28,15 +28,15 @@ impl Client { Some(csfle) => csfle, None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; - // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing - // thread for crypto finalization. An `Option` would also work here, but `Result` means we - // can return a helpful error if things get into a broken state rather than panicing. let mut result = None; let num_cpus = std::thread::available_parallelism()?.get(); let crypt_threads = rayon::ThreadPoolBuilder::new() .num_threads(num_cpus) .build() .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?; + // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing + // thread for crypto finalization. An `Option` would also work here, but `Result` means we + // can return a helpful error if things get into a broken state rather than panicing. let mut ctx = Ok(ctx); loop { let state = result_ref(&ctx)?.state()?; From e1b9ad829a807560549f1d288786e5b885bdade3 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 4 Aug 2022 15:21:48 -0400 Subject: [PATCH 25/35] point back to main mongocrypt repo --- Cargo.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2d086fb2a..e82b59047 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,8 +83,7 @@ hex = "0.4.0" hmac = "0.12.1" lazy_static = "1.4.0" md-5 = "0.10.1" -# mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } -mongocrypt = { git = "https://github.com/abr-egn/libmongocrypt-rust.git", branch = "RUST-1384/send-ctx", optional = true } +mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } openssl = { version = "0.10.38", optional = true } openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } From 293cfb9fa44933de522fa7f555ef681a87f1d98e Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 8 Aug 2022 11:17:31 -0400 Subject: [PATCH 26/35] move thread pool to client state --- src/client/csfle.rs | 8 ++++++++ src/client/csfle/state_machine.rs | 7 +------ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/client/csfle.rs b/src/client/csfle.rs index d6c902793..613ca3e8a 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -8,6 +8,7 @@ use std::{ use derivative::Derivative; use mongocrypt::Crypt; +use rayon::ThreadPool; use crate::{ error::{Error, Result}, @@ -36,6 +37,7 @@ pub(super) struct ClientState { mongocryptd_client: Option, aux_clients: AuxClients, opts: AutoEncryptionOptions, + crypto_threads: ThreadPool, } #[derive(Debug)] @@ -53,12 +55,18 @@ impl ClientState { let crypt = Self::make_crypt(&opts)?; let mongocryptd_client = Self::spawn_mongocryptd_if_needed(&opts, &crypt).await?; let aux_clients = Self::make_aux_clients(client, &opts)?; + let num_cpus = std::thread::available_parallelism()?.get(); + let crypto_threads = rayon::ThreadPoolBuilder::new() + .num_threads(num_cpus) + .build() + .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?; Ok(Self { crypt, mongocryptd_client, aux_clients, opts, + crypto_threads, }) } diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index ee3a320f2..ee9d60a2c 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -29,11 +29,6 @@ impl Client { None => return Err(Error::internal("no csfle state for mongocrypt ctx")), }; let mut result = None; - let num_cpus = std::thread::available_parallelism()?.get(); - let crypt_threads = rayon::ThreadPoolBuilder::new() - .num_threads(num_cpus) - .build() - .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?; // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing // thread for crypto finalization. An `Option` would also work here, but `Result` means we // can return a helpful error if things get into a broken state rather than panicing. @@ -125,7 +120,7 @@ impl Client { &mut ctx, Err(Error::internal("crypto context not present")), )?; - crypt_threads.spawn(move || { + csfle.crypto_threads.spawn(move || { let result = thread_ctx.finalize().map(|doc| doc.to_owned()); let _ = tx.send((thread_ctx, result)); }); From d7acff990b6561d589d1806d1def12171077a336 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 10 Aug 2022 12:29:08 -0400 Subject: [PATCH 27/35] review updates --- src/client/csfle.rs | 8 ++------ src/client/csfle/state_machine.rs | 10 +++++++++- src/coll/mod.rs | 7 ++++++- src/operation/insert/mod.rs | 10 +++++++++- 4 files changed, 26 insertions(+), 9 deletions(-) diff --git a/src/client/csfle.rs b/src/client/csfle.rs index 613ca3e8a..7db1c87b2 100644 --- a/src/client/csfle.rs +++ b/src/client/csfle.rs @@ -32,7 +32,6 @@ use super::WeakClient; #[derivative(Debug)] pub(super) struct ClientState { #[derivative(Debug = "ignore")] - #[allow(dead_code)] pub(crate) crypt: Crypt, mongocryptd_client: Option, aux_clients: AuxClients, @@ -42,12 +41,9 @@ pub(super) struct ClientState { #[derive(Debug)] struct AuxClients { - #[allow(dead_code)] key_vault_client: WeakClient, - #[allow(dead_code)] metadata_client: Option, - #[allow(dead_code)] - internal_client: Option, + _internal_client: Option, } impl ClientState { @@ -180,7 +176,7 @@ impl ClientState { Ok(AuxClients { key_vault_client, metadata_client, - internal_client, + _internal_client: internal_client, }) } } diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index ee9d60a2c..ed8f677f7 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -39,7 +39,15 @@ impl Client { State::NeedMongoCollinfo => { let ctx = result_mut(&mut ctx)?; let filter = raw_to_doc(ctx.mongo_op()?)?; - let db = self.database(db.as_ref().ok_or_else(|| { + let metadata_client = csfle + .aux_clients + .metadata_client + .as_ref() + .and_then(|w| w.upgrade()) + .ok_or_else(|| { + Error::internal("metadata_client required for NeedMongoCollinfo state") + })?; + let db = metadata_client.database(db.as_ref().ok_or_else(|| { Error::internal("db required for NeedMongoCollinfo state") })?); let mut cursor = db.list_collections(filter, None).await?; diff --git a/src/coll/mod.rs b/src/coll/mod.rs index 6916a5d84..a50f8ae55 100644 --- a/src/coll/mod.rs +++ b/src/coll/mod.rs @@ -1189,6 +1189,10 @@ where } let ordered = options.as_ref().and_then(|o| o.ordered).unwrap_or(true); + #[cfg(feature = "csfle")] + let encrypted = self.client().auto_encryption_opts().await.is_some(); + #[cfg(not(feature = "csfle"))] + let encrypted = false; let mut cumulative_failure: Option = None; let mut error_labels: HashSet = Default::default(); @@ -1198,7 +1202,7 @@ where while n_attempted < ds.len() { let docs: Vec<&T> = ds.iter().skip(n_attempted).map(Borrow::borrow).collect(); - let insert = Insert::new(self.namespace(), docs, options.clone()); + let insert = Insert::new(self.namespace(), docs, options.clone(), encrypted); match self .client() @@ -1325,6 +1329,7 @@ where self.namespace(), vec![doc], options.map(InsertManyOptions::from_insert_one_options), + false, ); self.client() .execute_operation(insert, session) diff --git a/src/operation/insert/mod.rs b/src/operation/insert/mod.rs index 7f6a48ea4..fd2aebb7f 100644 --- a/src/operation/insert/mod.rs +++ b/src/operation/insert/mod.rs @@ -30,6 +30,7 @@ pub(crate) struct Insert<'a, T> { documents: Vec<&'a T>, inserted_ids: Vec, options: Option, + encrypted: bool, } impl<'a, T> Insert<'a, T> { @@ -37,12 +38,14 @@ impl<'a, T> Insert<'a, T> { ns: Namespace, documents: Vec<&'a T>, options: Option, + encrypted: bool, ) -> Self { Self { ns, options, documents, inserted_ids: vec![], + encrypted, } } @@ -63,6 +66,11 @@ impl<'a, T: Serialize> OperationWithDefaults for Insert<'a, T> { fn build(&mut self, description: &StreamDescription) -> Result> { let mut docs = RawArrayBuf::new(); let mut size = 0; + let batch_size_limit = if self.encrypted { + 2_097_152 + } else { + description.max_bson_object_size as u64 + }; for (i, d) in self .documents @@ -97,7 +105,7 @@ impl<'a, T: Serialize> OperationWithDefaults for Insert<'a, T> { let doc_size = bson_util::array_entry_size_bytes(i, doc.as_bytes().len()); - if (size + doc_size) <= description.max_bson_object_size as u64 { + if (size + doc_size) <= batch_size_limit { if self.inserted_ids.len() <= i { self.inserted_ids.push(id); } From 6bbd8c1fa816151d407eba43c7f63712de0bd70d Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Thu, 11 Aug 2022 10:42:05 -0400 Subject: [PATCH 28/35] fix tests and include csfle in feature testing --- .evergreen/feature-combinations.sh | 2 +- src/coll/mod.rs | 3 +-- src/operation/insert/mod.rs | 8 ++++++++ 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/.evergreen/feature-combinations.sh b/.evergreen/feature-combinations.sh index 92fcc9eb4..20dc50e2f 100755 --- a/.evergreen/feature-combinations.sh +++ b/.evergreen/feature-combinations.sh @@ -6,5 +6,5 @@ export FEATURE_COMBINATIONS=( '' # default features '--no-default-features --features async-std-runtime,sync' # features that conflict w/ default features - '--features tokio-sync,zstd-compression,snappy-compression,zlib-compression,openssl-tls,aws-auth' # additive features + '--features tokio-sync,zstd-compression,snappy-compression,zlib-compression,openssl-tls,aws-auth,csfle' # additive features ) diff --git a/src/coll/mod.rs b/src/coll/mod.rs index a50f8ae55..ed9bfee0f 100644 --- a/src/coll/mod.rs +++ b/src/coll/mod.rs @@ -1202,7 +1202,7 @@ where while n_attempted < ds.len() { let docs: Vec<&T> = ds.iter().skip(n_attempted).map(Borrow::borrow).collect(); - let insert = Insert::new(self.namespace(), docs, options.clone(), encrypted); + let insert = Insert::new_encrypted(self.namespace(), docs, options.clone(), encrypted); match self .client() @@ -1329,7 +1329,6 @@ where self.namespace(), vec![doc], options.map(InsertManyOptions::from_insert_one_options), - false, ); self.client() .execute_operation(insert, session) diff --git a/src/operation/insert/mod.rs b/src/operation/insert/mod.rs index fd2aebb7f..c7d36dc12 100644 --- a/src/operation/insert/mod.rs +++ b/src/operation/insert/mod.rs @@ -38,6 +38,14 @@ impl<'a, T> Insert<'a, T> { ns: Namespace, documents: Vec<&'a T>, options: Option, + ) -> Self { + Self::new_encrypted(ns, documents, options, false) + } + + pub(crate) fn new_encrypted( + ns: Namespace, + documents: Vec<&'a T>, + options: Option, encrypted: bool, ) -> Self { Self { From 859967a3fde710dcc23a3acfd2b83f50cdeb9f3a Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Fri, 12 Aug 2022 14:25:35 -0400 Subject: [PATCH 29/35] actually use tls_options --- Cargo.toml | 3 ++- src/client/csfle/state_machine.rs | 12 ++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e82b59047..16f05ec77 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,8 @@ hex = "0.4.0" hmac = "0.12.1" lazy_static = "1.4.0" md-5 = "0.10.1" -mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } +#mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } +mongocrypt = { path = "../libmongocrypt-rust/mongocrypt", optional = true } openssl = { version = "0.10.38", optional = true } openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index ed8f677f7..cc22d2708 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -9,7 +9,7 @@ use tokio::{ }; use crate::{ - client::options::{ServerAddress, TlsOptions}, + client::options::ServerAddress, cmap::options::StreamOptions, error::{Error, Result}, operation::{RawOutput, RunCommand}, @@ -100,10 +100,18 @@ impl Client { .try_for_each_concurrent(None, |mut kms_ctx| async move { let endpoint = kms_ctx.endpoint()?; let addr = ServerAddress::parse(endpoint)?; + let provider = kms_ctx.kms_provider()?; + let tls_options = csfle + .opts() + .tls_options + .as_ref() + .and_then(|tls| tls.get(&provider)) + .cloned() + .unwrap_or_default(); let mut stream = AsyncStream::connect( StreamOptions::builder() .address(addr) - .tls_options(TlsOptions::default()) + .tls_options(tls_options) .build(), ) .await?; From 0dedc5b770d24aa4058708951739c56e29acf610 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Fri, 12 Aug 2022 14:40:42 -0400 Subject: [PATCH 30/35] fix mongocrypt dep --- Cargo.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 16f05ec77..e82b59047 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,8 +83,7 @@ hex = "0.4.0" hmac = "0.12.1" lazy_static = "1.4.0" md-5 = "0.10.1" -#mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } -mongocrypt = { path = "../libmongocrypt-rust/mongocrypt", optional = true } +mongocrypt = { git = "https://github.com/mongodb/libmongocrypt-rust.git", branch = "main", optional = true } openssl = { version = "0.10.38", optional = true } openssl-probe = { version = "0.1.5", optional = true } os_info = { version = "3.0.1", default-features = false } From 8169c6aa40716c15790d37ce33aaf53f75ceb2fc Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 15 Aug 2022 11:00:25 -0400 Subject: [PATCH 31/35] use advance/current --- src/client/csfle/state_machine.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/client/csfle/state_machine.rs b/src/client/csfle/state_machine.rs index cc22d2708..1532b1ad7 100644 --- a/src/client/csfle/state_machine.rs +++ b/src/client/csfle/state_machine.rs @@ -84,8 +84,8 @@ impl Client { .database(&kv_ns.db) .collection::(&kv_ns.coll); let mut cursor = kv_coll.find(filter, None).await?; - while let Some(result) = cursor.try_next().await? { - ctx.mongo_feed(&result)? + while cursor.advance().await? { + ctx.mongo_feed(cursor.current())?; } ctx.mongo_done()?; } From 40cffb99c9d80cead3cdad6525467eb21f902a7d Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Mon, 15 Aug 2022 11:06:45 -0400 Subject: [PATCH 32/35] comment rawoutput --- src/operation/raw_output.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/operation/raw_output.rs b/src/operation/raw_output.rs index 020878fa0..d7b2b9534 100644 --- a/src/operation/raw_output.rs +++ b/src/operation/raw_output.rs @@ -5,6 +5,8 @@ use crate::{ use super::Operation; +/// Forwards all implementation to the wrapped `Operation`, but returns the response unparsed and +/// unvalidated as a `RawCommandResponse`. pub(crate) struct RawOutput(pub(crate) Op); impl Operation for RawOutput { From d27039396e4557610237a1002a1c401f3a193807 Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 17 Aug 2022 11:38:04 -0400 Subject: [PATCH 33/35] review updates --- src/operation/mod.rs | 2 +- src/operation/run_command/test.rs | 19 ------------------- 2 files changed, 1 insertion(+), 20 deletions(-) diff --git a/src/operation/mod.rs b/src/operation/mod.rs index f3b599f6e..414f71370 100644 --- a/src/operation/mod.rs +++ b/src/operation/mod.rs @@ -393,7 +393,7 @@ macro_rules! remove_empty_write_concern { pub(crate) use remove_empty_write_concern; // A mirror of the `Operation` trait, with default behavior where appropriate. Should only be -// implemented by leaf operation types. +// implemented by operation types that do not delegate to other operations. pub(crate) trait OperationWithDefaults { /// The output type of this operation. type O; diff --git a/src/operation/run_command/test.rs b/src/operation/run_command/test.rs index ad2b89a44..593f03ba4 100644 --- a/src/operation/run_command/test.rs +++ b/src/operation/run_command/test.rs @@ -9,25 +9,6 @@ use crate::{ operation::{test::handle_response_test, Operation}, }; -#[test] -fn build() { - let mut op = RunCommand::new("foo".into(), doc! { "hello": 1 }, None, None).unwrap(); - assert!(op.selection_criteria().is_none()); - - let command = op.build(&StreamDescription::new_testing()).unwrap(); - - assert_eq!(command.name, "hello"); - assert_eq!(command.target_db, "foo"); - assert_eq!( - command - .body - .get("hello") - .unwrap() - .and_then(|raw| crate::bson_util::get_int(&raw.try_into().unwrap())), - Some(1) - ); -} - #[test] fn handle_success() { let op = RunCommand::new("foo".into(), doc! { "hello": 1 }, None, None).unwrap(); From 614bd789e32d1a81d9ae6c30ba5351b43d6a00ec Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 17 Aug 2022 11:40:07 -0400 Subject: [PATCH 34/35] fix test imports --- src/operation/run_command/test.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/operation/run_command/test.rs b/src/operation/run_command/test.rs index 593f03ba4..d7e959961 100644 --- a/src/operation/run_command/test.rs +++ b/src/operation/run_command/test.rs @@ -1,12 +1,9 @@ -use std::convert::TryInto; - use bson::Timestamp; use super::RunCommand; use crate::{ bson::doc, - cmap::StreamDescription, - operation::{test::handle_response_test, Operation}, + operation::test::handle_response_test, }; #[test] From 0b007b51ea564dd0ad98ee485fa78c4442a265cf Mon Sep 17 00:00:00 2001 From: Abraham Egnor Date: Wed, 17 Aug 2022 13:51:30 -0400 Subject: [PATCH 35/35] fix rustfmt --- src/operation/run_command/test.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/operation/run_command/test.rs b/src/operation/run_command/test.rs index d7e959961..467ccf5e1 100644 --- a/src/operation/run_command/test.rs +++ b/src/operation/run_command/test.rs @@ -1,10 +1,7 @@ use bson::Timestamp; use super::RunCommand; -use crate::{ - bson::doc, - operation::test::handle_response_test, -}; +use crate::{bson::doc, operation::test::handle_response_test}; #[test] fn handle_success() {