Add functions for AAD calculation and adjust code for API changes

cpp/src/parquet/util/crypto.cc

@@ -249,16 +249,16 @@ int SignedFooterEncrypt(const uint8_t* plaintext, int plaintext_len, uint8_t* ke
     return gcm_encrypt(plaintext, plaintext_len, key, key_len, nonce, nonce_len, aad, aad_len, ciphertext);
 }
 
-int Encrypt(Encryption::type alg_id, bool metadata, const uint8_t* plaintext,
+int Encrypt(ParquetCipher::type alg_id, bool metadata, const uint8_t* plaintext,
             int plaintext_len, uint8_t* key, int key_len, uint8_t* aad, int aad_len,
             uint8_t* ciphertext) {
-  if (Encryption::AES_GCM_V1 != alg_id && Encryption::AES_GCM_CTR_V1 != alg_id) {
+  if (ParquetCipher::AES_GCM_V1 != alg_id && ParquetCipher::AES_GCM_CTR_V1 != alg_id) {
     std::stringstream ss;
     ss << "Crypto algorithm " << alg_id << " is not supported";
     throw ParquetException(ss.str());
   }
 
-  if (metadata || (Encryption::AES_GCM_V1 == alg_id)) {
+  if (metadata || (ParquetCipher::AES_GCM_V1 == alg_id)) {
     uint8_t nonce[nonceLen];
     memset(nonce, 0, nonceLen);
 
@@ -393,16 +393,16 @@ int ctr_decrypt(const uint8_t* ciphertext, int ciphertext_len, uint8_t* key, int
   return plaintext_len;
 }
 
-int Decrypt(Encryption::type alg_id, bool metadata, const uint8_t* ciphertext,
+int Decrypt(ParquetCipher::type alg_id, bool metadata, const uint8_t* ciphertext,
             int ciphertext_len, uint8_t* key, int key_len, uint8_t* aad, int aad_len,
             uint8_t* plaintext) {
-  if (Encryption::AES_GCM_V1 != alg_id && Encryption::AES_GCM_CTR_V1 != alg_id) {
+  if (ParquetCipher::AES_GCM_V1 != alg_id && ParquetCipher::AES_GCM_CTR_V1 != alg_id) {
     std::stringstream ss;
     ss << "Crypto algorithm " << alg_id << " is not supported";
     throw ParquetException(ss.str());
   }
 
-  if (metadata || (Encryption::AES_GCM_V1 == alg_id)) {
+  if (metadata || (ParquetCipher::AES_GCM_V1 == alg_id)) {
     return gcm_decrypt(ciphertext, ciphertext_len, key, key_len, aad, aad_len, plaintext);
   }
 
@@ -410,6 +410,16 @@ int Decrypt(Encryption::type alg_id, bool metadata, const uint8_t* ciphertext,
   return ctr_decrypt(ciphertext, ciphertext_len, key, key_len, plaintext);
 }
 
+static std::string shortToBytesLE(int16_t input) {
+  int8_t output[2];
+  memset(output, 0, 2);
+  output[1] = (int8_t)(0xff & (input >> 8));
+  output[0] = (int8_t)(0xff & (input));
+  std::string output_str(reinterpret_cast<char const*>(output), 2) ;
+
+  return output_str;
+}
+
 int Decrypt(std::shared_ptr<EncryptionProperties> encryption_props, bool metadata,
             const uint8_t* ciphertext, int ciphertext_len, uint8_t* plaintext) {
   return Decrypt(encryption_props->algorithm(), metadata, ciphertext, ciphertext_len,
@@ -418,4 +428,39 @@ int Decrypt(std::shared_ptr<EncryptionProperties> encryption_props, bool metadat
                  plaintext);
 }
 
+std::string createModuleAAD(const std::string& fileAAD, int8_t module_type,
+			    int16_t row_group_ordinal, int16_t column_ordinal,
+			    int16_t page_ordinal) {
+  
+  int8_t type_ordinal_bytes[1];
+  type_ordinal_bytes[0] = module_type;
+  std::string type_ordinal_bytes_str(reinterpret_cast<char const*>(type_ordinal_bytes), 1) ;
+  if (Footer == module_type) {
+    std::string result = fileAAD + type_ordinal_bytes_str;
+    return result;
+  }
+  std::string row_group_ordinal_bytes = shortToBytesLE(row_group_ordinal);
+  std::string column_ordinal_bytes = shortToBytesLE(column_ordinal);
+  if (DataPage != module_type && DataPageHeader != module_type) {
+    std::string result =  fileAAD + type_ordinal_bytes_str + row_group_ordinal_bytes
+      + column_ordinal_bytes;
+    return result;
+  }
+  std::string page_ordinal_bytes = shortToBytesLE(page_ordinal);
+  std::string result = fileAAD + type_ordinal_bytes_str + row_group_ordinal_bytes
+    + column_ordinal_bytes + page_ordinal_bytes;;
+  return result;
+}
+  
+std::string createFooterAAD(const std::string& aad_prefix_bytes) {
+  return createModuleAAD(aad_prefix_bytes, Footer, (int16_t) -1, (int16_t) -1, (int16_t) -1);
+}
+  
+// Update last two bytes with new page ordinal (instead of creating new page AAD from scratch)
+void quickUpdatePageAAD(const std::string &AAD, int16_t new_page_ordinal) {
+  std::string page_ordinal_bytes = shortToBytesLE(new_page_ordinal);
+  int length = (int)AAD.size();
+  std::memcpy((int16_t*)(AAD.c_str()+length-2), (int16_t*)(page_ordinal_bytes.c_str()), 2);
+}
+
 }  // namespace parquet_encryption

cpp/src/parquet/util/crypto.h

@@ -23,28 +23,47 @@
 #include "parquet/properties.h"
 #include "parquet/types.h"
 
-using parquet::Encryption;
+using parquet::ParquetCipher;
 using parquet::EncryptionProperties;
 
 namespace parquet_encryption {
 
+// Module types
+const int8_t Footer = 0;
+const int8_t ColumnMetaData = 1;
+const int8_t DataPage = 2;
+const int8_t DictionaryPage = 3;
+const int8_t DataPageHeader = 4;
+const int8_t DictionaryPageHeader = 5;
+const int8_t ColumnIndex = 6;
+const int8_t OffsetIndex = 7;
+
 int SignedFooterEncrypt(const uint8_t* plaintext,
             int plaintext_len, uint8_t* key, int key_len, uint8_t* aad, int aad_len,
             uint8_t* nonce, int nonce_len, uint8_t* ciphertext);
 
-int Encrypt(Encryption::type alg_id, bool metadata, const uint8_t* plaintext,
+int Encrypt(ParquetCipher::type alg_id, bool metadata, const uint8_t* plaintext,
             int plaintext_len, uint8_t* key, int key_len, uint8_t* aad, int aad_len,
             uint8_t* ciphertext);
 
 int Encrypt(std::shared_ptr<EncryptionProperties> encryption_props, bool metadata,
             const uint8_t* plaintext, int plaintext_len, uint8_t* ciphertext);
 
-int Decrypt(Encryption::type alg_id, bool metadata, const uint8_t* ciphertext,
+int Decrypt(ParquetCipher::type alg_id, bool metadata, const uint8_t* ciphertext,
             int ciphertext_len, uint8_t* key, int key_len, uint8_t* aad, int aad_len,
             uint8_t* plaintext);
 
 int Decrypt(std::shared_ptr<EncryptionProperties> encryption_props, bool metadata,
             const uint8_t* ciphertext, int ciphertext_len, uint8_t* plaintext);
+
+std::string createModuleAAD(const std::string& fileAAD, int8_t module_type,
+			    int16_t row_group_ordinal, int16_t column_ordinal,
+			    int16_t page_ordinal);
+
+std::string createFooterAAD(const std::string& aad_prefix_bytes);
+
+void quickUpdatePageAAD(const std::string &AAD, int16_t new_page_ordinal);
+
 }  // namespace parquet_encryption
 
 #endif  // PARQUET_UTIL_CRYPTO_H