Merge pull request #15 from ccfelius/secrets

Adding local state, quick fixes for various bugs and new tests

src/core/crypto/crypto_primitives.cpp

@@ -86,6 +86,8 @@ void AESStateSSL::SetEncryptionAlgorithm(string_t s_algorithm) {
 
 void AESStateSSL::GenerateRandomData(data_ptr_t data, idx_t len) {
   // generate random bytes for nonce
+  unsigned char seed[] = "my_custom_seed";
+  RAND_seed(seed, sizeof(seed));
   RAND_bytes(data, len);
 }
 

src/core/functions/common.cpp

@@ -0,0 +1,34 @@
+#include "simple_encryption/common.hpp"
+#include "simple_encryption/core/functions/common.hpp"
+
+namespace simple_encryption {
+
+namespace core {
+
+SimpleEncryptionFunctionLocalState::SimpleEncryptionFunctionLocalState(ClientContext &context) : arena(BufferAllocator::Get(context)) {
+}
+
+unique_ptr<FunctionLocalState>
+SimpleEncryptionFunctionLocalState::Init(ExpressionState &state, const BoundFunctionExpression &expr, FunctionData *bind_data) {
+  return make_uniq<SimpleEncryptionFunctionLocalState>(state.GetContext());
+}
+
+unique_ptr<FunctionLocalState> SimpleEncryptionFunctionLocalState::InitCast(CastLocalStateParameters &parameters) {
+  return make_uniq<SimpleEncryptionFunctionLocalState>(*parameters.context.get());
+}
+
+SimpleEncryptionFunctionLocalState &SimpleEncryptionFunctionLocalState::ResetAndGet(CastParameters &parameters) {
+  auto &local_state = parameters.local_state->Cast<SimpleEncryptionFunctionLocalState>();
+  local_state.arena.Reset();
+  return local_state;
+}
+
+SimpleEncryptionFunctionLocalState &SimpleEncryptionFunctionLocalState::ResetAndGet(ExpressionState &state) {
+  auto &local_state = ExecuteFunctionState::GetFunctionState(state)->Cast<SimpleEncryptionFunctionLocalState>();
+  local_state.arena.Reset();
+  return local_state;
+}
+
+} // namespace core
+
+} // namespace simple_encryption

src/core/functions/scalar/encrypt_to_etype.cpp

@@ -61,6 +61,7 @@ ProcessAndCastEncrypt(shared_ptr<EncryptionState> encryption_state,
 
   // convert to Base64 into a newly allocated string in the result vector
   T base64_data = StringVector::EmptyString(*result_vector, base64_size);
+  memset(base64_data.GetDataWriteable(), 0, 12);
   Blob::ToBase64(encrypted_data, base64_data.GetDataWriteable());
 
   return base64_data;
@@ -113,7 +114,6 @@ shared_ptr<SimpleEncryptionState>
 GetSimpleEncryptionState(ExpressionState &state) {
 
   auto &info = GetEncryptionBindInfo(state);
-
   return info.context.registered_state->Get<SimpleEncryptionState>(
       "simple_encryption");
 }
@@ -309,6 +309,8 @@ void EncryptToEtype(LogicalType result_struct, Vector &input_vector,
   Vector struct_vector(result_struct, size);
   result.ReferenceAndSetType(struct_vector);
 
+//  ValidityMask &result_validity = FlatVector::Validity(result);
+
   if ((simple_encryption_state->counter == 0) || (HasSpace(simple_encryption_state, size) == false)) {
     // generate new random IV and reset counter (if strart or if there is no space left)
     SetIV(simple_encryption_state);
@@ -345,6 +347,7 @@ void EncryptToEtype(LogicalType result_struct, Vector &input_vector,
         return ENCRYPTED_TYPE{simple_encryption_state->iv[0],
                               simple_encryption_state->iv[1], encrypted_data};
       });
+
 }
 
 

src/include/simple_encryption/core/functions/common.hpp

@@ -0,0 +1,24 @@
+#pragma once
+#include "simple_encryption/common.hpp"
+
+namespace simple_encryption {
+
+namespace core {
+
+struct SimpleEncryptionFunctionLocalState : FunctionLocalState {
+public:
+
+  ArenaAllocator arena;
+
+public:
+  explicit SimpleEncryptionFunctionLocalState(ClientContext &context);
+  static unique_ptr<FunctionLocalState> Init(ExpressionState &state, const BoundFunctionExpression &expr,
+                                             FunctionData *bind_data);
+  static unique_ptr<FunctionLocalState> InitCast(CastLocalStateParameters &context);
+  static SimpleEncryptionFunctionLocalState &ResetAndGet(ExpressionState &state);
+  static SimpleEncryptionFunctionLocalState &ResetAndGet(CastParameters &parameters);
+};
+
+} // namespace core
+
+} // namespace simple_encryption

src/simple_encryption_state.cpp

@@ -1,4 +1,4 @@
-#define MAX_BUFFER_SIZE 1024
+#define MAX_BUFFER_SIZE 128
 #include "include/simple_encryption_state.hpp"
 #include "duckdb.hpp"
 #include "mbedtls_wrapper.hpp"
@@ -30,8 +30,7 @@ SimpleEncryptionState::SimpleEncryptionState(shared_ptr<ClientContext> context)
 
   // allocate encryption buffer
   // maybe do this in a better way (i.e. use buffer manager?)
-  uint8_t encryption_buffer[MAX_BUFFER_SIZE];
-  buffer_p = encryption_buffer;
+  buffer_p = static_cast<uint8_t *>(duckdb_malloc(MAX_BUFFER_SIZE));
 
   // clear the iv
   iv[0] = iv[1] = 0;

test/sql/encryption/bulk_encryption.test

@@ -1,28 +1,35 @@
-# name: test/sql/bulk_encryption.test
-# description: test simple__encryption extension
-# group: [simple_encryption]
+# name: test/sql/secrets/secrets_encryption.test
+# description: Test secret creation for internal encryption
+# group: [simple-encryption/secrets]
 
-# Require statement will ensure this test is run with this extension loaded
 require simple_encryption
 
-# So, when we do  more then 2048 values (i.e. vector size) it crashes
+# Ensure any currently stored secrets don't interfere with the test
 statement ok
-CREATE TABLE test_1 AS SELECT 1 AS value FROM range(2048);
+set allow_persistent_secrets=false;
 
+# Create an internal secret (for internal encryption of columns)
 statement ok
-SELECT encrypt(value, '0123456789112345') AS encrypted_value FROM test_1;
+CREATE SECRET key_1 (
+	TYPE ENCRYPTION,
+    MASTER_KEY '0123456789112345',
+    LENGTH 16
+);
 
 statement ok
-ALTER TABLE test_1 ADD COLUMN encrypted_values STRUCT(nonce_hi UBIGINT, nonce_lo UBIGINT, value INTEGER);
+CREATE TABLE test_1 AS SELECT 1 AS value FROM range(10000);
+
+statement ok
+ALTER TABLE test_1 ADD COLUMN encrypted_values STRUCT(nonce_hi UBIGINT, nonce_lo UBIGINT, value INTEGER) DEFAULT (STRUCT_PACK(nonce_hi := 0, nonce_lo := 0, value := 0));
 
 statement ok
 ALTER TABLE test_1 ADD COLUMN decrypted_values INTEGER;
 
 statement ok
-UPDATE test_1 SET encrypted_values = encrypt(value, '0123456789112345');
+UPDATE test_1 SET encrypted_values = encrypt(value, 'key_1', 'random_message');
 
 statement ok
-UPDATE test_1 SET decrypted_values = decrypt(encrypted_values, '0123456789112345');
+UPDATE test_1 SET decrypted_values = decrypt(encrypted_values, 'key_1', 'random_message');
 
 query I
 SELECT decrypted_values FROM test_1 LIMIT 10;

test/sql/encryption/string_encryption.test

@@ -0,0 +1,43 @@
+# name: test/sql/secrets/secrets_encryption.test
+# description: Test secret creation for internal encryption
+# group: [simple-encryption/secrets]
+
+require simple_encryption
+
+# Ensure any currently stored secrets don't interfere with the test
+statement ok
+set allow_persistent_secrets=false;
+
+# Create an internal secret (for internal encryption of columns)
+statement ok
+CREATE SECRET key_1 (
+	TYPE ENCRYPTION,
+    MASTER_KEY '0123456789112345',
+    LENGTH 16
+);
+
+#statement ok
+#SELECT encrypt('testtest', 'key_1', 'random_message');
+
+statement ok
+CREATE TABLE rd_data AS
+    SELECT
+        SUBSTRING(MD5(RANDOM()::TEXT), 1, 5) AS rd_values
+    FROM
+        range(10);
+
+statement ok
+ALTER TABLE rd_data
+  ADD COLUMN encrypted_value VARCHAR;
+
+statement ok
+ALTER TABLE rd_data
+  ADD COLUMN decrypted_value VARCHAR;
+
+statement ok
+UPDATE rd_data
+SET encrypted_value = encrypt(rd_values, 'key_1', 'random_message');
+
+statement ok
+UPDATE rd_data
+SET decrypted_value = decrypt(encrypted_value, 'key_1', 'random_message');

test/sql/secrets/secrets_encryption.test

@@ -2,8 +2,8 @@
 # description: Test secret creation for internal encryption
 # group: [simple-encryption/secrets]
 
-#statement ok
-#PRAGMA enable_verification;
+statement ok
+PRAGMA enable_verification;
 
 require simple_encryption