adamreeve commented on code in PR #7286:
URL: https://github.com/apache/arrow-rs/pull/7286#discussion_r2076762353
##########
parquet/src/encryption/decrypt.rs:
##########
@@ -136,49 +167,158 @@ impl CryptoContext {
}
}
-/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
-#[derive(Debug, Clone, PartialEq)]
-pub struct FileDecryptionProperties {
+#[derive(Clone, PartialEq)]
+struct ExplicitDecryptionKeys {
footer_key: Vec<u8>,
column_keys: HashMap<String, Vec<u8>>,
+}
+
+#[derive(Clone)]
+enum DecryptionKeys {
+ Explicit(ExplicitDecryptionKeys),
+ ViaRetriever(Arc<dyn KeyRetriever>),
+}
+
+impl PartialEq for DecryptionKeys {
+ fn eq(&self, other: &Self) -> bool {
+ match (self, other) {
+ (DecryptionKeys::Explicit(keys),
DecryptionKeys::Explicit(other_keys)) => {
+ keys.footer_key == other_keys.footer_key
+ && keys.column_keys == other_keys.column_keys
+ }
+ (DecryptionKeys::ViaRetriever(_), DecryptionKeys::ViaRetriever(_))
=> true,
+ _ => false,
+ }
+ }
+}
+
+/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
+#[derive(Clone, PartialEq)]
+pub struct FileDecryptionProperties {
+ keys: DecryptionKeys,
pub(crate) aad_prefix: Option<Vec<u8>>,
}
impl FileDecryptionProperties {
- /// Returns a new FileDecryptionProperties builder
+ /// Returns a new [`FileDecryptionProperties`] builder that will use the
provided key to
+ /// decrypt footer metadata.
pub fn builder(footer_key: Vec<u8>) -> DecryptionPropertiesBuilder {
DecryptionPropertiesBuilder::new(footer_key)
}
+
+ /// Returns a new [`FileDecryptionProperties`] builder that uses a
[`KeyRetriever`]
+ /// to get decryption keys based on key metadata.
+ pub fn with_key_retriever(key_retriever: Arc<dyn KeyRetriever>) ->
DecryptionPropertiesBuilder {
+ DecryptionPropertiesBuilder::new_with_key_retriever(key_retriever)
+ }
+
+ /// Get the encryption key for decrypting a file's footer,
+ /// and also column data if uniform encryption is used.
+ pub(crate) fn footer_key(&self, key_metadata: Option<&[u8]>) ->
Result<Cow<Vec<u8>>> {
+ match &self.keys {
+ DecryptionKeys::Explicit(keys) =>
Ok(Cow::Borrowed(&keys.footer_key)),
+ DecryptionKeys::ViaRetriever(retriever) => {
+ let key =
retriever.retrieve_key(key_metadata.unwrap_or_default())?;
+ Ok(Cow::Owned(key))
+ }
+ }
+ }
+
+ /// Get the column-specific encryption key for decrypting column data and
metadata within a file
+ pub(crate) fn column_key(
+ &self,
+ column_name: &str,
+ key_metadata: Option<&[u8]>,
+ ) -> Result<Cow<Vec<u8>>> {
+ match &self.keys {
+ DecryptionKeys::Explicit(keys) => match
keys.column_keys.get(column_name) {
+ None => Err(general_err!(
+ "No column decryption key set for column '{}'",
+ column_name
+ )),
+ Some(key) => Ok(Cow::Borrowed(key)),
+ },
+ DecryptionKeys::ViaRetriever(retriever) => {
+ let key =
retriever.retrieve_key(key_metadata.unwrap_or_default())?;
+ Ok(Cow::Owned(key))
+ }
+ }
+ }
+}
+
+impl std::fmt::Debug for FileDecryptionProperties {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "FileDecryptionProperties {{ }}")
+ }
}
+/// Builder for [`FileDecryptionProperties`]
pub struct DecryptionPropertiesBuilder {
- footer_key: Vec<u8>,
+ footer_key: Option<Vec<u8>>,
+ key_retriever: Option<Arc<dyn KeyRetriever>>,
column_keys: HashMap<String, Vec<u8>>,
aad_prefix: Option<Vec<u8>>,
Review Comment:
A bit late but I made a PR to follow up on this here:
https://github.com/apache/arrow-rs/pull/7477
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
