| 44 | } |
| 45 | |
| 46 | void KeyToolkit::RotateMasterKeys( |
| 47 | const KmsConnectionConfig& kms_connection_config, |
| 48 | const std::string& parquet_file_path, |
| 49 | const std::shared_ptr<::arrow::fs::FileSystem>& file_system, bool double_wrapping, |
| 50 | double cache_lifetime_seconds) { |
| 51 | // If process wrote files with double-wrapped keys, clean KEK cache (since master keys |
| 52 | // are changing). Only once for each key rotation cycle; not for every file. |
| 53 | const auto now = internal::CurrentTimePoint(); |
| 54 | auto lock = last_cache_clean_for_key_rotation_time_mutex_.Lock(); |
| 55 | if (now > last_cache_clean_for_key_rotation_time_ + |
| 56 | std::chrono::duration<double>(kCacheCleanPeriodForKeyRotation)) { |
| 57 | kek_write_cache_per_token().Clear(); |
| 58 | last_cache_clean_for_key_rotation_time_ = now; |
| 59 | } |
| 60 | lock.Unlock(); |
| 61 | |
| 62 | std::shared_ptr<FileKeyMaterialStore> key_material_store = |
| 63 | FileSystemKeyMaterialStore::Make(parquet_file_path, file_system, false); |
| 64 | |
| 65 | // Unwrapper for decrypting encrypted keys |
| 66 | FileKeyUnwrapper file_key_unwrapper(this, kms_connection_config, cache_lifetime_seconds, |
| 67 | key_material_store); |
| 68 | |
| 69 | // Create a temporary store to hold new key material during rotation, |
| 70 | // and wrapper that will write material to this store when getting key metadata. |
| 71 | std::shared_ptr<FileKeyMaterialStore> temp_key_material_store = |
| 72 | FileSystemKeyMaterialStore::Make(parquet_file_path, file_system, true); |
| 73 | FileKeyWrapper file_key_wrapper(this, kms_connection_config, temp_key_material_store, |
| 74 | cache_lifetime_seconds, double_wrapping); |
| 75 | |
| 76 | std::vector<std::string> file_key_id_set = key_material_store->GetKeyIDSet(); |
| 77 | |
| 78 | // When writing new encryption material, we re-use the same key identifiers |
| 79 | // so that key metadata does not need to change. |
| 80 | // Start with footer key (to get KMS ID, URL, if needed). |
| 81 | // We can rely on the footer key using a standardised key identifier. |
| 82 | std::string footer_key_id_str = std::string(KeyMaterial::kFooterKeyIdInFile); |
| 83 | std::string key_material_string = key_material_store->GetKeyMaterial(footer_key_id_str); |
| 84 | KeyWithMasterId key = |
| 85 | file_key_unwrapper.GetDataEncryptionKey(KeyMaterial::Parse(key_material_string)); |
| 86 | file_key_wrapper.GetEncryptionKeyMetadata(key.data_key(), key.master_id(), true, |
| 87 | std::string(KeyMaterial::kFooterKeyIdInFile)); |
| 88 | |
| 89 | // Rotate column keys |
| 90 | for (const auto& key_id_in_file : file_key_id_set) { |
| 91 | if (key_id_in_file == std::string(KeyMaterial::kFooterKeyIdInFile)) { |
| 92 | continue; |
| 93 | } |
| 94 | key_material_string = key_material_store->GetKeyMaterial(key_id_in_file); |
| 95 | KeyWithMasterId column_key = |
| 96 | file_key_unwrapper.GetDataEncryptionKey(KeyMaterial::Parse(key_material_string)); |
| 97 | file_key_wrapper.GetEncryptionKeyMetadata( |
| 98 | column_key.data_key(), column_key.master_id(), false, key_id_in_file); |
| 99 | } |
| 100 | |
| 101 | // Save material to the temporary store then move it to the original store location |
| 102 | temp_key_material_store->SaveMaterial(); |
| 103 | key_material_store->RemoveMaterial(); |