Support for scalar subquery (#157)

This PR implements the scalar subquery expression, which is triggered whenever a subquery returns a scalar value. There were two main problems that needed to be solved.

First, support for matching the scalar subquery expression is necessary. Spark implements this by wrapping a SparkPlan within the expression and calls executeCollect. Then it constructs a literal with that value. However, this is problematic for us because that value should not be decrypted by the driver and serialized into an expression, since it's an intermediate value.

Therefore, the second issue to be addressed here is supporting an encrypted literal. This is implemented in this PR by serializing an encrypted ciphertext into a base64 encoded string, and wrapping a Decrypt expression on top of it. This expression is then evaluated in the enclave and returns a literal. Note that, in order to test our implementation, we also implement a Decrypt expression in Scala. However, this should never be evaluated on the driver side and serialized into a plaintext literal. This is because Decrypt is designated as a Nondeterministic expression, and therefore will always evaluate on the workers.

Author: wzheng

src/enclave/Enclave/ExpressionEvaluation.h

@@ -288,6 +288,49 @@ class FlatbuffersExpressionEvaluator {
         static_cast<const tuix::Literal *>(expr->expr())->value(), builder);
     }
 
+    case tuix::ExprUnion_Decrypt:
+    {
+      auto decrypt_expr = static_cast<const tuix::Decrypt *>(expr->expr());
+      const tuix::Field *value =
+        flatbuffers::GetTemporaryPointer(builder, eval_helper(row, decrypt_expr->value()));
+
+      if (value->value_type() != tuix::FieldUnion_StringField) {
+        throw std::runtime_error(
+          std::string("tuix::Decrypt only accepts a string input, not ")
+          + std::string(tuix::EnumNameFieldUnion(value->value_type())));
+      }
+
+      bool result_is_null = value->is_null();
+      if (!result_is_null) {
+        auto str_field = static_cast<const tuix::StringField *>(value->value());
+
+        std::vector<uint8_t> str_vec(
+            flatbuffers::VectorIterator<uint8_t, uint8_t>(str_field->value()->Data(),
+                                                          static_cast<uint32_t>(0)),
+            flatbuffers::VectorIterator<uint8_t, uint8_t>(str_field->value()->Data(),
+                                                          static_cast<uint32_t>(str_field->length())));
+
+        std::string ciphertext(str_vec.begin(), str_vec.end());
+        std::string ciphertext_decoded = ciphertext_base64_decode(ciphertext);
+
+        uint8_t *plaintext = new uint8_t[dec_size(ciphertext_decoded.size())];
+        decrypt(reinterpret_cast<const uint8_t *>(ciphertext_decoded.data()), ciphertext_decoded.size(), plaintext);
+
+        BufferRefView<tuix::Rows> buf(plaintext, ciphertext_decoded.size());
+        buf.verify();
+
+        const tuix::Rows *rows = buf.root();
+        const tuix::Field *field = rows->rows()->Get(0)->field_values()->Get(0);
+        auto ret = flatbuffers_copy<tuix::Field>(field, builder);
+         
+        delete plaintext;
+        return ret;
+      } else {
+        throw std::runtime_error(std::string("tuix::Decrypt does not accept a NULL string\n"));
+      }
+      
+    }
+
     case tuix::ExprUnion_Cast:
     {
       auto cast = static_cast<const tuix::Cast *>(expr->expr());

src/enclave/Enclave/util.cpp

@@ -142,3 +142,79 @@ int secs_to_tm(long long t, struct tm *tm) {
 
   return 0;
 }
+
+// Code adapted from https://stackoverflow.com/questions/180947/base64-decode-snippet-in-c
+/*
+   Copyright (C) 2004-2008 Rene Nyffenegger
+
+   This source code is provided 'as-is', without any express or implied
+   warranty. In no event will the author be held liable for any damages
+   arising from the use of this software.
+
+   Permission is granted to anyone to use this software for any purpose,
+   including commercial applications, and to alter it and redistribute it
+   freely, subject to the following restrictions:
+
+   1. The origin of this source code must not be misrepresented; you must not
+      claim that you wrote the original source code. If you use this source code
+      in a product, an acknowledgment in the product documentation would be
+      appreciated but is not required.
+
+   2. Altered source versions must be plainly marked as such, and must not be
+      misrepresented as being the original source code.
+
+   3. This notice may not be removed or altered from any source distribution.
+
+   Rene Nyffenegger [email protected]
+
+*/
+
+static const std::string base64_chars =
+             "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+             "abcdefghijklmnopqrstuvwxyz"
+             "0123456789+/";
+
+static inline bool is_base64(unsigned char c) {
+  return (isalnum(c) || (c == '+') || (c == '/'));
+}
+
+std::string ciphertext_base64_decode(const std::string &encoded_string) {
+  int in_len = encoded_string.size();
+  int i = 0;
+  int j = 0;
+  int in_ = 0;
+  uint8_t char_array_4[4], char_array_3[3];
+  std::string ret;
+
+  while (in_len-- && ( encoded_string[in_] != '=') && is_base64(encoded_string[in_])) {
+    char_array_4[i++] = encoded_string[in_]; in_++;
+    if (i ==4) {
+      for (i = 0; i <4; i++)
+        char_array_4[i] = base64_chars.find(char_array_4[i]);
+
+      char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+      char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+      char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+
+      for (i = 0; (i < 3); i++)
+        ret += char_array_3[i];
+      i = 0;
+    }
+  }
+
+  if (i) {
+    for (j = i; j <4; j++)
+      char_array_4[j] = 0;
+
+    for (j = 0; j <4; j++)
+      char_array_4[j] = base64_chars.find(char_array_4[j]);
+
+    char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+    char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+    char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+
+    for (j = 0; (j < i - 1); j++) ret += char_array_3[j];
+  }
+
+  return ret;
+}

src/enclave/Enclave/util.h

@@ -41,4 +41,6 @@ int pow_2(int value);
 
 int secs_to_tm(long long t, struct tm *tm);
 
+std::string ciphertext_base64_decode(const std::string &encoded_string);
+
 #endif // UTIL_H

src/flatbuffers/Expr.fbs

@@ -40,7 +40,8 @@ union ExprUnion {
     CreateArray,
     Upper,
     DateAdd,
-    DateAddInterval
+    DateAddInterval,
+    Decrypt
 }
 
 table Expr {
@@ -221,4 +222,8 @@ table ClosestPoint {
 
 table Upper {
     child:Expr;
-}
+}
+
+table Decrypt {
+    value:Expr;
+}

src/main/scala/edu/berkeley/cs/rise/opaque/Utils.scala

@@ -21,7 +21,9 @@ import java.io.File
 import java.io.FileNotFoundException
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
+import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom
+import java.util.Base64
 import java.util.UUID
 
 import javax.crypto._
@@ -92,6 +94,8 @@ import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.catalyst.util.ArrayBasedMapData
 import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.catalyst.util.MapData
+import org.apache.spark.sql.execution.SubqueryExec
+import org.apache.spark.sql.execution.ScalarSubquery
 import org.apache.spark.sql.execution.aggregate.ScalaUDAF
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.StorageLevel
@@ -102,6 +106,7 @@ import edu.berkeley.cs.rise.opaque.execution.Block
 import edu.berkeley.cs.rise.opaque.execution.OpaqueOperatorExec
 import edu.berkeley.cs.rise.opaque.execution.SGXEnclave
 import edu.berkeley.cs.rise.opaque.expressions.ClosestPoint
+import edu.berkeley.cs.rise.opaque.expressions.Decrypt
 import edu.berkeley.cs.rise.opaque.expressions.DotProduct
 import edu.berkeley.cs.rise.opaque.expressions.VectorAdd
 import edu.berkeley.cs.rise.opaque.expressions.VectorMultiply
@@ -589,6 +594,7 @@ object Utils extends Logging {
             tuix.StringField.createValueVector(builder, Array.empty),
             0),
           isNull)
+      case _ => throw new OpaqueException(s"FlatbuffersCreateField failed to match on ${value} of type {value.getClass.getName()}, ${dataType}")
     }
   }
 
