parquet/encryption/test_in_memory_kms.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 <unordered_map>
0021
0022 #include "arrow/util/base64.h"
0023
0024 #include "parquet/encryption/kms_client_factory.h"
0025 #include "parquet/encryption/local_wrap_kms_client.h"
0026 #include "parquet/platform.h"
0027
0028 namespace parquet::encryption {
0029
0030 // This is a mock class, built for testing only. Don't use it as an example of
0031 // LocalWrapKmsClient implementation.
0032 class TestOnlyLocalWrapInMemoryKms : public LocalWrapKmsClient {
0033  public:
0034   explicit TestOnlyLocalWrapInMemoryKms(const KmsConnectionConfig& kms_connection_config);
0035
0036   static void InitializeMasterKeys(
0037       const std::unordered_map<std::string, ::arrow::util::SecureString>&
0038           master_keys_map);
0039
0040  protected:
0041   const ::arrow::util::SecureString& GetMasterKeyFromServer(
0042       const std::string& master_key_identifier) override;
0043
0044  private:
0045   static std::unordered_map<std::string, ::arrow::util::SecureString> master_key_map_;
0046 };
0047
0048 // This is a mock class, built for testing only. Don't use it as an example of KmsClient
0049 // implementation.
0050 class TestOnlyInServerWrapKms : public KmsClient {
0051  public:
0052   static void InitializeMasterKeys(
0053       const std::unordered_map<std::string, ::arrow::util::SecureString>&
0054           master_keys_map);
0055
0056   std::string WrapKey(const ::arrow::util::SecureString& key_bytes,
0057                       const std::string& master_key_identifier) override;
0058
0059   ::arrow::util::SecureString UnwrapKey(
0060       const std::string& wrapped_key, const std::string& master_key_identifier) override;
0061
0062   static void StartKeyRotation(
0063       const std::unordered_map<std::string, ::arrow::util::SecureString>&
0064           new_master_keys_map);
0065   static void FinishKeyRotation();
0066
0067  private:
0068   ::arrow::util::SecureString GetMasterKeyFromServer(
0069       const std::string& master_key_identifier);
0070
0071   // Different wrapping and unwrapping key maps to imitate versioning
0072   // and support key rotation.
0073   static std::unordered_map<std::string, ::arrow::util::SecureString>
0074       unwrapping_master_key_map_;
0075   static std::unordered_map<std::string, ::arrow::util::SecureString>
0076       wrapping_master_key_map_;
0077 };
0078
0079 // This is a mock class, built for testing only. Don't use it as an example of
0080 // KmsClientFactory implementation.
0081 class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory {
0082  public:
0083   TestOnlyInMemoryKmsClientFactory(
0084       bool wrap_locally,
0085       const std::unordered_map<std::string, ::arrow::util::SecureString>& master_keys_map)
0086       : KmsClientFactory(wrap_locally) {
0087     TestOnlyLocalWrapInMemoryKms::InitializeMasterKeys(master_keys_map);
0088     TestOnlyInServerWrapKms::InitializeMasterKeys(master_keys_map);
0089   }
0090
0091   std::shared_ptr<KmsClient> CreateKmsClient(
0092       const KmsConnectionConfig& kms_connection_config) {
0093     if (wrap_locally_) {
0094       return std::make_shared<TestOnlyLocalWrapInMemoryKms>(kms_connection_config);
0095     } else {
0096       return std::make_shared<TestOnlyInServerWrapKms>();
0097     }
0098   }
0099 };
0100
0101 }  // namespace parquet::encryption