Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add flag to encrypt_util for length delimited binary format #1509

Merged
merged 1 commit into from
Dec 20, 2024
Merged

add flag to encrypt_util for length delimited binary format #1509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions ipa-core/src/cli/crypto/hybrid_decrypt.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,9 +226,14 @@ mod tests {

let output_dir = tempdir().unwrap();
let network_file = hybrid_sample_data::test_keys().network_config();
HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir.path(),
network_file.path(),
false,
)
.encrypt()
.unwrap();

let decrypt_output = output_dir.path().join("output");
let enc1 = output_dir.path().join("DOES_NOT_EXIST.enc");
Expand Down Expand Up @@ -258,9 +263,14 @@ mod tests {

let network_file = hybrid_sample_data::test_keys().network_config();
let output_dir = tempdir().unwrap();
HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir.path(),
network_file.path(),
false,
)
.encrypt()
.unwrap();

let decrypt_output = output_dir.path().join("output");
let enc1 = output_dir.path().join("helper1.enc");
Expand Down
148 changes: 120 additions & 28 deletions ipa-core/src/cli/crypto/hybrid_encrypt.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,15 +52,38 @@
/// Path to helper network configuration file
#[arg(long)]
network: PathBuf,
/// a flag to produce length delimited binary instead of newline delimited hex
#[arg(long)]
length_delimited: bool,

Check warning on line 57 in ipa-core/src/cli/crypto/hybrid_encrypt.rs

View check run for this annotation

Codecov / codecov/patch

ipa-core/src/cli/crypto/hybrid_encrypt.rs#L57 

Added line #L57 was not covered by tests
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of a boolean argument it might be clearer to have descriptive options (something like the FileFormat choices -- possibly clap can even be taught to parse into FileFormat).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I generally think of docopt as the canonical way to design CLIs, and this would probably fit as a mutually exclusive element better than an element with an arg. Clap probably has a way of defining mutually exclusive elements (and as you say, clap could likely be taught to coerce that into the enum.) This would be even more important if there were more than two options.

I likely won't get to this before I'm out, so up to @andyleiserson and @akoshelev if you want to refactor before merging.

}

#[derive(Copy, Clone)]
enum FileFormat {
LengthDelimitedBinary,
NewlineDelimitedHex,
}

