Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/client/dbps_api_client.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -270,7 +270,8 @@ DecryptApiResponse DBPSApiClient::Decrypt(
CompressionCodec::type encrypted_compression,
const std::string& key_id,
const std::string& user_id,
const std::string& application_context
const std::string& application_context,
const std::map<std::string, std::string>& encryption_metadata
) {
DecryptJsonRequest json_request;
json_request.column_name_ = column_name;
Expand All @@ -283,6 +284,7 @@ DecryptApiResponse DBPSApiClient::Decrypt(
json_request.key_id_ = key_id;
json_request.user_id_ = user_id;
json_request.application_context_ = application_context;
json_request.encryption_metadata_ = encryption_metadata;
json_request.reference_id_ = GenerateReferenceId();

DecryptApiResponse api_response;
Expand Down
3 changes: 2 additions & 1 deletion src/client/dbps_api_client.h
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,8 @@ class DBPSApiClient {
CompressionCodec::type encrypted_compression,
const std::string& key_id,
const std::string& user_id,
const std::string& application_context
const std::string& application_context,
const std::map<std::string, std::string>& encryption_metadata
);

private:
Expand Down
11 changes: 8 additions & 3 deletions src/client/dbps_api_client_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,7 @@ TEST(DBPSApiClient, DecryptWithValidData) {
"encryption": {"key_id": "test_key_123"},
"access": {"user_id": "test_user_456"},
"application_context": "{\"user_id\": \"test_user_456\"}",
"encryption_metadata": {},
"debug": {"reference_id": "1755831549871"}
})";

Expand Down Expand Up @@ -410,7 +411,8 @@ TEST(DBPSApiClient, DecryptWithValidData) {
CompressionCodec::UNCOMPRESSED, // encrypted_compression
"test_key_123", // key_id
"test_user_456", // user_id
"{\"user_id\": \"test_user_456\"}" // application_context
"{\"user_id\": \"test_user_456\"}", // application_context
std::map<std::string, std::string>{} // encryption_metadata
);

// Verify the response
Expand Down Expand Up @@ -486,7 +488,8 @@ TEST(DBPSApiClient, DecryptWithInvalidData) {
CompressionCodec::UNCOMPRESSED, // encrypted_compression
"test_key_123", // key_id
"test_user_456", // user_id
"{\"user_id\": \"test_user_456\"}" // application_context
"{\"user_id\": \"test_user_456\"}", // application_context
std::map<std::string, std::string>{} // encryption_metadata
);

// Verify the response indicates failure
Expand Down Expand Up @@ -581,6 +584,7 @@ TEST(DBPSApiClient, DecryptWithInvalidJsonResponse) {
"encryption": {"key_id": "test_key_123"},
"access": {"user_id": "test_user_456"},
"application_context": "{\"user_id\": \"test_user_456\"}",
"encryption_metadata": {},
"debug": {"reference_id": "1755831549871"}
})";

Expand Down Expand Up @@ -616,7 +620,8 @@ TEST(DBPSApiClient, DecryptWithInvalidJsonResponse) {
CompressionCodec::UNCOMPRESSED, // encrypted_compression
"test_key_123", // key_id
"test_user_456", // user_id
"{\"user_id\": \"test_user_456\"}" // application_context
"{\"user_id\": \"test_user_456\"}", // application_context
std::map<std::string, std::string>{} // encryption_metadata
);

