Add GetExternalFileEncryptionProperties method to C++ CryptoFactory

Add GetExternalFileEncryptionProperties method to C++ CryptoFactory

Author: sofia-tekdatum

cpp/src/parquet/CMakeLists.txt

@@ -413,7 +413,9 @@ if(PARQUET_REQUIRE_ENCRYPTION)
                    encryption/write_configurations_test.cc
                    encryption/read_configurations_test.cc
                    encryption/properties_test.cc
-                   encryption/test_encryption_util.cc)
+                   encryption/crypto_factory_test.cc
+                   encryption/test_encryption_util.cc
+                   encryption/test_in_memory_kms.cc)
   add_parquet_test(encryption-key-management-test
                    SOURCES
                    encryption/key_management_test.cc

cpp/src/parquet/encryption/crypto_factory.cc

@@ -28,24 +28,21 @@
 
 namespace parquet::encryption {
 
-void CryptoFactory::RegisterKmsClientFactory(
-    std::shared_ptr<KmsClientFactory> kms_client_factory) {
-  key_toolkit_->RegisterKmsClientFactory(std::move(kms_client_factory));
-}
+/// Extracting functionality common to both GetFileEncryptionProperties and
+/// GetExternalFileEncryptionProperties here for reuse.
+namespace {
 
-std::shared_ptr<FileEncryptionProperties> CryptoFactory::GetFileEncryptionProperties(
-    const KmsConnectionConfig& kms_connection_config,
-    const EncryptionConfiguration& encryption_config, const std::string& file_path,
-    const std::shared_ptr<::arrow::fs::FileSystem>& file_system) {
-  if (!encryption_config.uniform_encryption && encryption_config.column_keys.empty()) {
-    throw ParquetException("Either column_keys or uniform_encryption must be set");
-  } else if (encryption_config.uniform_encryption &&
-             !encryption_config.column_keys.empty()) {
-    throw ParquetException("Cannot set both column_keys and uniform_encryption");
-  }
-  const std::string& footer_key_id = encryption_config.footer_key;
-  const std::string& column_key_str = encryption_config.column_keys;
+// Struct to simplify the returned objects in GetFileKeyUtils.
+struct FileKeyUtils {
+  std::shared_ptr<FileKeyMaterialStore> key_material_store;
+  FileKeyWrapper key_wrapper;
+};
 
+FileKeyUtils GetFileKeyUtils(
+    const std::shared_ptr<KeyToolkit>& key_toolkit,
+    const KmsConnectionConfig& kms_connection_config,
+    const EncryptionConfiguration& encryption_config,
+    const std::string& file_path, const std::shared_ptr<::arrow::fs::FileSystem>& file_system) {
   std::shared_ptr<FileKeyMaterialStore> key_material_store = nullptr;
   if (!encryption_config.internal_key_material) {
     try {
@@ -58,18 +55,46 @@ std::shared_ptr<FileEncryptionProperties> CryptoFactory::GetFileEncryptionProper
     }
   }
 
-  FileKeyWrapper key_wrapper(key_toolkit_.get(), kms_connection_config,
+  FileKeyWrapper key_wrapper(key_toolkit.get(), kms_connection_config,
                              key_material_store, encryption_config.cache_lifetime_seconds,
                              encryption_config.double_wrapping);
 
-  int32_t dek_length_bits = encryption_config.data_key_length_bits;
+  return {key_material_store, std::move(key_wrapper)};
+}
+
+int ValidateAndGetKeyLength(int32_t dek_length_bits) {
   if (!internal::ValidateKeyLength(dek_length_bits)) {
     std::ostringstream ss;
     ss << "Wrong data key length : " << dek_length_bits;
     throw ParquetException(ss.str());
   }
+  return dek_length_bits / 8;
+}
 
-  int dek_length = dek_length_bits / 8;
+}  // Anonymous namespace
+
+void CryptoFactory::RegisterKmsClientFactory(
+    std::shared_ptr<KmsClientFactory> kms_client_factory) {
+  key_toolkit_->RegisterKmsClientFactory(std::move(kms_client_factory));
+}
+
+std::shared_ptr<FileEncryptionProperties> CryptoFactory::GetFileEncryptionProperties(
+    const KmsConnectionConfig& kms_connection_config,
+    const EncryptionConfiguration& encryption_config, const std::string& file_path,
+    const std::shared_ptr<::arrow::fs::FileSystem>& file_system) {
+  if (!encryption_config.uniform_encryption && encryption_config.column_keys.empty()) {
+    throw ParquetException("Either column_keys or uniform_encryption must be set");
+  } else if (encryption_config.uniform_encryption &&
+             !encryption_config.column_keys.empty()) {
+    throw ParquetException("Cannot set both column_keys and uniform_encryption");
+  }
+  const std::string& footer_key_id = encryption_config.footer_key;
+  const std::string& column_key_str = encryption_config.column_keys;
+
+  auto [key_material_store, key_wrapper] = GetFileKeyUtils(
+      key_toolkit_, kms_connection_config, encryption_config, file_path, file_system);
+
+  int dek_length = ValidateAndGetKeyLength(encryption_config.data_key_length_bits);
 
   std::string footer_key(dek_length, '\0');
   RandBytes(reinterpret_cast<uint8_t*>(footer_key.data()), footer_key.size());
@@ -98,6 +123,92 @@ std::shared_ptr<FileEncryptionProperties> CryptoFactory::GetFileEncryptionProper
   return properties_builder.build();
 }
 
+std::shared_ptr<ExternalFileEncryptionProperties>
+CryptoFactory::GetExternalFileEncryptionProperties(
+      const KmsConnectionConfig& kms_connection_config,
+      const ExternalEncryptionConfiguration& external_encryption_config,
+      const std::string& file_path, const std::shared_ptr<::arrow::fs::FileSystem>& file_system) {
+  // Validate the same rules as FileEncryptionProperties but considering per_column_encryption too.
+  // If uniform_encryption is not set then either column_keys or per_column_encryption must have
+  // values.
+  // If uniform_encryption is set, then both column_keys and per_column_encryption must be empty.
+  bool no_columns_encrypted = external_encryption_config.column_keys.empty() &&
+                              external_encryption_config.per_column_encryption.empty();
+  if (!external_encryption_config.uniform_encryption && no_columns_encrypted) {
+    throw ParquetException(
+      "Either uniform_encryption must be set or column encryption must be specified in either "
+      "column_keys or per_column_encryption");
+  } else if (external_encryption_config.uniform_encryption && !no_columns_encrypted) {
+    throw ParquetException("Cannot set both column encryption and uniform_encryption");
+  }
+
+  auto [key_material_store, key_wrapper] = GetFileKeyUtils(
+      key_toolkit_, kms_connection_config, external_encryption_config, file_path, file_system);
+
+  int dek_length = ValidateAndGetKeyLength(external_encryption_config.data_key_length_bits);
+
+  std::string footer_key(dek_length, '\0');
+  RandBytes(reinterpret_cast<uint8_t*>(footer_key.data()), footer_key.size());
+
+  std::string footer_key_metadata =
+      key_wrapper.GetEncryptionKeyMetadata(footer_key, external_encryption_config.footer_key, true);
+  
+  ExternalFileEncryptionProperties::Builder external_properties_builder =
+      ExternalFileEncryptionProperties::Builder(external_encryption_config.footer_key);
+  external_properties_builder.footer_key_metadata(footer_key_metadata);
+  external_properties_builder.algorithm(external_encryption_config.encryption_algorithm);
+
+  if (!external_encryption_config.uniform_encryption && 
+      external_encryption_config.plaintext_footer) {
+        external_properties_builder.set_plaintext_footer();
+  }
+  
+  ColumnPathToEncryptionPropertiesMap encrypted_columns;
+  if (!external_encryption_config.column_keys.empty()) {
+    encrypted_columns = GetColumnEncryptionProperties(
+        dek_length, external_encryption_config.column_keys, &key_wrapper);
+  }
+  if (!external_encryption_config.per_column_encryption.empty()) {
+    for (const auto& pair : external_encryption_config.per_column_encryption) {
+      const std::string& column_name = pair.first;
+      const ColumnEncryptionAttributes& attributes = pair.second;
+
+      // Validate column names are not in both column_keys and per_column_encryption maps.
+      if (encrypted_columns.find(column_name) != encrypted_columns.end()) {
+        std::stringstream string_stream;
+        string_stream << "Multiple keys defined for column [" << column_name << "]. ";
+        string_stream << "Keys found in column_keys and in per_column_encryption.";
+        throw ParquetException(string_stream.str());
+      }
+
+      // TODO(sbrenes): Check whether the attributes.parquet_cipher == EXTERNAL.
+      // If so, do not use KMS to resolve the column_key, just forward it.
+      std::string column_key(dek_length, '\0');
+      RandBytes(reinterpret_cast<uint8_t*>(column_key.data()), column_key.size());
+      std::string column_key_metadata =
+          key_wrapper.GetEncryptionKeyMetadata(column_key, attributes.key_id, false);
+      
+      std::shared_ptr<ColumnEncryptionProperties> column_properties =
+          ColumnEncryptionProperties::Builder(column_name)
+              .key(column_key)
+              ->key_metadata(column_key_metadata)
+              ->parquet_cipher(attributes.parquet_cipher)
+              ->build();
+      
+      encrypted_columns.insert({column_name, column_properties});
+    }
+  }
+  if (!encrypted_columns.empty()) {
+    external_properties_builder.encrypted_columns(encrypted_columns);
+  }
+
+  if (key_material_store != nullptr) {
+    key_material_store->SaveMaterial();
+  }
+
+  return external_properties_builder.build_external();
+}
+
 ColumnPathToEncryptionPropertiesMap CryptoFactory::GetColumnEncryptionProperties(
     int dek_length, const std::string& column_keys, FileKeyWrapper* key_wrapper) {
   ColumnPathToEncryptionPropertiesMap encrypted_columns;

cpp/src/parquet/encryption/crypto_factory.h

@@ -115,13 +115,15 @@ struct PARQUET_EXPORT ExternalEncryptionConfiguration : public EncryptionConfigu
   /// algorithm specified in the encryption_algorithm field. 
   /// If a column name appears in the new per_column_encryption map, it will be encrypted using the
   /// per column specific algorithm and key.
-  /// If a column name appears in both, the per_column_encryption values will take precedence.
+  /// If a column name appears in both, an exception will be thrown.
   std::unordered_map<std::string, ColumnEncryptionAttributes> per_column_encryption;
 
   /// External encryption services may use additional context provided by the application to
   /// enforce robust access control. The values sent to the external service depend on each
   /// implementation. 
-  /// This values must be a valid JSON-formatted string.
+  /// This value must be a valid JSON-formatted string.
+  /// Validation of the string will be done by the external encryption service, Arrow will only
+  /// forward this value.
   /// Format: "{\"user_id\": \"abc123\", \"location\": {\"lat\": 9.7489, \"lon\": -83.7534}}"
   std::string app_context;
 
@@ -153,13 +155,21 @@ class PARQUET_EXPORT CryptoFactory {
   void RegisterKmsClientFactory(std::shared_ptr<KmsClientFactory> kms_client_factory);
 
   /// Get the encryption properties for a Parquet file.
-  /// If external key material is used then a file system and path to the
+  /// If key material from outside the file is used, then a file system and path to the
   /// parquet file must be provided.
   std::shared_ptr<FileEncryptionProperties> GetFileEncryptionProperties(
       const KmsConnectionConfig& kms_connection_config,
       const EncryptionConfiguration& encryption_config, const std::string& file_path = "",
       const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR);
 
+  /// Get the external encryption properties for a Parquet file. Used when encryption
+  /// will be provided by an external service.
+  std::shared_ptr<ExternalFileEncryptionProperties> GetExternalFileEncryptionProperties(
+      const KmsConnectionConfig& kms_connection_config,
+      const ExternalEncryptionConfiguration& external_encryption_config,
+      const std::string& file_path = "",
+      const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR);
+
   /// Get decryption properties for a Parquet file.
   /// If external key material is used then a file system and path to the
   /// parquet file must be provided.