ceph/src/arrow/cpp/src/parquet/encryption/file_key_unwrapper.cc

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 #include <iostream>
  19
  20 #include "arrow/util/utf8.h"
  21
  22 #include "parquet/encryption/file_key_unwrapper.h"
  23 #include "parquet/encryption/key_metadata.h"
  24
  25 namespace parquet {
  26 namespace encryption {
  27
  28 using internal::KeyWithMasterId;
  29
  30 FileKeyUnwrapper::FileKeyUnwrapper(KeyToolkit* key_toolkit,
  31                                    const KmsConnectionConfig& kms_connection_config,
  32                                    double cache_lifetime_seconds)
  33     : key_toolkit_(key_toolkit),
  34       kms_connection_config_(kms_connection_config),
  35       cache_entry_lifetime_seconds_(cache_lifetime_seconds) {
  36   kek_per_kek_id_ = key_toolkit_->kek_read_cache_per_token().GetOrCreateInternalCache(
  37       kms_connection_config.key_access_token(), cache_entry_lifetime_seconds_);
  38 }
  39
  40 std::string FileKeyUnwrapper::GetKey(const std::string& key_metadata_bytes) {
  41   // key_metadata is expected to be in UTF8 encoding
  42   ::arrow::util::InitializeUTF8();
  43   if (!::arrow::util::ValidateUTF8(
  44           reinterpret_cast<const uint8_t*>(key_metadata_bytes.data()),
  45           key_metadata_bytes.size())) {
  46     throw ParquetException("key metadata should be in UTF8 encoding");
  47   }
  48   KeyMetadata key_metadata = KeyMetadata::Parse(key_metadata_bytes);
  49
  50   if (!key_metadata.key_material_stored_internally()) {
  51     throw ParquetException("External key material store is not supported yet.");
  52   }
  53
  54   const KeyMaterial& key_material = key_metadata.key_material();
  55
  56   return GetDataEncryptionKey(key_material).data_key();
  57 }
  58
  59 KeyWithMasterId FileKeyUnwrapper::GetDataEncryptionKey(const KeyMaterial& key_material) {
  60   auto kms_client = GetKmsClientFromConfigOrKeyMaterial(key_material);
  61
  62   bool double_wrapping = key_material.is_double_wrapped();
  63   const std::string& master_key_id = key_material.master_key_id();
  64   const std::string& encoded_wrapped_dek = key_material.wrapped_dek();
  65
  66   std::string data_key;
  67   if (!double_wrapping) {
  68     data_key = kms_client->UnwrapKey(encoded_wrapped_dek, master_key_id);
  69   } else {
  70     // Get Key Encryption Key
  71     const std::string& encoded_kek_id = key_material.kek_id();
  72     const std::string& encoded_wrapped_kek = key_material.wrapped_kek();
  73
  74     std::string kek_bytes = kek_per_kek_id_->GetOrInsert(
  75         encoded_kek_id, [kms_client, encoded_wrapped_kek, master_key_id]() {
  76           return kms_client->UnwrapKey(encoded_wrapped_kek, master_key_id);
  77         });
  78
  79     // Decrypt the data key
  80     std::string aad = ::arrow::util::base64_decode(encoded_kek_id);
  81     data_key = internal::DecryptKeyLocally(encoded_wrapped_dek, kek_bytes, aad);
  82   }
  83
  84   return KeyWithMasterId(data_key, master_key_id);
  85 }
  86
  87 std::shared_ptr<KmsClient> FileKeyUnwrapper::GetKmsClientFromConfigOrKeyMaterial(
  88     const KeyMaterial& key_material) {
  89   std::string& kms_instance_id = kms_connection_config_.kms_instance_id;
  90   if (kms_instance_id.empty()) {
  91     kms_instance_id = key_material.kms_instance_id();
  92     if (kms_instance_id.empty()) {
  93       throw ParquetException(
  94           "KMS instance ID is missing both in both kms connection configuration and file "
  95           "key material");
  96     }
  97   }
  98
  99   std::string& kms_instance_url = kms_connection_config_.kms_instance_url;
 100   if (kms_instance_url.empty()) {
 101     kms_instance_url = key_material.kms_instance_url();
 102     if (kms_instance_url.empty()) {
 103       throw ParquetException(
 104           "KMS instance ID is missing both in both kms connection configuration and file "
 105           "key material");
 106     }
 107   }
 108
 109   return key_toolkit_->GetKmsClient(kms_connection_config_,
 110                                     cache_entry_lifetime_seconds_);
 111 }
 112
 113 }  // namespace encryption
 114 }  // namespace parquet