1 files changed, 121 insertions, 0 deletions
diff --git a/src/s3select/include/internal_file_decryptor.h b/src/s3select/include/internal_file_decryptor.h
new file mode 100644
index 000000000..011c4acbe
--- /dev/null
+++ b/src/s3select/include/internal_file_decryptor.h
@@ -0,0 +1,121 @@
+// 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.
+
+#pragma once
+
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "parquet/schema.h"
+
+namespace parquet {
+
+namespace encryption {
+class AesDecryptor;
+class AesEncryptor;
+}  // namespace encryption
+
+class FileDecryptionProperties;
+
+class PARQUET_EXPORT Decryptor {
+ public:
+  Decryptor(encryption::AesDecryptor* decryptor, const std::string& key,
+            const std::string& file_aad, const std::string& aad,
+            ::arrow::MemoryPool* pool);
+
+  const std::string& file_aad() const { return file_aad_; }
+  void UpdateAad(const std::string& aad) { aad_ = aad; }
+  ::arrow::MemoryPool* pool() { return pool_; }
+
+  int CiphertextSizeDelta();
+  int Decrypt(const uint8_t* ciphertext, int ciphertext_len, uint8_t* plaintext);
+
+ private:
+  encryption::AesDecryptor* aes_decryptor_;
+  std::string key_;
+  std::string file_aad_;
+  std::string aad_;
+  ::arrow::MemoryPool* pool_;
+};
+
+class InternalFileDecryptor {
+ public:
+  explicit InternalFileDecryptor(FileDecryptionProperties* properties,
+                                 const std::string& file_aad,
+                                 ParquetCipher::type algorithm,
+                                 const std::string& footer_key_metadata,
+                                 ::arrow::MemoryPool* pool);
+
+  std::string& file_aad() { return file_aad_; }
+
+  std::string GetFooterKey();
+
+  ParquetCipher::type algorithm() { return algorithm_; }
+
+  std::string& footer_key_metadata() { return footer_key_metadata_; }
+
+  FileDecryptionProperties* properties() { return properties_; }
+
+  void WipeOutDecryptionKeys();
+
+  ::arrow::MemoryPool* pool() { return pool_; }
+
+  std::shared_ptr<Decryptor> GetFooterDecryptor();
+  std::shared_ptr<Decryptor> GetFooterDecryptorForColumnMeta(const std::string& aad = "");
+  std::shared_ptr<Decryptor> GetFooterDecryptorForColumnData(const std::string& aad = "");
+  std::shared_ptr<Decryptor> GetColumnMetaDecryptor(
+      const std::string& column_path, const std::string& column_key_metadata,
+      const std::string& aad = "");
+  std::shared_ptr<Decryptor> GetColumnDataDecryptor(
+      const std::string& column_path, const std::string& column_key_metadata,
+      const std::string& aad = "");
+
+ private:
+  FileDecryptionProperties* properties_;
+  // Concatenation of aad_prefix (if exists) and aad_file_unique
+  std::string file_aad_;
+  std::map<std::string, std::shared_ptr<Decryptor>> column_data_map_;
+  std::map<std::string, std::shared_ptr<Decryptor>> column_metadata_map_;
+
+  std::shared_ptr<Decryptor> footer_metadata_decryptor_;
+  std::shared_ptr<Decryptor> footer_data_decryptor_;
+  ParquetCipher::type algorithm_;
+  std::string footer_key_metadata_;
+  std::vector<encryption::AesDecryptor*> all_decryptors_;
+
+  /// Key must be 16, 24 or 32 bytes in length. Thus there could be up to three
+  // types of meta_decryptors and data_decryptors.
+  std::unique_ptr<encryption::AesDecryptor> meta_decryptor_[3];
+  std::unique_ptr<encryption::AesDecryptor> data_decryptor_[3];
+
+  ::arrow::MemoryPool* pool_;
+
+  std::shared_ptr<Decryptor> GetFooterDecryptor(const std::string& aad, bool metadata);
+  std::shared_ptr<Decryptor> GetColumnDecryptor(const std::string& column_path,
+                                                const std::string& column_key_metadata,
+                                                const std::string& aad,
+                                                bool metadata = false);
+
+  encryption::AesDecryptor* GetMetaAesDecryptor(size_t key_size);
+  encryption::AesDecryptor* GetDataAesDecryptor(size_t key_size);
+
+  int MapKeyLenToDecryptorArrayIndex(int key_len);
+};
+
+}  // namespace parquet