Refactoring BasicXorEncryptor to make more readable.

src/processing/encryptors/basic_xor_encryptor.cpp

@@ -31,12 +31,12 @@ using namespace dbps::external;
 // Functions for encrypting and decrypting byte arrays.
 // ---------------------------------------------------------------------------
 
-static std::vector<uint8_t> EncryptByteArray(tcb::span<const uint8_t> data, const std::string& key_id) {
+std::vector<uint8_t> BasicXorEncryptor::XorEncrypt(tcb::span<const uint8_t> data, const std::string& key_id) {
     if (data.empty()) {
         return {};
     }
     if (key_id.empty()) {
-        throw std::invalid_argument("EncryptByteArray: key must not be empty for non-empty data");
+        throw std::invalid_argument("XorEncrypt: key must not be empty for non-empty data");
     }
     std::vector<uint8_t> out(data.size());
     std::hash<std::string> hasher;
@@ -48,20 +48,20 @@ static std::vector<uint8_t> EncryptByteArray(tcb::span<const uint8_t> data, cons
     return out;
 }
 
-static std::vector<uint8_t> DecryptByteArray(tcb::span<const uint8_t> data, const std::string& key_id) {
-    return EncryptByteArray(data, key_id);
+std::vector<uint8_t> BasicXorEncryptor::XorDecrypt(tcb::span<const uint8_t> data, const std::string& key_id) {
+    return XorEncrypt(data, key_id);
 }
 
 // ---------------------------------------------------------------------------
 // Block encryption
 // ---------------------------------------------------------------------------
 
 std::vector<uint8_t> BasicXorEncryptor::EncryptBlock(tcb::span<const uint8_t> data) {
-    return EncryptByteArray(data, key_id_);
+    return XorEncrypt(data, key_id_);
 }
 
 std::vector<uint8_t> BasicXorEncryptor::DecryptBlock(tcb::span<const uint8_t> data) {
-    return DecryptByteArray(data, key_id_);
+    return XorDecrypt(data, key_id_);
 }
 
 // ---------------------------------------------------------------------------
@@ -71,79 +71,120 @@ std::vector<uint8_t> BasicXorEncryptor::DecryptBlock(tcb::span<const uint8_t> da
 //   Fixed:    [0x01][uint32 count][uint32 elem_size] <contiguous encrypted elements>
 //   Variable: [0x00][uint32 count]                   <length-prefixed encrypted elements>
 //
+// NOTE ON TYPE-SPECIFIC ENCRYPTION:
+//
+// While this version of EncryptValueList is aware of types, the implementation ignores the type
+// of the TypeValuesBuffer. It encrypts each element by interpreting the elements a plain bytes,
+// not utilizing the specific type of the buffer elements (INT32, INT64, etc.).
+// 
+// A more sophisticated implementation could take advantage of the specific element type
+// to call type-specific encryption functions, if the underlying library supports them.
+//
+// Further, a context-aware encryptor can additionally take advantage of the app_context, 
+// user_id, and column_name to further customize the encryption process, refined access control, etc.
+//
+// For example, a type-specific encryption per element could look like:
+//   assert(input_buffer.IsTypedBufferI32());
+//   size_t i = 0;
+//   for (const int32_t element : input_buffer) {    // e.g. int32_t for TypedBufferI32
+//       auto encrypted = library.EncryptInt32(element, key_id);
+//       output_buffer.SetElement(i++, encrypted);
+//   }
+//
 // ---------------------------------------------------------------------------
 
-std::vector<uint8_t> BasicXorEncryptor::EncryptValueList(
-    const TypedValuesBuffer& typed_buffer) {
+template <typename TypedBuffer>
+std::vector<uint8_t> BasicXorEncryptor::EncryptTypedElements(
+    const TypedBuffer& input_buffer, const std::string& key_id) {
+    constexpr bool is_fixed = TypedBuffer::is_fixed_sized;
+    constexpr size_t prefix_length = is_fixed ? kFixedHeaderLength : kVariableHeaderLength;
+    const size_t num_elements = input_buffer.GetNumElements();
 
-    std::cout << "EncryptValueList context: column=" << column_name_
-              << " user=" << user_id_ << " key=" << key_id_
-              << " datatype=" << dbps::enum_utils::to_string(datatype_) << std::endl;
-
-    // TODO: Make EncryptValueList/DecryptValueList more readable:
-    // - Implement a simpler way of accessing Variant TypedBuffer if possible.
-    // - Replace unnecesary use of lambdas and removing other unnecessary complexity.
-    return std::visit([&](const auto& input_buffer) -> std::vector<uint8_t> {
-        using BufferType = std::decay_t<decltype(input_buffer)>;
-        constexpr bool is_fixed = BufferType::is_fixed_sized;
-        const size_t num_elements = input_buffer.GetNumElements();
-        constexpr size_t prefix_length = is_fixed ? kFixedHeaderLength : kVariableHeaderLength;
-
-        // Empty buffer, return empty vector with header.
-        if (num_elements == 0) {
-            std::vector<uint8_t> result(prefix_length);
-            WriteHeader(result, {is_fixed, 0, 0});
-            return result;
-        }
+    // If there are no elements, return an empty buffer with the header.
+    if (num_elements == 0) {
+        std::vector<uint8_t> result(prefix_length);
+        WriteHeader(result, {is_fixed, 0, 0});
+        return result;
+    }
 
-        auto encrypt_into = [&](auto& output_buffer) -> std::vector<uint8_t> {
-            size_t output_index = 0;
-            for (const auto raw_bytes : input_buffer.raw_elements()) {
-                auto encrypted = EncryptByteArray(raw_bytes, key_id_);
-                output_buffer.SetElement(output_index, tcb::span<const uint8_t>(encrypted));
-                output_index++;
-            }
-            return output_buffer.FinalizeAndTakeBuffer();
-        };
-
-        std::vector<uint8_t> final_buffer;
-        size_t element_size = 0;
-        if constexpr (is_fixed) {
-            element_size = input_buffer.GetElementSize();
-            TypedBufferRawBytesFixedSized output_buffer{
-                num_elements, prefix_length, RawBytesFixedSizedCodec{element_size}};
-            final_buffer = encrypt_into(output_buffer);
-        } else {
-            auto reserved_bytes_hint = input_buffer.GetRawBufferSize();
-            TypedBufferRawBytesVariableSized output_buffer{
-                num_elements, reserved_bytes_hint, true, prefix_length};
-            final_buffer = encrypt_into(output_buffer);
+    // Encrypt the elements by traversing the typed input buffer and encrypt each element separately:
+    // - Create an output raw-bytes buffer to capture the encrypted elements as bytes.
+    // - For each element: read its raw bytes, encrypt them, write into the output buffer.
+    // - Finalize the output buffer into a contiguous byte vector.
+
+    std::vector<uint8_t> final_buffer;
+    size_t element_size = 0;
+
+    // Encrypt fixed-size elements
+    if constexpr (is_fixed) {
+        element_size = input_buffer.GetElementSize();
+        TypedBufferRawBytesFixedSized output_buffer{
+            num_elements, prefix_length, RawBytesFixedSizedCodec{element_size}};
+        size_t output_index = 0;
+        for (const auto raw_bytes : input_buffer.raw_elements()) {
+            auto encrypted = XorEncrypt(raw_bytes, key_id);
+            output_buffer.SetElement(output_index, tcb::span<const uint8_t>(encrypted));
+            output_index++;
         }
-        WriteHeader(final_buffer, {is_fixed,
-            static_cast<uint32_t>(num_elements),
-            static_cast<uint32_t>(element_size)});
-        return final_buffer;
+        final_buffer = output_buffer.FinalizeAndTakeBuffer();
+    }
+
+    // Encrypt variable-size elements
+    else {
+        auto reserved_bytes_hint = input_buffer.GetRawBufferSize();
+        TypedBufferRawBytesVariableSized output_buffer{
+            num_elements, reserved_bytes_hint, true, prefix_length};
+        size_t output_index = 0;
+        for (const auto raw_bytes : input_buffer.raw_elements()) {
+            auto encrypted = XorEncrypt(raw_bytes, key_id);
+            output_buffer.SetElement(output_index, tcb::span<const uint8_t>(encrypted));
+            output_index++;
+        }
+        final_buffer = output_buffer.FinalizeAndTakeBuffer();
+    }
 
+    // Write the header to the final buffer and return it.
+    WriteHeader(final_buffer, {is_fixed,
+        static_cast<uint32_t>(num_elements),
+        static_cast<uint32_t>(element_size)});
+    return final_buffer;
+}
+
+std::vector<uint8_t> BasicXorEncryptor::EncryptValueList(
+    const TypedValuesBuffer& typed_buffer) {
+
+    // Printable context string for logging
+    // std::string context_str = std::string("Context parameters:")
+    //    + "\n  column_name: " + column_name_
+    //    + "\n  user_id: " + user_id_
+    //    + "\n  key_id: " + key_id_
+    //    + "\n  application_context: " + application_context_
+    //    + "\n  datatype: " + std::string(dbps::enum_utils::to_string(datatype_));
+
+    // Printable typed buffer string for logging -- Commented out by default to avoid performance impact.
+    // std::string typed_buffer_str = PrintableTypedValuesBuffer(typed_buffer);
+
+    // std::visit extracts the concrete buffer type from the TypedValuesBuffer variant
+    // and forwards it to EncryptTypedElements, which handles all buffer types generically.
+    return std::visit([&](const auto& input_buffer) {
+        return EncryptTypedElements(input_buffer, key_id_);
     }, typed_buffer);
 }
 
 // ---------------------------------------------------------------------------
 // Value-level decryption  (bytes in -> TypedValuesBuffer out)
 //
-// Parses the header, then wraps the full span (with prefix_size) as a
-// TypedBufferRawBytes... read buffer so the buffer skips the header
-// automatically.  Output is the correctly-typed buffer matching datatype_.
+// Parses the header, then creates a correctly-typed buffer matching datatype_ (INT32, INT64, etc.).
+// Calls the byte-array decryptor to decrypt the elements into the typed buffer. 
 // ---------------------------------------------------------------------------
 
-// Helper function to decrypt fixed-size buffer into a specific output TypedBuffer type.
-template <typename OutputBuffer>
-static OutputBuffer DecryptFixedIntoBuffer(
-    const TypedBufferRawBytesFixedSized& encrypted_buffer,
-    const std::string& key_id,
-    OutputBuffer output_buffer) {
+// Helper function to decrypt fixed-size elements into the output TypedBuffer type.
+template <typename TypedBuffer>
+TypedBuffer BasicXorEncryptor::DecryptFixedSizedElementsIntoTypedBuffer(
+    const TypedBufferRawBytesFixedSized& encrypted_buffer, const std::string& key_id, TypedBuffer output_buffer) {
     size_t output_index = 0;
     for (const auto raw_bytes : encrypted_buffer.raw_elements()) {
-        auto decrypted_bytes = DecryptByteArray(raw_bytes, key_id);
+        auto decrypted_bytes = XorDecrypt(raw_bytes, key_id);
         output_buffer.SetRawElement(output_index, tcb::span<const uint8_t>(decrypted_bytes));
         output_index++;
     }
@@ -153,57 +194,57 @@ static OutputBuffer DecryptFixedIntoBuffer(
 TypedValuesBuffer BasicXorEncryptor::DecryptValueList(
     tcb::span<const uint8_t> encrypted_bytes) {
 
-    // During decryption, the number of elements can be obtained from the header,
-    // no need to calculate it from the encrypted_buffer.
     auto header = ReadHeader(encrypted_bytes);
     auto num_elements = static_cast<size_t>(header.num_elements);
 
+    // Decrypt fixed-size elements
     if (header.is_fixed) {
+
+        // Create a fixed-sized byte buffer for reading the encrypted elements.
         TypedBufferRawBytesFixedSized encrypted_buffer{
-            encrypted_bytes, kFixedHeaderLength,
-            RawBytesFixedSizedCodec{header.element_size}};
+            encrypted_bytes, kFixedHeaderLength, RawBytesFixedSizedCodec{header.element_size}};
 
-        // TODO: This is leaking Parquet-specific types into the encryptor, which should be agnostic of Parquet.
-        // This is needed because on the returned bytes we are not saving a type information.
-        // We could annotate the generating bytes by simply updating the 1st byte of the header to indicate the type.
+        // Populate a typed buffer with the decrypted elements in the corresponding type.
         switch (datatype_) {
             case Type::INT32:
-                return DecryptFixedIntoBuffer(encrypted_buffer, key_id_, TypedBufferI32{num_elements});
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_, TypedBufferI32{num_elements});
             case Type::INT64:
-                return DecryptFixedIntoBuffer(encrypted_buffer, key_id_, TypedBufferI64{num_elements});
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_, TypedBufferI64{num_elements});
             case Type::INT96:
-                return DecryptFixedIntoBuffer(encrypted_buffer, key_id_, TypedBufferInt96{num_elements});
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_, TypedBufferInt96{num_elements});
             case Type::FLOAT:
-                return DecryptFixedIntoBuffer(encrypted_buffer, key_id_, TypedBufferFloat{num_elements});
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_, TypedBufferFloat{num_elements});
             case Type::DOUBLE:
-                return DecryptFixedIntoBuffer(encrypted_buffer, key_id_, TypedBufferDouble{num_elements});
-            case Type::FIXED_LEN_BYTE_ARRAY: {
-                TypedBufferRawBytesFixedSized output_buffer{
-                    num_elements, 0, RawBytesFixedSizedCodec{header.element_size}};
-                size_t output_index = 0;
-                for (const auto element : encrypted_buffer) {
-                    auto decrypted_bytes = DecryptByteArray(element, key_id_);
-                    output_buffer.SetElement(output_index, tcb::span<const uint8_t>(decrypted_bytes));
-                    output_index++;
-                }
-                return output_buffer;
-            }
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_, TypedBufferDouble{num_elements});
+            case Type::FIXED_LEN_BYTE_ARRAY:
+                return DecryptFixedSizedElementsIntoTypedBuffer(
+                    encrypted_buffer, key_id_,
+                    TypedBufferRawBytesFixedSized{num_elements, 0, RawBytesFixedSizedCodec{header.element_size}});
             default:
                 throw InvalidInputException(
                     std::string("DecryptValueList: unsupported fixed-size datatype: ")
                     + std::string(dbps::enum_utils::to_string(datatype_)));
         }
-    } else {
-        TypedBufferRawBytesVariableSized encrypted_buffer{
-            encrypted_bytes, kVariableHeaderLength};
+    } 
+
+    // Decrypt variable-size elements
+    else {
+        // Create a variable-sized byte buffer for reading the encrypted elements.
+        TypedBufferRawBytesVariableSized encrypted_buffer{ encrypted_bytes, kVariableHeaderLength};
 
         switch (datatype_) {
+            // Create a BYTE-ARRAY typed buffer for storing the decrypted elements.
             case Type::BYTE_ARRAY: {
                 auto reserved_bytes_hint = encrypted_buffer.GetRawBufferSize();
                 TypedBufferRawBytesVariableSized output_buffer{num_elements, reserved_bytes_hint, true};
                 size_t output_index = 0;
                 for (const auto element : encrypted_buffer) {
-                    auto decrypted_bytes = DecryptByteArray(element, key_id_);
+                    auto decrypted_bytes = XorDecrypt(element, key_id_);
                     output_buffer.SetElement(output_index, tcb::span<const uint8_t>(decrypted_bytes));
                     output_index++;
                 }

src/processing/encryptors/basic_xor_encryptor.h

@@ -61,5 +61,18 @@ class DBPS_EXPORT BasicXorEncryptor : public DBPSEncryptor {
     std::vector<uint8_t> EncryptValueList(const TypedValuesBuffer& typed_buffer) override;
 
     TypedValuesBuffer DecryptValueList(tcb::span<const uint8_t> encrypted_bytes) override;
+
+private:
+    static std::vector<uint8_t> XorEncrypt(tcb::span<const uint8_t> data, const std::string& key_id);
+    static std::vector<uint8_t> XorDecrypt(tcb::span<const uint8_t> data, const std::string& key_id);
+
+    template <typename InputBuffer>
+    static std::vector<uint8_t> EncryptTypedElements(
+        const InputBuffer& input_buffer, const std::string& key_id);
+
+    template <typename TypedBuffer>
+    static TypedBuffer DecryptFixedSizedElementsIntoTypedBuffer(
+        const TypedBufferRawBytesFixedSized& encrypted_buffer,
+        const std::string& key_id, TypedBuffer output_buffer);
 };
 

src/processing/typed_buffer.h

@@ -41,6 +41,7 @@ class ByteBuffer {
 public:
     using value_type = typename Codec::value_type;
     static constexpr bool is_fixed_sized = Codec::is_fixed_sized;
+    static constexpr std::string_view type_name() noexcept { return Codec::type_name(); }
 
     // Constructor for read-only buffer for both fixed-sized and variable-sized elements.
     // Elements are stored contiguously in the span.

src/processing/typed_buffer_values.h

@@ -17,6 +17,8 @@
 
 #pragma once
 
+#include <sstream>
+#include <string>
 #include <variant>
 #include "typed_buffer_codecs.h"
 #include "typed_buffer.h"
@@ -56,4 +58,36 @@ using TypedValuesBuffer = std::variant<
     TypedBufferRawBytesFixedSized,
     TypedBufferRawBytesVariableSized
 >;
+
+// Printable string representation of the typed buffer
+inline std::string PrintableTypedValuesBuffer(const TypedValuesBuffer& buffer) {
+    return std::visit([](const auto& typed_buffer) -> std::string {
+        using BufferType = std::decay_t<decltype(typed_buffer)>;
+        using ValueType = typename BufferType::value_type;
+
+        std::ostringstream out;
+        const size_t num_elements = typed_buffer.GetNumElements();
+
+        out << BufferType::type_name() << " (" << num_elements << " elements):\n";
+
+        size_t i = 0;
+        for (const auto element : typed_buffer) {
+            if constexpr (std::is_same_v<ValueType, Int96>) {
+                out << "  [" << i << "] [" << element.lo << ", "
+                    << element.mid << ", " << element.hi << "]\n";
+            } else if constexpr (std::is_same_v<ValueType, std::string_view>) {
+                out << "  [" << i << "] \"" << element
+                    << "\" (length: " << element.size() << ")\n";
+            } else if constexpr (std::is_same_v<ValueType, tcb::span<const uint8_t>>) {
+                out << "  [" << i << "] <" << element.size() << " bytes>\n";
+            } else {
+                out << "  [" << i << "] " << element << "\n";
+            }
+            ++i;
+        }
+
+        return out.str();
+    }, buffer);
+}
+
 } // namespace dbps::processing