@@ -663,6 +669,50 @@ object Utils extends Logging {
 
   val MaxBlockSize = 1000
 
+  /**
+    * Encrypts/decrypts a given scalar value
+    **/
+  def encryptScalar(value: Any, dataType: DataType): String = {
+    // First serialize the scalar value
+    var builder = new FlatBufferBuilder
+    var rowOffsets = ArrayBuilder.make[Int]
+
+    val v = dataType match {
+      case StringType => UTF8String.fromString(value.asInstanceOf[String])
+      case _ => value
+    }
+
+    val isNull = (value == null)
+
+    // TODO: the NULL variable for field value could be set to true
+    builder.finish(
+      tuix.Rows.createRows(
+        builder,
+        tuix.Rows.createRowsVector(
+          builder,
+          Array(tuix.Row.createRow(
+            builder,
+            tuix.Row.createFieldValuesVector(
+              builder,
+              Array(flatbuffersCreateField(builder, v, dataType, false))),
+            isNull)))))
+
+    val plaintext = builder.sizedByteArray()
+    val ciphertext = encrypt(plaintext)
+    val ciphertext_str = Base64.getEncoder().encodeToString(ciphertext);
+    ciphertext_str
+  }
+
+  def decryptScalar(ciphertext: String): Any = {
+    val ciphertext_bytes = Base64.getDecoder().decode(ciphertext);
+    val plaintext = decrypt(ciphertext_bytes)
+    val rows = tuix.Rows.getRootAsRows(ByteBuffer.wrap(plaintext))
+    val row = rows.rows(0)
+    val field = row.fieldValues(0)
+    val value = flatbuffersExtractFieldValue(field)
+    value
+  }
+
   /**
    * Encrypts the given Spark SQL [[InternalRow]]s into a [[Block]] (a serialized
    * tuix.EncryptedBlocks).
@@ -822,6 +872,13 @@ object Utils extends Logging {
             tuix.ExprUnion.Literal,
             tuix.Literal.createLiteral(builder, valueOffset))
 
+        // This expression should never be evaluated on the driver
+        case (Decrypt(child, dataType), Seq(childOffset)) =>
+          tuix.Expr.createExpr(
+            builder,
+            tuix.ExprUnion.Decrypt,
+            tuix.Decrypt.createDecrypt(builder, childOffset))
+
         case (Alias(child, _), Seq(childOffset)) =>
           // TODO: Use an expression for aliases so we can refer to them elsewhere in the expression
           // tree. For now we just ignore them when evaluating expressions.
@@ -1112,6 +1169,36 @@ object Utils extends Logging {
           // TODO: Implement decimal serialization, followed by CheckOverflow
           childOffset
 
+        case (ScalarSubquery(SubqueryExec(name, child), exprId), Seq()) =>
+          val output = child.output(0)
+          val dataType = output match {
+            case AttributeReference(name, dataType, _, _) => dataType
+            case _ => throw new OpaqueException("Scalar subquery cannot match to AttributeReference")
+          }
+          // Need to deserialize the encrypted blocks to get the encrypted block
+          val blockList = child.asInstanceOf[OpaqueOperatorExec].collectEncrypted()
+          val encryptedBlocksList = blockList.map { block =>
+            val buf = ByteBuffer.wrap(block.bytes)
+            tuix.EncryptedBlocks.getRootAsEncryptedBlocks(buf)
+          }
+          val encryptedBlocks = encryptedBlocksList.find(_.blocksLength > 0).getOrElse(encryptedBlocksList(0))
+          if (encryptedBlocks.blocksLength == 0) {
+            // If empty, the returned result is null
+            flatbuffersSerializeExpression(builder, Literal(null, dataType), input)
+          } else {
+            assert(encryptedBlocks.blocksLength == 1)
+            val encryptedBlock = encryptedBlocks.blocks(0)
+            val ciphertextBuf = encryptedBlock.encRowsAsByteBuffer
+            val ciphertext = new Array[Byte](ciphertextBuf.remaining)
+            ciphertextBuf.get(ciphertext)
+            val ciphertext_str = Base64.getEncoder().encodeToString(ciphertext)
+            flatbuffersSerializeExpression(
+              builder,
+              Decrypt(Literal(UTF8String.fromString(ciphertext_str), StringType), dataType),
+              input
+            )
+          }
+
         case (_, Seq(childOffset)) =>
           throw new OpaqueException("Expression not supported: " + expr.toString())
       }

src/main/scala/edu/berkeley/cs/rise/opaque/execution/operators.scala

@@ -134,14 +134,19 @@ trait OpaqueOperatorExec extends SparkPlan {
    * method and persist the resulting RDD. [[ConvertToOpaqueOperators]] later eliminates the dummy
    * relation from the logical plan, but this only happens after InMemoryRelation has called this
    * method. We therefore have to silently return an empty RDD here.
-   */
+    */
+
   override def doExecute(): RDD[InternalRow] = {
     sqlContext.sparkContext.emptyRDD
     // throw new UnsupportedOperationException("use executeBlocked")
   }
 
+  def collectEncrypted(): Array[Block] = {
+    executeBlocked().collect
+  }
+
   override def executeCollect(): Array[InternalRow] = {
-    executeBlocked().collect().flatMap { block =>
+    collectEncrypted().flatMap { block =>
       Utils.decryptBlockFlatbuffers(block)
     }
   }

src/main/scala/edu/berkeley/cs/rise/opaque/expressions/ClosestPoint.scala

@@ -29,9 +29,6 @@ object ClosestPoint {
       * point - list of coordinates representing a point
       * centroids - list of lists of coordinates, each representing a point
   """)
-/**
- *
- */
 case class ClosestPoint(left: Expression, right: Expression)
     extends BinaryExpression with NullIntolerant with CodegenFallback with ExpectsInputTypes {
 

src/main/scala/edu/berkeley/cs/rise/opaque/expressions/Decrypt.scala

@@ -0,0 +1,49 @@
+package edu.berkeley.cs.rise.opaque.expressions
+
+import edu.berkeley.cs.rise.opaque.Utils
+
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.ExpressionDescription
+import org.apache.spark.sql.catalyst.expressions.NullIntolerant
+import org.apache.spark.sql.catalyst.expressions.Nondeterministic
+import org.apache.spark.sql.catalyst.expressions.UnaryExpression
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.types.DataTypes
+import org.apache.spark.sql.types.StringType
+import org.apache.spark.unsafe.types.UTF8String
+
+object Decrypt {
+  def decrypt(v: Column, dataType: DataType): Column = new Column(Decrypt(v.expr, dataType))
+}
+
+@ExpressionDescription(
+  usage = """
+    _FUNC_(child, outputDataType) - Decrypt the input evaluated expression, which should always be a string
+  """,
+  arguments = """
+    Arguments:
+      * child - an encrypted literal of string type
+      * outputDataType - the decrypted data type
+  """)
+case class Decrypt(child: Expression, outputDataType: DataType)
+    extends UnaryExpression with NullIntolerant with CodegenFallback with Nondeterministic {
+
+  override def dataType: DataType = outputDataType
+
+  protected def initializeInternal(partitionIndex: Int): Unit = { }
+
+  protected override def evalInternal(input: InternalRow): Any = {
+    val v = child.eval()
+    nullSafeEval(v)
+  }
+
+  protected override def nullSafeEval(input: Any): Any = {
+    // This function is implemented so that we can test against Spark;
+    // should never be used in production because we want to keep the literal encrypted
+    val v = input.asInstanceOf[UTF8String].toString
+    Utils.decryptScalar(v)
+  }
+}