ceph/src/arrow/cpp/src/parquet/encryption/test_in_memory_kms.h

   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 #pragma once
  19
  20 #include <unordered_map>
  21
  22 #include "arrow/util/base64.h"
  23
  24 #include "parquet/encryption/kms_client_factory.h"
  25 #include "parquet/encryption/local_wrap_kms_client.h"
  26 #include "parquet/platform.h"
  27
  28 namespace parquet {
  29 namespace encryption {
  30
  31 // This is a mock class, built for testing only. Don't use it as an example of
  32 // LocalWrapKmsClient implementation.
  33 class TestOnlyLocalWrapInMemoryKms : public LocalWrapKmsClient {
  34  public:
  35   explicit TestOnlyLocalWrapInMemoryKms(const KmsConnectionConfig& kms_connection_config);
  36
  37   static void InitializeMasterKeys(
  38       const std::unordered_map<std::string, std::string>& master_keys_map);
  39
  40  protected:
  41   std::string GetMasterKeyFromServer(const std::string& master_key_identifier) override;
  42
  43  private:
  44   static std::unordered_map<std::string, std::string> master_key_map_;
  45 };
  46
  47 // This is a mock class, built for testing only. Don't use it as an example of KmsClient
  48 // implementation.
  49 class TestOnlyInServerWrapKms : public KmsClient {
  50  public:
  51   static void InitializeMasterKeys(
  52       const std::unordered_map<std::string, std::string>& master_keys_map);
  53
  54   std::string WrapKey(const std::string& key_bytes,
  55                       const std::string& master_key_identifier) override;
  56
  57   std::string UnwrapKey(const std::string& wrapped_key,
  58                         const std::string& master_key_identifier) override;
  59
  60  private:
  61   std::string GetMasterKeyFromServer(const std::string& master_key_identifier);
  62
  63   static std::unordered_map<std::string, std::string> master_key_map_;
  64 };
  65
  66 // This is a mock class, built for testing only. Don't use it as an example of
  67 // KmsClientFactory implementation.
  68 class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory {
  69  public:
  70   TestOnlyInMemoryKmsClientFactory(
  71       bool wrap_locally,
  72       const std::unordered_map<std::string, std::string>& master_keys_map)
  73       : KmsClientFactory(wrap_locally) {
  74     TestOnlyLocalWrapInMemoryKms::InitializeMasterKeys(master_keys_map);
  75     TestOnlyInServerWrapKms::InitializeMasterKeys(master_keys_map);
  76   }
  77
  78   std::shared_ptr<KmsClient> CreateKmsClient(
  79       const KmsConnectionConfig& kms_connection_config) {
  80     if (wrap_locally_) {
  81       return std::make_shared<TestOnlyLocalWrapInMemoryKms>(kms_connection_config);
  82     } else {
  83       return std::make_shared<TestOnlyInServerWrapKms>();
  84     }
  85   }
  86 };
  87
  88 }  // namespace encryption
  89 }  // namespace parquet