// Verify the response indicates failure
Expand Down
15 changes: 10 additions & 5 deletions src/common/dbpa_interface.h
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,10 @@ class DBPS_EXPORT EncryptionResult {
// Success flag; false indicates an error.
virtual bool success() const = 0;

//TODO: revisit the default implementation.
virtual const std::optional<std::map<std::string, std::string>> encryption_metadata() const {
return std::nullopt;
}
// Encryption metadata (valid when success() == true)
// Map of string key-value pairs containing any extra parameters used during encryption that are needed for decryption,
// for example, the encryption_algorithm_version used.
virtual const std::optional<std::map<std::string, std::string>> encryption_metadata() const = 0;
Comment thread
argmarco-tkd marked this conversation as resolved.

// Error details (valid when success() == false).
virtual const std::string& error_message() const = 0;
Expand Down Expand Up @@ -137,7 +137,12 @@ class DBPS_EXPORT DataBatchProtectionAgentInterface {
span<const uint8_t> ciphertext,
std::map<std::string, std::string> encoding_attributes) = 0;

virtual const std::optional<std::map<std::string, std::string>> EncryptionMetadata() const {
/* Returns the encryption metadata provided during the class init() call.
* The encryption metadata is a map of string key-value pairs and is defined only for Decrypt usage.
* This metadata map is the one returned by the EncryptionResult.encryption_metadata() during the Encrypt call and indicates
* any extra parameters used during encryption that are needed for decryption, for example, the encryption_algorithm_version used.
*/
virtual const std::optional<std::map<std::string, std::string>> EncryptionMetadata() const {
return column_encryption_metadata_;
}

Expand Down
1 change: 1 addition & 0 deletions src/common/dbpa_interface_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ class MockEncryptionResult : public EncryptionResult {

std::size_t size() const override { return data_.size(); }
bool success() const override { return success_; }
const std::optional<std::map<std::string, std::string>> encryption_metadata() const override { return std::nullopt; }
const std::string& error_message() const override { static std::string empty; return empty; }
const std::map<std::string, std::string>& error_fields() const override { static std::map<std::string, std::string> empty; return empty; }
};
Expand Down
20 changes: 14 additions & 6 deletions src/common/dbpa_local.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,10 @@ using namespace dbps::enum_utils;

// LocalEncryptionResult implementation

LocalEncryptionResult::LocalEncryptionResult(std::vector<uint8_t> ciphertext)
: ciphertext_(std::move(ciphertext)), success_(true) {
LocalEncryptionResult::LocalEncryptionResult(std::vector<uint8_t> ciphertext, const std::map<std::string, std::string>& encryption_metadata)
: ciphertext_(std::move(ciphertext)),
success_(true),
encryption_metadata_(encryption_metadata) {
}

LocalEncryptionResult::LocalEncryptionResult(const std::string& error_stage, const std::string& error_message)
Expand All @@ -55,6 +57,10 @@ bool LocalEncryptionResult::success() const {
return success_;
}

const std::optional<std::map<std::string, std::string>> LocalEncryptionResult::encryption_metadata() const {
return encryption_metadata_.empty() ? std::nullopt : std::optional{encryption_metadata_};
}

const std::string& LocalEncryptionResult::error_message() const {
return error_message_;
}
Expand Down Expand Up @@ -191,7 +197,8 @@ std::unique_ptr<EncryptionResult> LocalDataBatchProtectionAgent::Encrypt(
compression_type_,
column_key_id_,
user_id_,
app_context_
app_context_,
{} // encryption_metadata, which is empty for the Encryption call.
);

// Convert plaintext span to vector for the sequencer
Expand All @@ -205,8 +212,8 @@ std::unique_ptr<EncryptionResult> LocalDataBatchProtectionAgent::Encrypt(
return std::make_unique<LocalEncryptionResult>(sequencer.error_stage_, sequencer.error_message_);
}

// Return successful result with encrypted data
return std::make_unique<LocalEncryptionResult>(std::move(sequencer.encrypted_result_));
// Return successful result with encrypted data and encryption_metadata
return std::make_unique<LocalEncryptionResult>(std::move(sequencer.encrypted_result_), sequencer.encryption_metadata_);
}

std::unique_ptr<DecryptionResult> LocalDataBatchProtectionAgent::Decrypt(
Expand Down Expand Up @@ -241,7 +248,8 @@ std::unique_ptr<DecryptionResult> LocalDataBatchProtectionAgent::Decrypt(
compression_type_,
column_key_id_,
user_id_,
app_context_
app_context_,
column_encryption_metadata_.value_or(std::map<std::string, std::string>{})
);

// Convert ciphertext span to vector for the sequencer
Expand Down
4 changes: 3 additions & 1 deletion src/common/dbpa_local.h
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ namespace dbps::external {
class DBPS_EXPORT LocalEncryptionResult : public EncryptionResult {
public:
// Constructor for successful encryption
LocalEncryptionResult(std::vector<uint8_t> ciphertext);
LocalEncryptionResult(std::vector<uint8_t> ciphertext, const std::map<std::string, std::string>& encryption_metadata = {});

// Constructor for failed encryption
LocalEncryptionResult(const std::string& error_stage, const std::string& error_message);
Expand All @@ -50,6 +50,7 @@ class DBPS_EXPORT LocalEncryptionResult : public EncryptionResult {
span<const uint8_t> ciphertext() const override;
std::size_t size() const override;
bool success() const override;
const std::optional<std::map<std::string, std::string>> encryption_metadata() const override;
const std::string& error_message() const override;
const std::map<std::string, std::string>& error_fields() const override;

Expand All @@ -58,6 +59,7 @@ class DBPS_EXPORT LocalEncryptionResult : public EncryptionResult {
private:
std::vector<uint8_t> ciphertext_;
bool success_;
std::map<std::string, std::string> encryption_metadata_;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if it should be in this header file or in another, but can you add a comment with an example of what may be contained in here?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added comments on the dbpa_interface.h file.

std::string error_message_;
std::map<std::string, std::string> error_fields_;
};
Expand Down
57 changes: 56 additions & 1 deletion src/common/dbpa_local_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@

using namespace dbps::external;

namespace {
// should rename to DBPS_ENCRYPTION_METADATA
const std::map<std::string, std::string> DBPS_ENCRYPTION_METADATA = {{"dbps_agent_version", "v0.01_unittest"}};
}

// Test fixture for LocalDataBatchProtectionAgent tests
class LocalDataBatchProtectionAgentTest : public ::testing::Test {
protected:
Expand Down Expand Up @@ -81,7 +86,7 @@ TEST_F(LocalDataBatchProtectionAgentTest, SuccessfulDecryption) {
std::string app_context = R"({"user_id": "test_user"})";

EXPECT_NO_THROW(agent.init("test_column", connection_config, app_context, "test_key",
Type::UNDEFINED, std::nullopt, CompressionCodec::UNCOMPRESSED, std::nullopt));
Type::UNDEFINED, std::nullopt, CompressionCodec::UNCOMPRESSED, DBPS_ENCRYPTION_METADATA));

std::vector<uint8_t> test_data = {1, 2, 3, 4};
std::map<std::string, std::string> encoding_attributes = {{"page_encoding", "PLAIN"}, {"page_type", "DICTIONARY_PAGE"}};
Expand All @@ -92,6 +97,56 @@ TEST_F(LocalDataBatchProtectionAgentTest, SuccessfulDecryption) {
EXPECT_GT(result->size(), 0);
}

// Test roundtrip encryption/decryption
TEST_F(LocalDataBatchProtectionAgentTest, RoundTripEncryptDecrypt) {
LocalDataBatchProtectionAgent encrypt_agent;

std::map<std::string, std::string> connection_config;
std::string app_context = R"({"user_id": "test_user"})";

EXPECT_NO_THROW(encrypt_agent.init("test_column", connection_config, app_context, "test_key",
Type::UNDEFINED, std::nullopt, CompressionCodec::UNCOMPRESSED, std::nullopt));

// Original data to encrypt
std::vector<uint8_t> original_data = {1, 2, 3, 4, 5};
std::map<std::string, std::string> encoding_attributes = {{"page_encoding", "PLAIN"}, {"page_type", "DICTIONARY_PAGE"}};

// Encrypt the data
auto encrypt_result = encrypt_agent.Encrypt(original_data, encoding_attributes);

ASSERT_NE(encrypt_result, nullptr);
ASSERT_TRUE(encrypt_result->success());
EXPECT_GT(encrypt_result->size(), 0);

// Verify encryption_metadata is present in the result
auto encryption_metadata = encrypt_result->encryption_metadata();
ASSERT_TRUE(encryption_metadata.has_value());
ASSERT_EQ(1, encryption_metadata->size());
ASSERT_TRUE(encryption_metadata->at("dbps_agent_version").length() > 0); // Non-empty string

// Get the ciphertext
auto ciphertext_span = encrypt_result->ciphertext();
std::vector<uint8_t> ciphertext(ciphertext_span.begin(), ciphertext_span.end());

// Create a new agent for decryption with the encryption_metadata from the encryption result
LocalDataBatchProtectionAgent decrypt_agent;
EXPECT_NO_THROW(decrypt_agent.init("test_column", connection_config, app_context, "test_key",
Type::UNDEFINED, std::nullopt, CompressionCodec::UNCOMPRESSED, encryption_metadata));

// Decrypt the ciphertext
auto decrypt_result = decrypt_agent.Decrypt(ciphertext, encoding_attributes);

ASSERT_NE(decrypt_result, nullptr);
ASSERT_TRUE(decrypt_result->success());

// Verify the decrypted data matches the original
auto plaintext_span = decrypt_result->plaintext();
std::vector<uint8_t> decrypted_data(plaintext_span.begin(), plaintext_span.end());

ASSERT_EQ(original_data.size(), decrypted_data.size());
EXPECT_EQ(original_data, decrypted_data);
}

// Test encryption without initialization
TEST_F(LocalDataBatchProtectionAgentTest, EncryptWithoutInit) {
LocalDataBatchProtectionAgent agent;
Expand Down
47 changes: 32 additions & 15 deletions src/common/dbpa_remote.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -47,24 +47,39 @@ bool RemoteEncryptionResult::success() const {
return response_ && response_->Success();
}

const std::optional<std::map<std::string, std::string>> RemoteEncryptionResult::encryption_metadata() const {
if (!parsed_encryption_metadata_.has_value() && response_ && response_->Success()) {
const auto& response_attrs = response_->GetResponseAttributes();
// For the RemoteEncryptionResult, encryption_metadata_ is forwarded from the API response.
// The attribute in the API response is also named encryption_metadata.
// If the API reponse is empty, the parsed value is set to NULL for compatibility with the class interface.
if (!response_attrs.encryption_metadata_.empty()) {
parsed_encryption_metadata_ = response_attrs.encryption_metadata_;
} else {
parsed_encryption_metadata_ = std::nullopt;
}
Comment on lines +52 to +60

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's document what's going on here - it's a bit hard to grasp

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. You can let me know if the comment is sufficient and I can expand more if needed.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's a clearer, thanks - but the whole concept of 'cached metadata' is a bit confusing. maybe let's add a short note on that? (in particular, why is it 'cached' ?)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cached was probably a bad choice, so changed the variable names to parsed. I think that captures better the usage.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thank you! it's much clearer now.

}
return parsed_encryption_metadata_;
}

const std::string& RemoteEncryptionResult::error_message() const {
if (cached_error_message_.empty() && response_) {
if (parsed_error_message_.empty() && response_) {
if (!response_->Success()) {
cached_error_message_ = response_->ErrorMessage();
parsed_error_message_ = response_->ErrorMessage();
} else {
cached_error_message_ = "Successful encryption";
parsed_error_message_ = "Successful encryption";
}
}
return cached_error_message_;
return parsed_error_message_;
}

const std::map<std::string, std::string>& RemoteEncryptionResult::error_fields() const {
if (cached_error_fields_.empty() && response_) {
if (parsed_error_fields_.empty() && response_) {
if (!response_->Success()) {
cached_error_fields_ = response_->ErrorFields();
parsed_error_fields_ = response_->ErrorFields();
}
}
return cached_error_fields_;
return parsed_error_fields_;
}

RemoteDecryptionResult::RemoteDecryptionResult(std::unique_ptr<DecryptApiResponse> response)
Expand All @@ -90,23 +105,23 @@ bool RemoteDecryptionResult::success() const {
}

const std::string& RemoteDecryptionResult::error_message() const {
if (cached_error_message_.empty() && response_) {
if (parsed_error_message_.empty() && response_) {
if (!response_->Success()) {
cached_error_message_ = response_->ErrorMessage();
parsed_error_message_ = response_->ErrorMessage();
} else {
cached_error_message_ = "Successful decryption";
parsed_error_message_ = "Successful decryption";
}
}
return cached_error_message_;
return parsed_error_message_;
}

const std::map<std::string, std::string>& RemoteDecryptionResult::error_fields() const {
if (cached_error_fields_.empty() && response_) {
if (parsed_error_fields_.empty() && response_) {
if (!response_->Success()) {
cached_error_fields_ = response_->ErrorFields();
parsed_error_fields_ = response_->ErrorFields();
}
}
return cached_error_fields_;
return parsed_error_fields_;
}

// Helper functions for validating that fields of the request <> response match.
Expand Down Expand Up @@ -318,7 +333,9 @@ std::unique_ptr<DecryptionResult> RemoteDataBatchProtectionAgent::Decrypt(span<c
compression_type_,
column_key_id_,
user_id_,
app_context_
app_context_,
// For the Decrypt use case, column_encryption_metadata_ will be set on the constructor, however adding a default empty map for type safety.
column_encryption_metadata_.value_or(std::map<std::string, std::string>{})
);

// Validate that response fields match request fields
Expand Down
14 changes: 8 additions & 6 deletions src/common/dbpa_remote.h
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ class DBPS_EXPORT RemoteEncryptionResult : public EncryptionResult {
span<const uint8_t> ciphertext() const override;
std::size_t size() const override;
bool success() const override;
const std::optional<std::map<std::string, std::string>> encryption_metadata() const override;
const std::string& error_message() const override;
const std::map<std::string, std::string>& error_fields() const override;

Expand All @@ -57,11 +58,12 @@ class DBPS_EXPORT RemoteEncryptionResult : public EncryptionResult {
private:
std::unique_ptr<EncryptApiResponse> response_;

// Cache error message/fields to avoid repeated parsing of API response.
// Stored error message/fields/encryption_metadata to avoid repeated parsing of API response.
// Defined as mutable because it has lazy evaluation (updated only once on the getter methods)
// The caching of this is possible because the API response `response_` doesn't change after construction.
mutable std::string cached_error_message_;
mutable std::map<std::string, std::string> cached_error_fields_;
mutable std::string parsed_error_message_;
mutable std::map<std::string, std::string> parsed_error_fields_;
mutable std::optional<std::map<std::string, std::string>> parsed_encryption_metadata_;
};

/**
Expand All @@ -84,10 +86,10 @@ class DBPS_EXPORT RemoteDecryptionResult : public DecryptionResult {
private:
std::unique_ptr<DecryptApiResponse> response_;

// Cached error message/fields to avoid repeated parsing of API response.
// Stored error message/fields to avoid repeated parsing of API response.
// Defined as mutable because it has lazy evaluation (updated only once on the getter methods)
mutable std::string cached_error_message_;
mutable std::map<std::string, std::string> cached_error_fields_;
mutable std::string parsed_error_message_;
mutable std::map<std::string, std::string> parsed_error_fields_;
};

/**
Expand Down
Loading