parquet/encryption/key_toolkit.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <memory>
0021 #include <string>
0022
0023 #include "parquet/encryption/key_encryption_key.h"
0024 #include "parquet/encryption/kms_client.h"
0025 #include "parquet/encryption/kms_client_factory.h"
0026 #include "parquet/encryption/two_level_cache_with_expiration.h"
0027 #include "parquet/platform.h"
0028
0029 namespace parquet::encryption {
0030
0031 static constexpr uint64_t kCacheCleanPeriodForKeyRotation = 60 * 60;  // 1 hour
0032
0033 // KeyToolkit is a utility that keeps various tools for key management (such as key
0034 // rotation, kms client instantiation, cache control, etc), plus a number of auxiliary
0035 // classes for internal use.
0036 class PARQUET_EXPORT KeyToolkit {
0037  public:
0038   KeyToolkit() { last_cache_clean_for_key_rotation_time_ = {}; }
0039
0040   /// KMS client two level cache: token -> KMSInstanceId -> KmsClient
0041   TwoLevelCacheWithExpiration<std::shared_ptr<KmsClient>>& kms_client_cache_per_token() {
0042     return kms_client_cache_;
0043   }
0044   /// Key encryption key two level cache for wrapping: token -> MasterEncryptionKeyId ->
0045   /// KeyEncryptionKey
0046   TwoLevelCacheWithExpiration<KeyEncryptionKey>& kek_write_cache_per_token() {
0047     return key_encryption_key_write_cache_;
0048   }
0049
0050   /// Key encryption key two level cache for unwrapping: token -> KeyEncryptionKeyId ->
0051   /// KeyEncryptionKeyBytes
0052   TwoLevelCacheWithExpiration<::arrow::util::SecureString>& kek_read_cache_per_token() {
0053     return key_encryption_key_read_cache_;
0054   }
0055
0056   std::shared_ptr<KmsClient> GetKmsClient(
0057       const KmsConnectionConfig& kms_connection_config, double cache_entry_lifetime_ms);
0058
0059   /// Flush any caches that are tied to the (compromised) access_token
0060   void RemoveCacheEntriesForToken(const std::string& access_token);
0061
0062   void RemoveCacheEntriesForAllTokens();
0063
0064   void RegisterKmsClientFactory(std::shared_ptr<KmsClientFactory> kms_client_factory) {
0065     if (kms_client_factory_ != NULLPTR) {
0066       throw ParquetException("KMS client factory has already been registered.");
0067     }
0068     kms_client_factory_ = std::move(kms_client_factory);
0069   }
0070
0071   /// Key rotation. In the single wrapping mode, decrypts data keys with old master keys,
0072   /// then encrypts them with new master keys. In the double wrapping mode, decrypts KEKs
0073   /// (key encryption keys) with old master keys, generates new KEKs and encrypts them
0074   /// with new master keys. Works only if key material is not stored internally in file
0075   /// footers. Not supported in local key wrapping mode. Method can be run by multiple
0076   /// threads, but each thread must work on different files.
0077   void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config,
0078                         const std::string& parquet_file_path,
0079                         const std::shared_ptr<::arrow::fs::FileSystem>& file_system,
0080                         bool double_wrapping, double cache_lifetime_seconds);
0081
0082  private:
0083   TwoLevelCacheWithExpiration<std::shared_ptr<KmsClient>> kms_client_cache_;
0084   TwoLevelCacheWithExpiration<KeyEncryptionKey> key_encryption_key_write_cache_;
0085   TwoLevelCacheWithExpiration<::arrow::util::SecureString> key_encryption_key_read_cache_;
0086   std::shared_ptr<KmsClientFactory> kms_client_factory_;
0087   mutable ::arrow::util::Mutex last_cache_clean_for_key_rotation_time_mutex_;
0088   internal::TimePoint last_cache_clean_for_key_rotation_time_;
0089 };
0090
0091 // "data encryption key" and "master key identifier" are paired together as output when
0092 // parsing from "key material"
0093 class PARQUET_EXPORT KeyWithMasterId {
0094  public:
0095   KeyWithMasterId(::arrow::util::SecureString key_bytes, std::string master_id)
0096       : key_bytes_(std::move(key_bytes)), master_id_(std::move(master_id)) {}
0097
0098   const ::arrow::util::SecureString& data_key() const { return key_bytes_; }
0099   const std::string& master_id() const { return master_id_; }
0100
0101  private:
0102   ::arrow::util::SecureString key_bytes_;
0103   std::string master_id_;
0104 };
0105
0106 }  // namespace parquet::encryption