text: // Licensed to the Apache Software Foundation (ASF) under one
  // or more contributor license agreements.  See the NOTICE file
  // distributed with this work for additional information
  // regarding copyright ownership.  The ASF licenses this file
  // to you under the Apache License, Version 2.0 (the
  // "License"); you may not use this file except in compliance
  // with the License.  You may obtain a copy of the License at
  //
  //   http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing,
  // software distributed under the License is distributed on an
  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, either express or implied.  See the License for the
  // specific language governing permissions and limitations
  // under the License.
  
  #include "parquet/encryption.h"
  
  #include <string.h>
  
  #include <map>
  #include <utility>
  
  #include "arrow/util/logging.h"
  #include "arrow/util/utf8.h"
  #include "parquet/encryption_internal.h"
  
  namespace parquet {
  
  // integer key retriever
  void IntegerKeyIdRetriever::PutKey(uint32_t key_id, const std::string& key) {
    key_map_.insert({key_id, key});
  }
  
  std::string IntegerKeyIdRetriever::GetKey(const std::string& key_metadata) const {
    uint32_t key_id;
    memcpy(reinterpret_cast<uint8_t*>(&key_id), key_metadata.c_str(), 4);
  
    return key_map_.at(key_id);
  }
  
  // string key retriever
  void StringKeyIdRetriever::PutKey(const std::string& key_id, const std::string& key) {
    key_map_.insert({key_id, key});
  }
  
  std::string StringKeyIdRetriever::GetKey(const std::string& key_id) const {
    return key_map_.at(key_id);
  }
  
  ColumnEncryptionProperties::Builder* ColumnEncryptionProperties::Builder::key(
      std::string column_key) {
    if (column_key.empty()) return this;
  
    DCHECK(key_.empty());
    key_ = column_key;
    return this;
  }
  
  ColumnEncryptionProperties::Builder* ColumnEncryptionProperties::Builder::key_metadata(
      const std::string& key_metadata) {
    DCHECK(!key_metadata.empty());
    DCHECK(key_metadata_.empty());
    key_metadata_ = key_metadata;
    return this;
  }
  
  ColumnEncryptionProperties::Builder* ColumnEncryptionProperties::Builder::key_id(
      const std::string& key_id) {
    // key_id is expected to be in UTF8 encoding
    ::arrow::util::InitializeUTF8();
    const uint8_t* data = reinterpret_cast<const uint8_t*>(key_id.c_str());
    if (!::arrow::util::ValidateUTF8(data, key_id.size())) {
      throw ParquetException("key id should be in UTF8 encoding");
    }
  
    DCHECK(!key_id.empty());
    this->key_metadata(key_id);
    return this;
  }
  
  FileDecryptionProperties::Builder* FileDecryptionProperties::Builder::column_keys(
      const ColumnPathToDecryptionPropertiesMap& column_decryption_properties) {
    if (column_decryption_properties.size() == 0) return this;
  
    if (column_decryption_properties_.size() != 0)
      throw ParquetException("Column properties already set");
  
    for (const auto& element : column_decryption_properties) {
      if (element.second->is_utilized()) {
        throw ParquetException("Column properties utilized in another file");
      }
      element.second->set_utilized();
    }
  
    column_decryption_properties_ = column_decryption_properties;
    return this;
  }
  
  void FileDecryptionProperties::WipeOutDecryptionKeys() {
    footer_key_.clear();
  
    for (const auto& element : column_decryption_properties_) {
      element.second->WipeOutDecryptionKey();
    }
  }
  
  bool FileDecryptionProperties::is_utilized() {
    if (footer_key_.empty() && column_decryption_properties_.size() == 0 &&
        aad_prefix_.empty())
      return false;
  
    return utilized_;
  }
  
  std::shared_ptr<FileDecryptionProperties> FileDecryptionProperties::DeepClone(
      std::string new_aad_prefix) {
    std::string footer_key_copy = footer_key_;
    ColumnPathToDecryptionPropertiesMap column_decryption_properties_map_copy;
  
    for (const auto& element : column_decryption_properties_) {
      column_decryption_properties_map_copy.insert(
          {element.second->column_path(), element.second->DeepClone()});
    }
  
    if (new_aad_prefix.empty()) new_aad_prefix = aad_prefix_;
    return std::shared_ptr<FileDecryptionProperties>(new FileDecryptionProperties(
        footer_key_copy, key_retriever_, check_plaintext_footer_integrity_, new_aad_prefix,
        aad_prefix_verifier_, column_decryption_properties_map_copy,
        plaintext_files_allowed_));
  }
  
  FileDecryptionProperties::Builder* FileDecryptionProperties::Builder::footer_key(
      const std::string footer_key) {
    if (footer_key.empty()) {
      return this;
    }
    DCHECK(footer_key_.empty());
    footer_key_ = footer_key;
    return this;
  }
  
  FileDecryptionProperties::Builder* FileDecryptionProperties::Builder::key_retriever(
      const std::shared_ptr<DecryptionKeyRetriever>& key_retriever) {
    if (key_retriever == nullptr) return this;
  
    DCHECK(key_retriever_ == nullptr);
    key_retriever_ = key_retriever;
    return this;
  }
  
  FileDecryptionProperties::Builder* FileDecryptionProperties::Builder::aad_prefix(
      const std::string& aad_prefix) {
    if (aad_prefix.empty()) {
      return this;
    }
    DCHECK(aad_prefix_.empty());
    aad_prefix_ = aad_prefix;
    return this;
  }
  
  FileDecryptionProperties::Builder* FileDecryptionProperties::Builder::aad_prefix_verifier(
      std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier) {
    if (aad_prefix_verifier == nullptr) return this;
  
    DCHECK(aad_prefix_verifier_ == nullptr);
    aad_prefix_verifier_ = std::move(aad_prefix_verifier);
    return this;
  }
  
  ColumnDecryptionProperties::Builder* ColumnDecryptionProperties::Builder::key(
      const std::string& key) {
    if (key.empty()) return this;
  
    DCHECK(!key.empty());
    key_ = key;
    return this;
  }
  
  std::shared_ptr<ColumnDecryptionProperties> ColumnDecryptionProperties::Builder::build() {
    return std::shared_ptr<ColumnDecryptionProperties>(
        new ColumnDecryptionProperties(column_path_, key_));
  }
  
  void ColumnDecryptionProperties::WipeOutDecryptionKey() { key_.clear(); }
  
  std::shared_ptr<ColumnDecryptionProperties> ColumnDecryptionProperties::DeepClone() {
    std::string key_copy = key_;
    return std::shared_ptr<ColumnDecryptionProperties>(
        new ColumnDecryptionProperties(column_path_, key_copy));
  }
  
  FileEncryptionProperties::Builder* FileEncryptionProperties::Builder::footer_key_metadata(
      const std::string& footer_key_metadata) {
    if (footer_key_metadata.empty()) return this;
  
    DCHECK(footer_key_metadata_.empty());
    footer_key_metadata_ = footer_key_metadata;
    return this;
  }
  
  FileEncryptionProperties::Builder* FileEncryptionProperties::Builder::encrypted_columns(
      const ColumnPathToEncryptionPropertiesMap& encrypted_columns) {
    if (encrypted_columns.size() == 0) return this;
  
    if (encrypted_columns_.size() != 0)
      throw ParquetException("Column properties already set");
  
    for (const auto& element : encrypted_columns) {
      if (element.second->is_utilized()) {
        throw ParquetException("Column properties utilized in another file");
      }
      element.second->set_utilized();
    }
    encrypted_columns_ = encrypted_columns;
    return this;
  }
  
  void FileEncryptionProperties::WipeOutEncryptionKeys() {
    footer_key_.clear();
    for (const auto& element : encrypted_columns_) {
      element.second->WipeOutEncryptionKey();
    }
  }
  
  std::shared_ptr<FileEncryptionProperties> FileEncryptionProperties::DeepClone(
      std::string new_aad_prefix) {
    std::string footer_key_copy = footer_key_;
    ColumnPathToEncryptionPropertiesMap encrypted_columns_map_copy;
  
    for (const auto& element : encrypted_columns_) {
      encrypted_columns_map_copy.insert(
          {element.second->column_path(), element.second->DeepClone()});
    }
  
    if (new_aad_prefix.empty()) new_aad_prefix = aad_prefix_;
    return std::shared_ptr<FileEncryptionProperties>(new FileEncryptionProperties(
        algorithm_.algorithm, footer_key_copy, footer_key_metadata_, encrypted_footer_,
        new_aad_prefix, store_aad_prefix_in_file_, encrypted_columns_map_copy));
  }
  
  FileEncryptionProperties::Builder* FileEncryptionProperties::Builder::aad_prefix(
      const std::string& aad_prefix) {
    if (aad_prefix.empty()) return this;
  
    DCHECK(aad_prefix_.empty());
    aad_prefix_ = aad_prefix;
    store_aad_prefix_in_file_ = true;
    return this;
  }
  
  FileEncryptionProperties::Builder*
  FileEncryptionProperties::Builder::disable_aad_prefix_storage() {
    DCHECK(!aad_prefix_.empty());
  
    store_aad_prefix_in_file_ = false;
    return this;
  }
  
  ColumnEncryptionProperties::ColumnEncryptionProperties(bool encrypted,
                                                         const std::string& column_path,
                                                         const std::string& key,
                                                         const std::string& key_metadata)
      : column_path_(column_path) {
    // column encryption properties object (with a column key) can be used for writing only
    // one file.
    // Upon completion of file writing, the encryption keys in the properties will be wiped
    // out (set to 0 in memory).
    utilized_ = false;
  
    DCHECK(!column_path.empty());
    if (!encrypted) {
      DCHECK(key.empty() && key_metadata.empty());
    }
  
    if (!key.empty()) {
      DCHECK(key.length() == 16 || key.length() == 24 || key.length() == 32);
    }
  
    encrypted_with_footer_key_ = (encrypted && key.empty());
    if (encrypted_with_footer_key_) {
      DCHECK(key_metadata.empty());
    }
  
    encrypted_ = encrypted;
    key_metadata_ = key_metadata;
    key_ = key;
  }
  
  ColumnDecryptionProperties::ColumnDecryptionProperties(const std::string& column_path,
                                                         const std::string& key)
      : column_path_(column_path) {
    utilized_ = false;
    DCHECK(!column_path.empty());
  
    if (!key.empty()) {
      DCHECK(key.length() == 16 || key.length() == 24 || key.length() == 32);
    }
  
    key_ = key;
  }
  
  std::string FileDecryptionProperties::column_key(const std::string& column_path) const {
    if (column_decryption_properties_.find(column_path) !=
        column_decryption_properties_.end()) {
      auto column_prop = column_decryption_properties_.at(column_path);
      if (column_prop != nullptr) {
        return column_prop->key();
      }
    }
    return empty_string_;
  }
  
  FileDecryptionProperties::FileDecryptionProperties(
      const std::string& footer_key, std::shared_ptr<DecryptionKeyRetriever> key_retriever,
      bool check_plaintext_footer_integrity, const std::string& aad_prefix,
      std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier,
      const ColumnPathToDecryptionPropertiesMap& column_decryption_properties,
      bool plaintext_files_allowed) {
    DCHECK(!footer_key.empty() || nullptr != key_retriever ||
           0 != column_decryption_properties.size());
  
    if (!footer_key.empty()) {
      DCHECK(footer_key.length() == 16 || footer_key.length() == 24 ||
             footer_key.length() == 32);
    }
    if (footer_key.empty() && check_plaintext_footer_integrity) {
      DCHECK(nullptr != key_retriever);
    }
    aad_prefix_verifier_ = std::move(aad_prefix_verifier);
    footer_key_ = footer_key;
    check_plaintext_footer_integrity_ = check_plaintext_footer_integrity;
    key_retriever_ = std::move(key_retriever);
    aad_prefix_ = aad_prefix;
    column_decryption_properties_ = column_decryption_properties;
    plaintext_files_allowed_ = plaintext_files_allowed;
    utilized_ = false;
  }
  
  FileEncryptionProperties::Builder* FileEncryptionProperties::Builder::footer_key_id(
      const std::string& key_id) {
    // key_id is expected to be in UTF8 encoding
    ::arrow::util::InitializeUTF8();
    const uint8_t* data = reinterpret_cast<const uint8_t*>(key_id.c_str());
    if (!::arrow::util::ValidateUTF8(data, key_id.size())) {
      throw ParquetException("footer key id should be in UTF8 encoding");
    }
  
    if (key_id.empty()) {
      return this;
    }
  
    return footer_key_metadata(key_id);
  }
  
  std::shared_ptr<ColumnEncryptionProperties>
  FileEncryptionProperties::column_encryption_properties(const std::string& column_path) {
    if (encrypted_columns_.size() == 0) {
      auto builder = std::make_shared<ColumnEncryptionProperties::Builder>(column_path);
      return builder->build();
    }
    if (encrypted_columns_.find(column_path) != encrypted_columns_.end()) {
      return encrypted_columns_[column_path];
    }
  
    return nullptr;
  }
  
  FileEncryptionProperties::FileEncryptionProperties(
      ParquetCipher::type cipher, const std::string& footer_key,
      const std::string& footer_key_metadata, bool encrypted_footer,
      const std::string& aad_prefix, bool store_aad_prefix_in_file,
      const ColumnPathToEncryptionPropertiesMap& encrypted_columns)
      : footer_key_(footer_key),
        footer_key_metadata_(footer_key_metadata),
        encrypted_footer_(encrypted_footer),
        aad_prefix_(aad_prefix),
        store_aad_prefix_in_file_(store_aad_prefix_in_file),
        encrypted_columns_(encrypted_columns) {
    // file encryption properties object can be used for writing only one file.
    // Upon completion of file writing, the encryption keys in the properties will be wiped
    // out (set to 0 in memory).
    utilized_ = false;
  
    DCHECK(!footer_key.empty());
    // footer_key must be either 16, 24 or 32 bytes.
    DCHECK(footer_key.length() == 16 || footer_key.length() == 24 ||
           footer_key.length() == 32);
  
    uint8_t aad_file_unique[kAadFileUniqueLength];
    memset(aad_file_unique, 0, kAadFileUniqueLength);
    encryption::RandBytes(aad_file_unique, sizeof(kAadFileUniqueLength));
    std::string aad_file_unique_str(reinterpret_cast<char const*>(aad_file_unique),
                                    kAadFileUniqueLength);
  
    bool supply_aad_prefix = false;
    if (aad_prefix.empty()) {
      file_aad_ = aad_file_unique_str;
    } else {
      file_aad_ = aad_prefix + aad_file_unique_str;
      if (!store_aad_prefix_in_file) supply_aad_prefix = true;
    }
    algorithm_.algorithm = cipher;
    algorithm_.aad.aad_file_unique = aad_file_unique_str;
    algorithm_.aad.supply_aad_prefix = supply_aad_prefix;
    if (!aad_prefix.empty() && store_aad_prefix_in_file) {
      algorithm_.aad.aad_prefix = aad_prefix;
    }
  }
  
  }  // namespace parquet
  