impl HybridEncryptArgs {
#[must_use]
pub fn new(input_file: &Path, output_dir: &Path, network: &Path) -> Self {
pub fn new(
input_file: &Path,
output_dir: &Path,
network: &Path,
length_delimited: bool,
) -> Self {
Self {
input_file: input_file.to_path_buf(),
output_dir: output_dir.to_path_buf(),
network: network.to_path_buf(),
length_delimited,
}
}

fn file_format(&self) -> FileFormat {
if self.length_delimited {
FileFormat::LengthDelimitedBinary
} else {
FileFormat::NewlineDelimitedHex
}
}

Expand Down Expand Up @@ -89,7 +112,8 @@
panic!("could not load network file")
};

let mut worker_pool = ReportWriter::new(key_registries, &self.output_dir);
let mut worker_pool =
ReportWriter::new(key_registries, &self.output_dir, self.file_format());
for (report_id, record) in input.iter::<TestHybridRecord>().enumerate() {
worker_pool.submit(report_id, record.share())?;
}
Expand Down Expand Up @@ -118,6 +142,7 @@
thread_count: usize,
file_writer: [SyncSender<EncryptorOutput>; 3],
key_registries: [KeyRegistry<PublicKeyOnly>; 3],
file_format: FileFormat,
) -> Self {
Self {
pool: (0..thread_count)
Expand All @@ -132,11 +157,23 @@
.spawn(move || {
for (i, helper_id, report) in rx {
let key_registry = &key_registries[helper_id];
let output = report.encrypt(
DEFAULT_KEY_ID,
key_registry,
&mut thread_rng(),
)?;
let mut output =
Vec::with_capacity(usize::from(report.encrypted_len() + 2));
match file_format {
FileFormat::NewlineDelimitedHex => report.encrypt_to(
DEFAULT_KEY_ID,
key_registry,
&mut thread_rng(),
&mut output,
)?,
FileFormat::LengthDelimitedBinary => report
.delimited_encrypt_to(
DEFAULT_KEY_ID,
key_registry,
&mut thread_rng(),
&mut output,
)?,
}
file_writer[helper_id].send((i, output))?;
}

Expand Down Expand Up @@ -178,7 +215,11 @@
}

impl ReportWriter {
pub fn new(key_registries: [KeyRegistry<PublicKeyOnly>; 3], output_dir: &Path) -> Self {
pub fn new(
key_registries: [KeyRegistry<PublicKeyOnly>; 3],
output_dir: &Path,
file_format: FileFormat,
) -> Self {
// create 3 worker threads to write data into 3 files
let workers = array::from_fn(|i| {
let output_filename = format!("helper{}.enc", i + 1);
Expand All @@ -188,12 +229,13 @@
.open(output_dir.join(&output_filename))
.unwrap_or_else(|e| panic!("unable write to {:?}. {}", &output_filename, e));

FileWriteWorker::new(file)
FileWriteWorker::new(file, file_format)
});
let encryptor_pool = EncryptorPool::with_worker_threads(
num_cpus::get(),
workers.each_ref().map(|x| x.sender.clone()),
key_registries,
file_format,
);

Self {
Expand Down Expand Up @@ -239,17 +281,26 @@
}

impl FileWriteWorker {
pub fn new(file: File) -> Self {
pub fn new(file: File, file_format: FileFormat) -> Self {
fn write_report<W: Write>(
writer: &mut W,
report: &[u8],
file_format: FileFormat,
) -> Result<(), BoxError> {
match file_format {
FileFormat::LengthDelimitedBinary => {
FileWriteWorker::write_report_length_delimited_binary(writer, report)
}
FileFormat::NewlineDelimitedHex => {
FileWriteWorker::write_report_newline_delimited_hex(writer, report)
}
}
}

let (tx, rx) = std::sync::mpsc::sync_channel(65535);
Self {
sender: tx,
handle: thread::spawn(move || {
fn write_report<W: Write>(writer: &mut W, report: &[u8]) -> Result<(), BoxError> {
let hex_output = hex::encode(report);
writeln!(writer, "{hex_output}")?;
Ok(())
}

// write low watermark. All reports below this line have been written
let mut lw = 0;
let mut pending_reports = BTreeMap::new();
Expand All @@ -271,7 +322,7 @@
"Internal error: received a duplicate report {report_id}"
);
while let Some(report) = pending_reports.remove(&lw) {
write_report(&mut writer, &report)?;
write_report(&mut writer, &report, file_format)?;
lw += 1;
if lw % 1_000_000 == 0 {
tracing::info!("Encrypted {}M reports", lw / 1_000_000);
Expand All @@ -282,6 +333,23 @@
}),
}
}

fn write_report_newline_delimited_hex<W: Write>(
writer: &mut W,
report: &[u8],
) -> Result<(), BoxError> {
let hex_output = hex::encode(report);
writeln!(writer, "{hex_output}")?;
Ok(())
}

fn write_report_length_delimited_binary<W: Write>(
writer: &mut W,
report: &[u8],
) -> Result<(), BoxError> {
writer.write_all(report)?;
Ok(())
}
}

#[cfg(all(test, unit_test))]
Expand Down Expand Up @@ -334,12 +402,26 @@
}
input_file.flush().unwrap();

let output_dir = tempdir().unwrap();
let output_dir_1 = tempdir().unwrap();
let output_dir_2 = tempdir().unwrap();
let network_file = sample_data::test_keys().network_config();

HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir_1.path(),
network_file.path(),
false,
)
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir_2.path(),
network_file.path(),
true,
)
.encrypt()
.unwrap();
}

#[test]
Expand All @@ -350,7 +432,7 @@
let output_dir = tempdir().unwrap();
let network_dir = tempdir().unwrap();
let network_file = network_dir.path().join("does_not_exist");
HybridEncryptArgs::new(input_file.path(), output_dir.path(), &network_file)
HybridEncryptArgs::new(input_file.path(), output_dir.path(), &network_file, true)
.encrypt()
.unwrap();
}
Expand All @@ -368,9 +450,14 @@
let mut network_file = NamedTempFile::new().unwrap();
writeln!(network_file.as_file_mut(), "{network_data}").unwrap();

HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir.path(),
network_file.path(),
true,
)
.encrypt()
.unwrap();
}

#[test]
Expand All @@ -392,8 +479,13 @@
let mut network_file = NamedTempFile::new().unwrap();
writeln!(network_file.as_file_mut(), "{network_data}").unwrap();

HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir.path(),
network_file.path(),
true,
)
.encrypt()
.unwrap();
}
}
11 changes: 8 additions & 3 deletions ipa-core/src/cli/crypto/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,9 +345,14 @@ mod tests {
let input = hybrid_sample_data::test_hybrid_data().take(10);
let input_file = hybrid_sample_data::write_csv(input).unwrap();
let network_file = hybrid_sample_data::test_keys().network_config();
HybridEncryptArgs::new(input_file.path(), output_dir.path(), network_file.path())
.encrypt()
.unwrap();
HybridEncryptArgs::new(
input_file.path(),
output_dir.path(),
network_file.path(),
false,
)
.encrypt()
.unwrap();

let decrypt_output = output_dir.path().join("output");
let enc1 = output_dir.path().join("helper1.enc");
Expand Down
Loading