parquet/encryption/file_system_key_material_store.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 <set>
0021 #include <string>
0022 #include <unordered_map>
0023
0024 #include "arrow/filesystem/filesystem.h"
0025
0026 #include "parquet/encryption/file_key_material_store.h"
0027
0028 namespace parquet::encryption {
0029
0030 /// A FileKeyMaterialStore that stores key material in a file system file in the same
0031 /// folder as the Parquet file.
0032 class PARQUET_EXPORT FileSystemKeyMaterialStore : public FileKeyMaterialStore {
0033  public:
0034   static constexpr const char kKeyMaterialFilePrefix[] = "_KEY_MATERIAL_FOR_";
0035   static constexpr const char kTempFilePrefix[] = "_TMP";
0036   static constexpr const char kKeyMaterialFileSuffix[] = ".json";
0037
0038   FileSystemKeyMaterialStore() {}
0039   FileSystemKeyMaterialStore(std::string key_material_file_path,
0040                              std::shared_ptr<::arrow::fs::FileSystem> file_system);
0041
0042   /// Creates a new file system key material store for a parquet file.
0043   /// When use_tmp_prefix is true, files are saved with an extra _TMP prefix so they don't
0044   /// conflict with existing external material files. This is useful during key rotation
0045   /// so that temporary key material files can be created while using the existing key
0046   /// material, before moving the key material to the non-temporary location.
0047   static std::shared_ptr<FileSystemKeyMaterialStore> Make(
0048       std::string parquet_file_path, std::shared_ptr<::arrow::fs::FileSystem> file_system,
0049       bool use_tmp_prefix);
0050
0051   /// Add key material for one encryption key.
0052   void AddKeyMaterial(std::string key_id_in_file, std::string key_material) {
0053     key_material_map_.emplace(std::move(key_id_in_file), std::move(key_material));
0054   }
0055
0056   /// Get key material
0057   std::string GetKeyMaterial(std::string key_id_in_file) {
0058     if (key_material_map_.empty()) {
0059       LoadKeyMaterialMap();
0060     }
0061     auto found = key_material_map_.find(key_id_in_file);
0062     return found->second;
0063   }
0064
0065   /// After key material was added for all keys in the given Parquet file,
0066   /// save material in persistent store.
0067   void SaveMaterial();
0068
0069   /// Remove key material from persistent store. Used in key rotation.
0070   void RemoveMaterial();
0071
0072   /// Move key material to another store. Used in key rotation.
0073   void MoveMaterialTo(std::shared_ptr<FileKeyMaterialStore> target_key_store);
0074
0075   ///  Returns the Set of all key IDs in this store (for the given Parquet file)
0076   std::vector<std::string> GetKeyIDSet();
0077
0078  private:
0079   std::string GetStorageFilePath() { return key_material_file_path_; }
0080
0081   std::string BuildKeyMaterialMapJson();
0082   void LoadKeyMaterialMap();
0083   std::string key_material_file_path_;
0084   std::shared_ptr<::arrow::fs::FileSystem> file_system_;
0085   /// Maps ID of a key in Parquet file and key material
0086   std::unordered_map<std::string, std::string> key_material_map_;
0087 };
0088
0089 }  // namespace parquet::encryption