gRPC reencryption code

src/enclave/App/App.cpp

@@ -239,6 +239,88 @@ JNIEXPORT void JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_Fin
   env->ReleaseByteArrayElements(shared_key_msg_input, shared_key_msg_bytes, 0);
 }
 
+/////////////////////////////// Shared Key Gen Begin ////////////////////////////////
+
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_GetPublicKey(
+  JNIEnv *env, jobject obj, jlong eid) {
+  (void)obj;
+  (void)eid;
+
+  uint8_t* report_msg = NULL;
+  size_t report_msg_size = 0;
+
+  oe_check_and_time("Get enclave public key",
+                     ecall_get_public_key((oe_enclave_t*)eid,
+                                           &report_msg,
+                                           &report_msg_size));
+
+  // Allocate memory
+  jbyteArray report_msg_bytes = env->NewByteArray(report_msg_size);
+  env->SetByteArrayRegion(report_msg_bytes, 0, report_msg_size, reinterpret_cast<jbyte *>(report_msg));
+
+  return report_msg_bytes;
+}
+
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_GetListEncrypted(
+  JNIEnv *env, jobject obj, jlong eid, jbyteArray shared_key_msg_input) {
+  (void)obj;
+
+  jboolean if_copy = false;
+  jbyte *shared_key_msg_bytes = env->GetByteArrayElements(shared_key_msg_input, &if_copy);
+  uint32_t shared_key_msg_size = static_cast<uint32_t>(env->GetArrayLength(shared_key_msg_input));
+
+  size_t report_msg_size = OE_SHARED_KEY_CIPHERTEXT_SIZE * (shared_key_msg_size / OE_PUBLIC_KEY_SIZE);
+  uint8_t* report_msg = new uint8_t[report_msg_size];
+
+  oe_check_and_time("Get List Encrypted",
+                    ecall_get_list_encrypted((oe_enclave_t*)eid,
+                                             reinterpret_cast<uint8_t *>(shared_key_msg_bytes),
+                                             shared_key_msg_size,
+                                             report_msg,
+                                             report_msg_size));
+
+  // Allocate memory
+  jbyteArray report_msg_bytes = env->NewByteArray(report_msg_size);
+  env->SetByteArrayRegion(report_msg_bytes, 0, report_msg_size, reinterpret_cast<jbyte *>(report_msg));
+
+  env->ReleaseByteArrayElements(shared_key_msg_input, (jbyte *) shared_key_msg_bytes, 0);
+
+  delete[] report_msg;
+
+  return report_msg_bytes;
+}
+
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_FinishSharedKey(
+  JNIEnv *env, jobject obj, jlong eid, jbyteArray shared_key_msg_input) {
+  (void)obj;
+
+  jboolean if_copy = false;
+
+  jbyte *shared_key_msg_bytes = env->GetByteArrayElements(shared_key_msg_input, &if_copy);
+  uint32_t shared_key_msg_size = static_cast<uint32_t>(env->GetArrayLength(shared_key_msg_input));
+
+  size_t report_msg_size = SGX_AESGCM_KEY_SIZE;
+  uint8_t* report_msg = new uint8_t[report_msg_size];
+
+  oe_check_and_time("Finish attestation",
+                    ecall_finish_shared_key((oe_enclave_t*)eid,
+                                             reinterpret_cast<uint8_t *>(shared_key_msg_bytes),
+                                             shared_key_msg_size,
+                                             report_msg,
+                                             report_msg_size));
+
+  // Allocate memory
+  jbyteArray report_msg_bytes = env->NewByteArray(report_msg_size);
+  env->SetByteArrayRegion(report_msg_bytes, 0, report_msg_size, reinterpret_cast<jbyte *>(report_msg));
+
+  env->ReleaseByteArrayElements(shared_key_msg_input, shared_key_msg_bytes, 0);
+  delete[] report_msg;
+
+  return report_msg_bytes;
+}
+
+/////////////////////////////// Shared Key Gen End ////////////////////////////////
+
 JNIEXPORT void JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_StopEnclave(
     JNIEnv *env, jobject obj, jlong eid) {
   (void)env;
@@ -345,6 +427,72 @@ JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEncla
   return ciphertext;
 }
 
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_ReEncryptUser(
+    JNIEnv *env, jobject obj, jlong eid, jbyteArray ciphertext, jstring username) {
+  (void)obj;
+
+  uint32_t clength = (uint32_t)env->GetArrayLength(ciphertext);
+
+  jboolean if_copy = false;
+  uint8_t *ciphertext_ptr = (uint8_t *)env->GetByteArrayElements(ciphertext, &if_copy);
+
+  uint8_t *new_ciphertext_copy = nullptr;
+  jsize new_clength = clength;
+
+  const char *username_str = env->GetStringUTFChars(username, nullptr);
+
+  if (ciphertext_ptr == nullptr) {
+    ocall_throw("Encrypt: JNI failed to get input byte array.");
+  } else {
+    new_ciphertext_copy = new uint8_t[new_clength];
+
+    oe_check("Encrypt", ecall_re_encrypt_user((oe_enclave_t *)eid, ciphertext_ptr, clength,
+                                      new_ciphertext_copy, (uint32_t)new_clength, username_str));
+  }
+
+  jbyteArray new_ciphertext = env->NewByteArray(new_clength);
+  env->SetByteArrayRegion(new_ciphertext, 0, new_clength, (jbyte *)new_ciphertext_copy);
+
+  env->ReleaseByteArrayElements(ciphertext, (jbyte *)ciphertext_ptr, 0);
+
+  delete[] new_ciphertext_copy;
+
+  return new_ciphertext;
+}
+
+
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_Decrypt(
+    JNIEnv *env, jobject obj, jlong eid, jbyteArray ciphertext) {
+  (void)obj;
+
+  uint32_t clength = (uint32_t)env->GetArrayLength(ciphertext);
+  jboolean if_copy = false;
+  uint8_t *ciphertext_ptr = (uint8_t *)env->GetByteArrayElements(ciphertext, &if_copy);
+
+  uint8_t *plaintext_copy = nullptr;
+  jsize plength = 0;
+
+  if (ciphertext_ptr == nullptr) {
+    ocall_throw("Encrypt: JNI failed to get input byte array.");
+  } else {
+    plength = clength - SGX_AESGCM_IV_SIZE - SGX_AESGCM_MAC_SIZE;
+    plaintext_copy = new uint8_t[plength];
+
+    oe_check("Decrypt", ecall_decrypt((oe_enclave_t *)eid, ciphertext_ptr, clength,
+                                      plaintext_copy, (uint32_t)plength));
+  }
+
+  jbyteArray plaintext = env->NewByteArray(plength);
+  env->SetByteArrayRegion(plaintext, 0, plength, (jbyte *)plaintext_copy);
+
+  env->ReleaseByteArrayElements(ciphertext, (jbyte *)ciphertext_ptr, 0);
+
+  delete[] plaintext_copy;
+
+  return plaintext;
+}
+
+
 JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_Sample(
     JNIEnv *env, jobject obj, jlong eid, jbyteArray input_rows) {
   (void)obj;

src/enclave/App/SGXEnclave.h

@@ -21,6 +21,9 @@ JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEncla
 JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_Encrypt(
     JNIEnv *, jobject, jlong, jbyteArray);
 
+JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_ReEncryptUser(
+    JNIEnv *, jobject, jlong, jbyteArray, jstring);
+
 JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_Decrypt(
     JNIEnv *, jobject, jlong, jbyteArray);
 
@@ -80,6 +83,15 @@ Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_GenerateReport(JNIEnv *, j
 JNIEXPORT void JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_FinishAttestation(
     JNIEnv *, jobject, jlong, jbyteArray);
 
+  JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_GetPublicKey(
+    JNIEnv *, jobject, jlong);
+
+  JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_GetListEncrypted(
+    JNIEnv *, jobject, jlong, jbyteArray);
+
+  JNIEXPORT jbyteArray JNICALL Java_edu_berkeley_cs_rise_opaque_execution_SGXEnclave_FinishSharedKey(
+    JNIEnv *, jobject, jlong, jbyteArray);
+
 #ifdef __cplusplus
 }
 #endif

src/enclave/Enclave/Crypto.cpp

@@ -1,9 +1,20 @@
+#include "Crypto.h"
+#include "Random.h"
+
 #include <stdexcept>
+#include <unordered_map>
+#include <vector>
+#include <iostream>
 
 #include "Crypto.h"
 #include "Random.h"
 #include "common.h"
 #include "util.h"
+#include <unordered_map>
+#include <vector>
+
+// Set this number before creating the enclave
+uint8_t num_clients = 1;
 
 /**
  * Symmetric key used to encrypt row data. This key is shared among the driver
@@ -17,16 +28,48 @@ unsigned char shared_key[SGX_AESGCM_KEY_SIZE] = {0};
 
 std::unique_ptr<KeySchedule> ks;
 
+// may user name to public key
+std::unordered_map<std::string, std::unique_ptr<KeySchedule>> client_key_schedules;
+std::unordered_map<std::string, std::vector<uint8_t>> client_keys;
+
 void initKeySchedule() {
+  // Use shared key to init key schedule
   ks.reset(new KeySchedule(reinterpret_cast<unsigned char *>(shared_key), SGX_AESGCM_KEY_SIZE));
 }
 
+void initKeySchedule(char* username) {
+  std::string user(username);
+  std::unique_ptr<KeySchedule> user_ks; 
+  unsigned char client_key[SGX_AESGCM_KEY_SIZE];
+
+  auto iter = client_keys.find(user);
+  // if (iter == client_keys.end()) {
+  //   ocall_throw("No client key for user: %s", username);
+  // } else {
+  memcpy(client_key, (uint8_t*) iter->second.data(), SGX_AESGCM_KEY_SIZE);
+  // }
+
+  user_ks.reset(new KeySchedule(reinterpret_cast<unsigned char *>(client_key), SGX_AESGCM_KEY_SIZE));
+  client_key_schedules[user] = std::move(user_ks);
+}
+
+void add_client_key(uint8_t *client_key_bytes, uint32_t client_key_size, char* username) {
+  if (client_key_size <= 0) {
+    throw std::runtime_error("Add client key failed: Invalid client key size");
+  }
+
+  std::vector<uint8_t> user_private_key(client_key_bytes, client_key_bytes + client_key_size);
+  std::string user(username);
+  client_keys[user] = user_private_key;
+
+  initKeySchedule(username);
+}
+
 void set_shared_key(uint8_t *shared_key_bytes, uint32_t shared_key_size) {
   if (shared_key_size <= 0) {
     throw std::runtime_error("Attempting to set a shared key with invalid key size.");
   }
   memcpy_s(shared_key, sizeof(shared_key), shared_key_bytes, shared_key_size);
-
   initKeySchedule();
 }
 
@@ -46,13 +89,36 @@ void encrypt(uint8_t *plaintext, uint32_t plaintext_length, uint8_t *ciphertext)
   AesGcm cipher(ks.get(), reinterpret_cast<uint8_t *>(iv_ptr), SGX_AESGCM_IV_SIZE);
   cipher.encrypt(plaintext, plaintext_length, ciphertext_ptr, plaintext_length);
   memcpy(mac_ptr, cipher.tag().t, SGX_AESGCM_MAC_SIZE);
+
+}
+
+void encrypt_user(uint8_t *plaintext, uint32_t plaintext_length, uint8_t *ciphertext, const char * user_name) {
+
+  for (auto& keypair : client_key_schedules) {
+
+    if (strcmp(keypair.first.c_str(), user_name) == 0) {
+      uint8_t *iv_ptr = ciphertext;
+      uint8_t *ciphertext_ptr = ciphertext + SGX_AESGCM_IV_SIZE;
+      sgx_aes_gcm_128bit_tag_t *mac_ptr =
+          (sgx_aes_gcm_128bit_tag_t *)(ciphertext + SGX_AESGCM_IV_SIZE + plaintext_length);
+      mbedtls_read_rand(reinterpret_cast<unsigned char *>(iv_ptr), SGX_AESGCM_IV_SIZE);
+
+      AesGcm cipher(keypair.second.get(), reinterpret_cast<uint8_t *>(iv_ptr), SGX_AESGCM_IV_SIZE);
+      cipher.encrypt(plaintext, plaintext_length, ciphertext_ptr, plaintext_length);
+      memcpy(mac_ptr, cipher.tag().t, SGX_AESGCM_MAC_SIZE);
+      return;
+    }
+  }
+  throw std::runtime_error("Couldn't find user and key");
 }
 
 void decrypt(const uint8_t *ciphertext, uint32_t ciphertext_length, uint8_t *plaintext) {
+
   if (!ks) {
     throw std::runtime_error("Cannot encrypt without a shared key. Ensure all "
                              "enclaves have completed attestation.");
   }
+
   uint32_t plaintext_length = dec_size(ciphertext_length);
 
   uint8_t *iv_ptr = (uint8_t *)ciphertext;
@@ -63,7 +129,27 @@ void decrypt(const uint8_t *ciphertext, uint32_t ciphertext_length, uint8_t *pla
   AesGcm decipher(ks.get(), iv_ptr, SGX_AESGCM_IV_SIZE);
   decipher.decrypt(ciphertext_ptr, plaintext_length, plaintext, plaintext_length);
   if (memcmp(mac_ptr, decipher.tag().t, SGX_AESGCM_MAC_SIZE) != 0) {
-    printf("Decrypt: invalid mac\n");
+    // Shared key doesn't work
+    // Perhaps we need to use a client key instead
+    int success = -1;
+    for (auto& keypair : client_key_schedules) {
+      std::vector<uint8_t> print_key = client_keys[keypair.first];
+      for(size_t i=0; i < print_key.size(); i++)
+        std::cout << print_key.at(i) << ' ';
+      std::cout << std::endl;
+
+      AesGcm decipher(keypair.second.get(), iv_ptr, SGX_AESGCM_IV_SIZE);
+      decipher.decrypt(ciphertext_ptr, plaintext_length, plaintext, plaintext_length);
+      if (memcmp(mac_ptr, decipher.tag().t, SGX_AESGCM_MAC_SIZE) == 0) {
+          std::cout << "We found the proper key, of user " << keypair.first << std::endl;
+          success = 0;
+          break;
+      }
+    }
+
+    if (success == -1) {
+        throw std::runtime_error("Couldn't decrypt -- proper key unknown\n");
+    }
   }
 }
 

src/enclave/Enclave/Crypto.h

@@ -19,8 +19,12 @@ extern const sgx_ec256_public_t g_sp_pub_key;
  * Set the symmetric key used to encrypt row data using message 4 of the remote
  * attestation process.
  */
+void add_client_key(uint8_t *client_key_bytes, uint32_t client_key_size, char* username);
+
 void set_shared_key(uint8_t *msg4, uint32_t msg4_size);
 
+void add_client_key(uint8_t *client_key_bytes, uint32_t client_key_size, char* username);
+
 /**
  * Encrypt the given plaintext using AES-GCM with a 128-bit key and write the
  * result to `ciphertext`. The encrypted data will be formatted as follows,
@@ -35,6 +39,8 @@ void set_shared_key(uint8_t *msg4, uint32_t msg4_size);
  */
 void encrypt(uint8_t *plaintext, uint32_t plaintext_length, uint8_t *ciphertext);
 
+void encrypt_user(uint8_t *plaintext, uint32_t plaintext_length, uint8_t *ciphertext, const char * user_name);
+
 /**
  * Decrypt the given ciphertext using AES-GCM with a 128-bit key and write the
  * result to `plaintext`. The encrypted data must be formatted as described in