From 2fb9b30b7778946720b2fc77f04c32970a49a6c6 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Mon, 28 Dec 2015 19:11:14 -0800
Subject: [PATCH 01/10] [SPARK-12567] add aes {encrypt, decrypt}

---
 .../catalyst/analysis/FunctionRegistry.scala  |   2 +
 .../spark/sql/catalyst/expressions/misc.scala | 129 +++++++++++++++++-
 .../expressions/MiscFunctionsSuite.scala      |  45 ++++++
 .../org/apache/spark/sql/functions.scala      |  34 +++++
 4 files changed, 209 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index d9009e3848e58..6b87dfba5d04f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -278,6 +278,8 @@ object FunctionRegistry {
     expression[ArrayContains]("array_contains"),
 
     // misc functions
+    expression[AesEncrypt]("aes_encrypt"),
+    expression[AesDecrypt]("aes_decrypt"),
     expression[Crc32]("crc32"),
     expression[Md5]("md5"),
     expression[Murmur3Hash]("hash"),
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
index f4ccadd9c563e..2471231042241 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
@@ -17,9 +17,12 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import java.security.{MessageDigest, NoSuchAlgorithmException}
+import java.security.{GeneralSecurityException, MessageDigest, NoSuchAlgorithmException}
 import java.util.zip.CRC32
 
+import javax.crypto.Cipher
+import javax.crypto.spec.SecretKeySpec
+
 import org.apache.commons.codec.digest.DigestUtils
 
 import org.apache.spark.sql.catalyst.InternalRow
@@ -442,3 +445,127 @@ case class PrintToStderr(child: Expression) extends UnaryExpression {
        """.stripMargin)
   }
 }
+
+/**
+ * A function that encrypts input using AES. Key lengths of 128, 192 or 256 bits can be used.
+ * 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
+ * Jurisdiction Policy Files are installed. If either argument is NULL or the key length is
+ * not one of the permitted values, the return value is NULL.
+ */
+@ExpressionDescription(
+  usage =
+    """_FUNC_(input, key) - Encrypts input using AES. Key lengths of 128, 192 or 256 bits can
+ be used. 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
+ Jurisdiction Policy Files are installed. If either argument is NULL or the key length is not one
+ of the permitted values, the return value is NULL.""",
+  extended = """> SELECT Base64(_FUNC_('ABC', '1234567890123456'));
+ 'y6Ss+zCYObpCbgfWfyNWTw=='"""
+)
+case class AesEncrypt(left: Expression, right: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def dataType: DataType = BinaryType
+  override def nullable: Boolean = true
+
+  override def inputTypes: Seq[DataType] = Seq(BinaryType, BinaryType)
+
+  protected override def nullSafeEval(input1: Any, input2: Any): Any = {
+    val cipher = Cipher.getInstance("AES")
+    val secretKey: SecretKeySpec = input2.asInstanceOf[Array[Byte]].length match {
+      case 16 | 24 | 32 =>
+        new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
+          input2.asInstanceOf[Array[Byte]].length, "AES")
+      case _ => null
+    }
+    try {
+      cipher.init(Cipher.ENCRYPT_MODE, secretKey)
+      cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0, input1.asInstanceOf[Array[Byte]].length)
+    } catch {
+      case e: GeneralSecurityException => null
+    }
+  }
+
+  override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
+    nullSafeCodeGen(ctx, ev, (str, key) => {
+      val Cipher = "javax.crypto.Cipher"
+      val SecretKeySpec = "javax.crypto.spec.SecretKeySpec"
+      s"""
+          try {
+            $Cipher cipher = $Cipher.getInstance("AES");
+            if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
+              cipher.init($Cipher.ENCRYPT_MODE, new $SecretKeySpec($key, 0, $key.length, "AES"));
+              ${ev.value} = cipher.doFinal($str, 0, $str.length);
+            } else {
+              ${ev.isNull} = true;
+            }
+          } catch (java.security.GeneralSecurityException e) {
+            org.apache.spark.unsafe.Platform.throwException(e);
+          }
+      """
+    })
+  }
+}
+
+/**
+ * A function that decrypts input using AES. Key lengths of 128, 192 or 256 bits can be used.
+ * 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
+ * Jurisdiction Policy Files are installed. If either argument is NULL or the key length is
+ * not one of the permitted values, the return value is NULL.
+ */
+@ExpressionDescription(
+  usage = """_FUNC_(input, key) - Decrypts input using AES. Key lengths of 128, 192 or 256 bits can
+ be used. 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
+ Jurisdiction Policy Files are installed. If either argument is NULL or the key length is not one
+ of the permitted values, the return value is NULL.""",
+  extended =
+    """> SELECT _FUNC_(UnBase64('y6Ss+zCYObpCbgfWfyNWTw=='),'1234567890123456');
+ 'ABC'"""
+)
+case class AesDecrypt(left: Expression, right: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def dataType: DataType = StringType
+  override def nullable: Boolean = true
+
+  override def inputTypes: Seq[DataType] = Seq(BinaryType, BinaryType)
+
+  protected override def nullSafeEval(input1: Any, input2: Any): Any = {
+    val cipher = Cipher.getInstance("AES")
+    val secretKey = input2.asInstanceOf[Array[Byte]].length match {
+      case 16 | 24 | 32 =>
+        new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
+          input2.asInstanceOf[Array[Byte]].length, "AES")
+      case _ => null
+    }
+
+    try {
+      cipher.init(Cipher.DECRYPT_MODE, secretKey)
+      UTF8String.fromBytes(
+        cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0,
+          input1.asInstanceOf[Array[Byte]].length))
+    } catch {
+      case e: GeneralSecurityException => null
+    }
+
+  }
+
+  override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
+    nullSafeCodeGen(ctx, ev, (str, key) => {
+      val Cipher = "javax.crypto.Cipher"
+      val SecretKeySpec = "javax.crypto.spec.SecretKeySpec"
+      s"""
+          try {
+            $Cipher cipher = $Cipher.getInstance("AES");
+            if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
+              cipher.init($Cipher.DECRYPT_MODE, new $SecretKeySpec($key, 0, $key.length, "AES"));
+              ${ev.value} = UTF8String.fromBytes(cipher.doFinal($str, 0, $str.length));
+            } else {
+              ${ev.isNull} = true;
+            }
+          } catch (java.security.GeneralSecurityException e) {
+            org.apache.spark.unsafe.Platform.throwException(e);
+          }
+      """
+    })
+  }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 75131a6170222..43fe92965ed00 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -132,4 +132,49 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       }
     }
   }
+
+  test("aesEncrypt") {
+    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
+      Literal("1234567890123456".getBytes))), "y6Ss+zCYObpCbgfWfyNWTw==")
+    checkEvaluation(Base64(AesEncrypt(Literal("".getBytes),
+      Literal("1234567890123456".getBytes))), "BQGHoM3lqYcsurCRq3PlUw==")
+
+    // Before testing this, installing Java Cryptography Extension (JCE)
+    // Unlimited Strength Jurisdiction Policy Files first
+//    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
+//      Literal("12345678901234561234567890123456".getBytes))), "nYfCuJeRd5eD60yXDw7WEA==")
+
+    // input is null
+    checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
+      Literal("1234567890123456".getBytes)), null)
+    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
+    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
+      Literal("1234567890".getBytes))), null)
+    // key is null
+    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
+      Literal.create(null, BinaryType))), null)
+    // both are null
+    checkEvaluation(Base64(AesEncrypt(Literal.create(null, BinaryType),
+      Literal.create(null, BinaryType))), null)
+  }
+
+  test("aesDecrypt") {
+    checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal("1234567890123456".getBytes)), "ABC")
+    checkEvaluation(AesDecrypt(UnBase64(Literal("BQGHoM3lqYcsurCRq3PlUw==")),
+      Literal("1234567890123456".getBytes)), "")
+
+    // input is null
+    checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
+      Literal("1234567890123456".getBytes)), null)
+    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
+    checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal("1234567890".getBytes)), null)
+    // key is null
+    checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal.create(null, BinaryType)), null)
+    // both are null
+    checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
+      Literal.create(null, BinaryType)), null)
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index b970eee4e31a7..c5c5234dad0e4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -1931,6 +1931,40 @@ object functions extends LegacyFunctions {
     new Murmur3Hash(cols.map(_.expr))
   }
 
+  /**
+   * Encrypts input using AES and Returns the result as a binary column.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if
+   * Java Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed.
+   * If either argument is NULL or the key length is not one of the permitted values,
+   * the result will also be null.
+   *
+   * @param input binary column to encrypt input
+   * @param key binary column of 128, 192 or 256 bits key
+   *
+   * @group misc_funcs
+   * @since 2.0.0
+   */
+  def aes_encrypt(input: Column, key: Column): Column = withExpr {
+    AesEncrypt(input.expr, key.expr)
+  }
+
+  /**
+   * Decrypts input using AES and Returns the result as a string column.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if
+   * Java Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed.
+   * If either argument is NULL or the key length is not one of the permitted values,
+   * the result will also be null.
+   *
+   * @param input binary column to decrypt input
+   * @param key binary column of 128, 192 or 256 bits key
+   *
+   * @group misc_funcs
+   * @since 2.0.0
+   */
+  def aes_decrypt(input: Column, key: Column): Column = withExpr {
+    AesEncrypt(input.expr, key.expr)
+  }
+
   //////////////////////////////////////////////////////////////////////////////////////////////
   // String functions
   //////////////////////////////////////////////////////////////////////////////////////////////

From c3e8331fb6529d1ddaecddaacc6aca5987e4093e Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Thu, 7 Jan 2016 06:12:20 -0800
Subject: [PATCH 02/10] codegen version has been consistent with  non-codegen
 version

---
 .../spark/sql/catalyst/expressions/misc.scala | 43 ++++++++-----------
 .../expressions/MiscFunctionsSuite.scala      | 14 +++---
 2 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
index 2471231042241..6e712702c7f54 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
@@ -19,7 +19,6 @@ package org.apache.spark.sql.catalyst.expressions
 
 import java.security.{GeneralSecurityException, MessageDigest, NoSuchAlgorithmException}
 import java.util.zip.CRC32
-
 import javax.crypto.Cipher
 import javax.crypto.spec.SecretKeySpec
 
@@ -453,14 +452,8 @@ case class PrintToStderr(child: Expression) extends UnaryExpression {
  * not one of the permitted values, the return value is NULL.
  */
 @ExpressionDescription(
-  usage =
-    """_FUNC_(input, key) - Encrypts input using AES. Key lengths of 128, 192 or 256 bits can
- be used. 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
- Jurisdiction Policy Files are installed. If either argument is NULL or the key length is not one
- of the permitted values, the return value is NULL.""",
-  extended = """> SELECT Base64(_FUNC_('ABC', '1234567890123456'));
- 'y6Ss+zCYObpCbgfWfyNWTw=='"""
-)
+  usage = "_FUNC_(input, key) - Encrypts input using AES.",
+  extended = "> SELECT Base64(_FUNC_('ABC', '1234567890123456'));\n 'y6Ss+zCYObpCbgfWfyNWTw=='")
 case class AesEncrypt(left: Expression, right: Expression)
   extends BinaryExpression with ImplicitCastInputTypes {
 
@@ -477,6 +470,7 @@ case class AesEncrypt(left: Expression, right: Expression)
           input2.asInstanceOf[Array[Byte]].length, "AES")
       case _ => null
     }
+
     try {
       cipher.init(Cipher.ENCRYPT_MODE, secretKey)
       cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0, input1.asInstanceOf[Array[Byte]].length)
@@ -492,14 +486,16 @@ case class AesEncrypt(left: Expression, right: Expression)
       s"""
           try {
             $Cipher cipher = $Cipher.getInstance("AES");
+            $SecretKeySpec secret;
             if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
-              cipher.init($Cipher.ENCRYPT_MODE, new $SecretKeySpec($key, 0, $key.length, "AES"));
-              ${ev.value} = cipher.doFinal($str, 0, $str.length);
+              secret = new $SecretKeySpec($key, 0, $key.length, "AES");
             } else {
-              ${ev.isNull} = true;
+              secret = null;
             }
+            cipher.init($Cipher.ENCRYPT_MODE, secret);
+            ${ev.value} = cipher.doFinal($str, 0, $str.length);
           } catch (java.security.GeneralSecurityException e) {
-            org.apache.spark.unsafe.Platform.throwException(e);
+            ${ev.isNull} = true;
           }
       """
     })
@@ -513,14 +509,8 @@ case class AesEncrypt(left: Expression, right: Expression)
  * not one of the permitted values, the return value is NULL.
  */
 @ExpressionDescription(
-  usage = """_FUNC_(input, key) - Decrypts input using AES. Key lengths of 128, 192 or 256 bits can
- be used. 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
- Jurisdiction Policy Files are installed. If either argument is NULL or the key length is not one
- of the permitted values, the return value is NULL.""",
-  extended =
-    """> SELECT _FUNC_(UnBase64('y6Ss+zCYObpCbgfWfyNWTw=='),'1234567890123456');
- 'ABC'"""
-)
+  usage = "_FUNC_(input, key) - Decrypts input using AES.",
+  extended = "> SELECT _FUNC_(UnBase64('y6Ss+zCYObpCbgfWfyNWTw=='),'1234567890123456');\n 'ABC'")
 case class AesDecrypt(left: Expression, right: Expression)
   extends BinaryExpression with ImplicitCastInputTypes {
 
@@ -546,7 +536,6 @@ case class AesDecrypt(left: Expression, right: Expression)
     } catch {
       case e: GeneralSecurityException => null
     }
-
   }
 
   override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
@@ -556,14 +545,16 @@ case class AesDecrypt(left: Expression, right: Expression)
       s"""
           try {
             $Cipher cipher = $Cipher.getInstance("AES");
+            $SecretKeySpec secret;
             if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
-              cipher.init($Cipher.DECRYPT_MODE, new $SecretKeySpec($key, 0, $key.length, "AES"));
-              ${ev.value} = UTF8String.fromBytes(cipher.doFinal($str, 0, $str.length));
+              secret = new $SecretKeySpec($key, 0, $key.length, "AES");
             } else {
-              ${ev.isNull} = true;
+              secret = null;
             }
+            cipher.init($Cipher.DECRYPT_MODE, secret);
+            ${ev.value} = UTF8String.fromBytes(cipher.doFinal($str, 0, $str.length));
           } catch (java.security.GeneralSecurityException e) {
-            org.apache.spark.unsafe.Platform.throwException(e);
+            ${ev.isNull} = true;
           }
       """
     })
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 43fe92965ed00..273ca5acddeb4 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -139,11 +139,6 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(Base64(AesEncrypt(Literal("".getBytes),
       Literal("1234567890123456".getBytes))), "BQGHoM3lqYcsurCRq3PlUw==")
 
-    // Before testing this, installing Java Cryptography Extension (JCE)
-    // Unlimited Strength Jurisdiction Policy Files first
-//    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
-//      Literal("12345678901234561234567890123456".getBytes))), "nYfCuJeRd5eD60yXDw7WEA==")
-
     // input is null
     checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
       Literal("1234567890123456".getBytes)), null)
@@ -177,4 +172,13 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
       Literal.create(null, BinaryType)), null)
   }
+
+  ignore("aesEncryptWith256bitsKey") {
+    // Before testing this, installing Java Cryptography Extension (JCE)
+    // Unlimited Strength Jurisdiction Policy Files first. Otherwise it
+    // will return `null`. Because Oracle JDK does not support 192 and 256
+    // bits key out of box.
+    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
+      Literal("12345678901234561234567890123456".getBytes))), "nYfCuJeRd5eD60yXDw7WEA==")
+  }
 }

From 1de0453ab64a94af97b1c65528762fcece198624 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Tue, 9 Feb 2016 03:57:14 -0800
Subject: [PATCH 03/10] re-throw the exception

---
 .../spark/sql/catalyst/expressions/misc.scala | 59 +++++--------------
 .../expressions/MiscFunctionsSuite.scala      |  6 --
 2 files changed, 14 insertions(+), 51 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
index 6e712702c7f54..fb176c47de446 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
@@ -458,25 +458,14 @@ case class AesEncrypt(left: Expression, right: Expression)
   extends BinaryExpression with ImplicitCastInputTypes {
 
   override def dataType: DataType = BinaryType
-  override def nullable: Boolean = true
-
   override def inputTypes: Seq[DataType] = Seq(BinaryType, BinaryType)
 
   protected override def nullSafeEval(input1: Any, input2: Any): Any = {
     val cipher = Cipher.getInstance("AES")
-    val secretKey: SecretKeySpec = input2.asInstanceOf[Array[Byte]].length match {
-      case 16 | 24 | 32 =>
-        new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
-          input2.asInstanceOf[Array[Byte]].length, "AES")
-      case _ => null
-    }
-
-    try {
-      cipher.init(Cipher.ENCRYPT_MODE, secretKey)
-      cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0, input1.asInstanceOf[Array[Byte]].length)
-    } catch {
-      case e: GeneralSecurityException => null
-    }
+    val secretKey: SecretKeySpec = new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
+      input2.asInstanceOf[Array[Byte]].length, "AES")
+    cipher.init(Cipher.ENCRYPT_MODE, secretKey)
+    cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0, input1.asInstanceOf[Array[Byte]].length)
   }
 
   override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
@@ -486,16 +475,11 @@ case class AesEncrypt(left: Expression, right: Expression)
       s"""
           try {
             $Cipher cipher = $Cipher.getInstance("AES");
-            $SecretKeySpec secret;
-            if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
-              secret = new $SecretKeySpec($key, 0, $key.length, "AES");
-            } else {
-              secret = null;
-            }
+            $SecretKeySpec secret = new $SecretKeySpec($key, 0, $key.length, "AES");
             cipher.init($Cipher.ENCRYPT_MODE, secret);
             ${ev.value} = cipher.doFinal($str, 0, $str.length);
           } catch (java.security.GeneralSecurityException e) {
-            ${ev.isNull} = true;
+            org.apache.spark.unsafe.Platform.throwException(e);
           }
       """
     })
@@ -515,27 +499,17 @@ case class AesDecrypt(left: Expression, right: Expression)
   extends BinaryExpression with ImplicitCastInputTypes {
 
   override def dataType: DataType = StringType
-  override def nullable: Boolean = true
-
   override def inputTypes: Seq[DataType] = Seq(BinaryType, BinaryType)
 
   protected override def nullSafeEval(input1: Any, input2: Any): Any = {
     val cipher = Cipher.getInstance("AES")
-    val secretKey = input2.asInstanceOf[Array[Byte]].length match {
-      case 16 | 24 | 32 =>
-        new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
-          input2.asInstanceOf[Array[Byte]].length, "AES")
-      case _ => null
-    }
+    val secretKey = new SecretKeySpec(input2.asInstanceOf[Array[Byte]], 0,
+      input2.asInstanceOf[Array[Byte]].length, "AES")
 
-    try {
-      cipher.init(Cipher.DECRYPT_MODE, secretKey)
-      UTF8String.fromBytes(
-        cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0,
-          input1.asInstanceOf[Array[Byte]].length))
-    } catch {
-      case e: GeneralSecurityException => null
-    }
+    cipher.init(Cipher.DECRYPT_MODE, secretKey)
+    UTF8String.fromBytes(
+      cipher.doFinal(input1.asInstanceOf[Array[Byte]], 0,
+        input1.asInstanceOf[Array[Byte]].length))
   }
 
   override def genCode(ctx: CodegenContext, ev: ExprCode): String = {
@@ -545,16 +519,11 @@ case class AesDecrypt(left: Expression, right: Expression)
       s"""
           try {
             $Cipher cipher = $Cipher.getInstance("AES");
-            $SecretKeySpec secret;
-            if ($key.length == 16 || $key.length == 24 || $key.length == 32) {
-              secret = new $SecretKeySpec($key, 0, $key.length, "AES");
-            } else {
-              secret = null;
-            }
+            $SecretKeySpec secret = new $SecretKeySpec($key, 0, $key.length, "AES");
             cipher.init($Cipher.DECRYPT_MODE, secret);
             ${ev.value} = UTF8String.fromBytes(cipher.doFinal($str, 0, $str.length));
           } catch (java.security.GeneralSecurityException e) {
-            ${ev.isNull} = true;
+            org.apache.spark.unsafe.Platform.throwException(e);
           }
       """
     })
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 273ca5acddeb4..53339f4947ecc 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -142,9 +142,6 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     // input is null
     checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
       Literal("1234567890123456".getBytes)), null)
-    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
-    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
-      Literal("1234567890".getBytes))), null)
     // key is null
     checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
       Literal.create(null, BinaryType))), null)
@@ -162,9 +159,6 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     // input is null
     checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
       Literal("1234567890123456".getBytes)), null)
-    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
-    checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
-      Literal("1234567890".getBytes)), null)
     // key is null
     checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
       Literal.create(null, BinaryType)), null)

From 6bc1b631647a8a8052756c4649acc52783737205 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Fri, 12 Feb 2016 22:13:24 -0800
Subject: [PATCH 04/10] address comment

- update the doc
- add error cases
  * key length is not one of the permitted values
  * input can not be decrypted
---
 .../spark/sql/catalyst/expressions/misc.scala | 20 +++++++------
 .../expressions/MiscFunctionsSuite.scala      | 30 ++++++++++++++++---
 .../org/apache/spark/sql/functions.scala      | 20 +++++++------
 3 files changed, 48 insertions(+), 22 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
index fb176c47de446..e8dcd4b238a85 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import java.security.{GeneralSecurityException, MessageDigest, NoSuchAlgorithmException}
+import java.security.{MessageDigest, NoSuchAlgorithmException}
 import java.util.zip.CRC32
 import javax.crypto.Cipher
 import javax.crypto.spec.SecretKeySpec
@@ -446,10 +446,11 @@ case class PrintToStderr(child: Expression) extends UnaryExpression {
 }
 
 /**
- * A function that encrypts input using AES. Key lengths of 128, 192 or 256 bits can be used.
- * 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
- * Jurisdiction Policy Files are installed. If either argument is NULL or the key length is
- * not one of the permitted values, the return value is NULL.
+ * A function that encrypts input using AES. Key lengths of 128, 192 or 256 bits can be used. 192
+ * and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength Jurisdic-
+ * tion Policy Files are installed. If either argument is NULL, the result will also be null. If
+ * input is invalid, key length is not one of the permitted values or using 192/256 bits key before
+ * installing JCE, an exception will be thrown.
  */
 @ExpressionDescription(
   usage = "_FUNC_(input, key) - Encrypts input using AES.",
@@ -487,10 +488,11 @@ case class AesEncrypt(left: Expression, right: Expression)
 }
 
 /**
- * A function that decrypts input using AES. Key lengths of 128, 192 or 256 bits can be used.
- * 192 and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength
- * Jurisdiction Policy Files are installed. If either argument is NULL or the key length is
- * not one of the permitted values, the return value is NULL.
+ * A function that decrypts input using AES. Key lengths of 128, 192 or 256 bits can be used. 192
+ * and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength Jurisdic-
+ * tion Policy Files are installed. If either argument is NULL, the result will also be null. If
+ * input is invalid, key length is not one of the permitted values or using 192/256 bits key before
+ * installing JCE, an exception will be thrown.
  */
 @ExpressionDescription(
   usage = "_FUNC_(input, key) - Decrypts input using AES.",
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 53339f4947ecc..d9e9fcd1d22d5 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -148,6 +148,11 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     // both are null
     checkEvaluation(Base64(AesEncrypt(Literal.create(null, BinaryType),
       Literal.create(null, BinaryType))), null)
+
+    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
+    intercept[java.security.InvalidKeyException] {
+      evaluate(AesEncrypt(Literal("ABC".getBytes), Literal("1234567890".getBytes)))
+    }
   }
 
   test("aesDecrypt") {
@@ -165,13 +170,30 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     // both are null
     checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
       Literal.create(null, BinaryType)), null)
+
+    // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
+    intercept[java.security.InvalidKeyException] {
+      evaluate(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+        Literal("1234567890".getBytes)))
+    }
+
+    // input can not be decrypted
+    intercept[javax.crypto.IllegalBlockSizeException] {
+      evaluate(AesDecrypt(UnBase64(Literal("y6Ss+zCsdYObpCbgfWfyNW3Twewr")),
+        Literal("1234567890123456".getBytes)));
+    }
+
+    // input can not be decrypted
+    intercept[javax.crypto.BadPaddingException] {
+      evaluate(AesDecrypt(UnBase64(Literal("t6Ss+zCYObpCbgfWfyNWTw==")),
+        Literal("1234567890123456".getBytes)));
+    }
   }
 
   ignore("aesEncryptWith256bitsKey") {
-    // Before testing this, installing Java Cryptography Extension (JCE)
-    // Unlimited Strength Jurisdiction Policy Files first. Otherwise it
-    // will return `null`. Because Oracle JDK does not support 192 and 256
-    // bits key out of box.
+    // Before testing this, installing Java Cryptography Extension (JCE) Unlimited Strength Juris-
+    // diction Policy Files first. Otherwise `java.security.InvalidKeyException` will be thrown.
+    // Because Oracle JDK does not support 192 and 256 bits key out of box.
     checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
       Literal("12345678901234561234567890123456".getBytes))), "nYfCuJeRd5eD60yXDw7WEA==")
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index c5c5234dad0e4..86708bdc7f05d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -1933,10 +1933,11 @@ object functions extends LegacyFunctions {
 
   /**
    * Encrypts input using AES and Returns the result as a binary column.
-   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if
-   * Java Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed.
-   * If either argument is NULL or the key length is not one of the permitted values,
-   * the result will also be null.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if Java
+   * Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed. If
+   * either argument is NULL, the result will also be null. If input is invalid, key length is not
+   * one of the permitted values or using 192/256 bits key before installing JCE, an exception will
+   * be thrown.
    *
    * @param input binary column to encrypt input
    * @param key binary column of 128, 192 or 256 bits key
@@ -1950,10 +1951,11 @@ object functions extends LegacyFunctions {
 
   /**
    * Decrypts input using AES and Returns the result as a string column.
-   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if
-   * Java Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed.
-   * If either argument is NULL or the key length is not one of the permitted values,
-   * the result will also be null.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if Java
+   * Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed. If
+   * either argument is NULL, the result will also be null. If input is invalid, key length is not
+   * one of the permitted values or using 192/256 bits key before installing JCE, an exception will
+   * be thrown.
    *
    * @param input binary column to decrypt input
    * @param key binary column of 128, 192 or 256 bits key
@@ -1962,7 +1964,7 @@ object functions extends LegacyFunctions {
    * @since 2.0.0
    */
   def aes_decrypt(input: Column, key: Column): Column = withExpr {
-    AesEncrypt(input.expr, key.expr)
+    AesDecrypt(input.expr, key.expr)
   }
 
   //////////////////////////////////////////////////////////////////////////////////////////////

From dc760570b594d100aeda7ce181edc8a3572a59c2 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Wed, 17 Feb 2016 04:25:38 -0800
Subject: [PATCH 05/10] Python API for aes_encrypt and aes_decrypt

---
 python/pyspark/sql/functions.py | 42 +++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 0d5708526701e..89422ffa323e2 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -1052,6 +1052,48 @@ def hash(*cols):
     return Column(jc)
 
 
+@since(2.0)
+def aes_encrypt(input, key):
+    """
+    Encrypts input of given column using AES. Key lengths of 128, 192 or 256 bits can be used. 192
+    and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength Jurisdic-
+    tion Policy Files are installed. If input is invalid, key length is not one of the permitted
+    values or using 192/256 bits key before installing JCE, an exception will be thrown.
+
+    >>> df = sqlContext.createDataFrame([('ABC','1234567890123456')], ['input','key'])
+    >>> df.select(base64(aes_encrypt(df.input, df.key)).alias('aes')).collect()
+    [Row(aes=u'y6Ss+zCYObpCbgfWfyNWTw==')]
+    >>> sqlContext.sql("SELECT base64(aes_encrypt('ABC', '1234567890123456')) AS aes").collect() \
+    == df.select(base64(aes_encrypt(df.input, df.key)).alias('aes')).collect()
+    True
+    """
+    sc = SparkContext._active_spark_context
+    jc = sc._jvm.functions.aes_encrypt(_to_java_column(input), _to_java_column(key))
+    return Column(jc)
+
+
+@since(2.0)
+def aes_decrypt(input, key):
+    """
+    Decrypts input of given column using AES. Key lengths of 128, 192 or 256 bits can be used. 192
+    and 256 bits keys can be used if Java Cryptography Extension (JCE) Unlimited Strength Jurisdic-
+    tion Policy Files are installed. If input is invalid, key length is not one of the permitted
+    values or using 192/256 bits key before installing JCE, an exception will be thrown.
+
+    >>> df = sqlContext.createDataFrame([(u'y6Ss+zCYObpCbgfWfyNWTw==','1234567890123456')], \
+    ['input','key'])
+    >>> df.select(aes_decrypt(unbase64(df.input), df.key).alias('aes')).collect()
+    [Row(aes=u'ABC')]
+    >>> sqlContext.sql("SELECT aes_decrypt(unbase64('y6Ss+zCYObpCbgfWfyNWTw=='), " + \
+    "'1234567890123456') AS aes").collect() == df.select(aes_decrypt(unbase64(df.input), \
+    df.key).alias('aes')).collect()
+    True
+    """
+    sc = SparkContext._active_spark_context
+    jc = sc._jvm.functions.aes_decrypt(_to_java_column(input), _to_java_column(key))
+    return Column(jc)
+
+
 # ---------------------- String/Binary functions ------------------------------
 
 _string_functions = {

From 170b86327dec1d6f63f262a947c7f01c1f2bf2cd Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Thu, 18 Feb 2016 03:38:20 -0800
Subject: [PATCH 06/10] add function decorator ignore_unicode_prefix

to avoid python test errors
---
 python/pyspark/sql/functions.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 89422ffa323e2..a57497e48409d 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -1052,6 +1052,7 @@ def hash(*cols):
     return Column(jc)
 
 
+@ignore_unicode_prefix
 @since(2.0)
 def aes_encrypt(input, key):
     """
@@ -1063,15 +1064,13 @@ def aes_encrypt(input, key):
     >>> df = sqlContext.createDataFrame([('ABC','1234567890123456')], ['input','key'])
     >>> df.select(base64(aes_encrypt(df.input, df.key)).alias('aes')).collect()
     [Row(aes=u'y6Ss+zCYObpCbgfWfyNWTw==')]
-    >>> sqlContext.sql("SELECT base64(aes_encrypt('ABC', '1234567890123456')) AS aes").collect() \
-    == df.select(base64(aes_encrypt(df.input, df.key)).alias('aes')).collect()
-    True
     """
     sc = SparkContext._active_spark_context
     jc = sc._jvm.functions.aes_encrypt(_to_java_column(input), _to_java_column(key))
     return Column(jc)
 
 
+@ignore_unicode_prefix
 @since(2.0)
 def aes_decrypt(input, key):
     """
@@ -1084,10 +1083,6 @@ def aes_decrypt(input, key):
     ['input','key'])
     >>> df.select(aes_decrypt(unbase64(df.input), df.key).alias('aes')).collect()
     [Row(aes=u'ABC')]
-    >>> sqlContext.sql("SELECT aes_decrypt(unbase64('y6Ss+zCYObpCbgfWfyNWTw=='), " + \
-    "'1234567890123456') AS aes").collect() == df.select(aes_decrypt(unbase64(df.input), \
-    df.key).alias('aes')).collect()
-    True
     """
     sc = SparkContext._active_spark_context
     jc = sc._jvm.functions.aes_decrypt(_to_java_column(input), _to_java_column(key))

From 557108730d9342a9c63d7c843aade28e155cbfa1 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Thu, 18 Feb 2016 06:00:26 -0800
Subject: [PATCH 07/10] add test cases under DataFrame API

---
 .../spark/sql/DataFrameFunctionsSuite.scala   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index aff9efe4b2b16..0381d5728077b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -206,6 +206,30 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
       Row(2743272264L, 2180413220L))
   }
 
+  test("misc aes encrypt function") {
+    val df = Seq(("ABC", "1234567890123456")).toDF("input", "key")
+    checkAnswer(
+      df.select(base64(aes_encrypt($"input", $"key"))),
+      Row("y6Ss+zCYObpCbgfWfyNWTw==")
+    )
+    checkAnswer(
+      sql("SELECT base64(aes_encrypt('', '1234567890123456'))"),
+      Row("BQGHoM3lqYcsurCRq3PlUw==")
+    )
+  }
+
+  test("misc aes decrypt function") {
+    val df = Seq(("y6Ss+zCYObpCbgfWfyNWTw==", "1234567890123456")).toDF("input", "key")
+    checkAnswer(
+      df.select((aes_decrypt(unbase64($"input"), $"key"))),
+      Row("ABC")
+    )
+    checkAnswer(
+      sql("SELECT aes_decrypt(unbase64('BQGHoM3lqYcsurCRq3PlUw=='), '1234567890123456')"),
+      Row("")
+    )
+  }
+
   test("string function find_in_set") {
     val df = Seq(("abc,b,ab,c,def", "abc,b,ab,c,def")).toDF("a", "b")
 

From 8e2960ce68676ea7cb09683b1e1556ccf2114053 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Thu, 18 Feb 2016 17:26:48 -0800
Subject: [PATCH 08/10] test codegen version

---
 .../expressions/MiscFunctionsSuite.scala      | 47 ++++++++++++++-----
 1 file changed, 35 insertions(+), 12 deletions(-)

diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index d9e9fcd1d22d5..8509227721522 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -134,32 +134,43 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("aesEncrypt") {
-    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
-      Literal("1234567890123456".getBytes))), "y6Ss+zCYObpCbgfWfyNWTw==")
-    checkEvaluation(Base64(AesEncrypt(Literal("".getBytes),
-      Literal("1234567890123456".getBytes))), "BQGHoM3lqYcsurCRq3PlUw==")
+    val expr1 = AesEncrypt(Literal("ABC".getBytes), Literal("1234567890123456".getBytes))
+    val expr2 = AesEncrypt(Literal("".getBytes), Literal("1234567890123456".getBytes))
+
+    checkEvaluation(Base64(expr1), "y6Ss+zCYObpCbgfWfyNWTw==")
+    checkEvaluation(Base64(expr2), "BQGHoM3lqYcsurCRq3PlUw==")
 
     // input is null
     checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
       Literal("1234567890123456".getBytes)), null)
     // key is null
-    checkEvaluation(Base64(AesEncrypt(Literal("ABC".getBytes),
-      Literal.create(null, BinaryType))), null)
+    checkEvaluation(AesEncrypt(Literal("ABC".getBytes),
+      Literal.create(null, BinaryType)), null)
     // both are null
-    checkEvaluation(Base64(AesEncrypt(Literal.create(null, BinaryType),
-      Literal.create(null, BinaryType))), null)
+    checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
+      Literal.create(null, BinaryType)), null)
 
     // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
     intercept[java.security.InvalidKeyException] {
       evaluate(AesEncrypt(Literal("ABC".getBytes), Literal("1234567890".getBytes)))
     }
+
+    // check codegen
+    val instance1 = UnsafeProjection.create(Base64(expr1) :: Nil)
+    assert(instance1.apply(null).getString(0) === "y6Ss+zCYObpCbgfWfyNWTw==")
+
+    val instance2 = UnsafeProjection.create(Base64(expr2) :: Nil)
+    assert(instance2.apply(null).getString(0) === "BQGHoM3lqYcsurCRq3PlUw==")
   }
 
   test("aesDecrypt") {
-    checkEvaluation(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
-      Literal("1234567890123456".getBytes)), "ABC")
-    checkEvaluation(AesDecrypt(UnBase64(Literal("BQGHoM3lqYcsurCRq3PlUw==")),
-      Literal("1234567890123456".getBytes)), "")
+    val expr1 = AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal("1234567890123456".getBytes))
+    val expr2 = AesDecrypt(UnBase64(Literal("BQGHoM3lqYcsurCRq3PlUw==")),
+      Literal("1234567890123456".getBytes))
+
+    checkEvaluation(expr1, "ABC")
+    checkEvaluation(expr2, "")
 
     // input is null
     checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
@@ -188,6 +199,18 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       evaluate(AesDecrypt(UnBase64(Literal("t6Ss+zCYObpCbgfWfyNWTw==")),
         Literal("1234567890123456".getBytes)));
     }
+
+    // check codegen
+    val instance1 = UnsafeProjection.create(expr1 :: Nil)
+    assert(instance1.apply(null).getString(0) === "ABC")
+
+    val instance2 = UnsafeProjection.create(expr2 :: Nil)
+    assert(instance2.apply(null).getString(0) === "")
+
+    intercept[java.security.InvalidKeyException] {
+      UnsafeProjection.create(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+        Literal("1234567890".getBytes)) :: Nil).apply(null)
+    }
   }
 
   ignore("aesEncryptWith256bitsKey") {

From 60b9da7cfe13a0825d114d5a44be61588c762aff Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Fri, 19 Feb 2016 01:49:46 -0800
Subject: [PATCH 09/10] fix

---
 .../expressions/MiscFunctionsSuite.scala      | 50 +++++++++----------
 1 file changed, 23 insertions(+), 27 deletions(-)

diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 8509227721522..4f091cc9eb345 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -150,17 +150,14 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(AesEncrypt(Literal.create(null, BinaryType),
       Literal.create(null, BinaryType)), null)
 
+    val expr3 = AesEncrypt(Literal("ABC".getBytes), Literal("1234567890".getBytes))
     // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
     intercept[java.security.InvalidKeyException] {
-      evaluate(AesEncrypt(Literal("ABC".getBytes), Literal("1234567890".getBytes)))
+      evaluate(expr3)
+    }
+    intercept[java.security.InvalidKeyException] {
+      UnsafeProjection.create(expr3::Nil).apply(null)
     }
-
-    // check codegen
-    val instance1 = UnsafeProjection.create(Base64(expr1) :: Nil)
-    assert(instance1.apply(null).getString(0) === "y6Ss+zCYObpCbgfWfyNWTw==")
-
-    val instance2 = UnsafeProjection.create(Base64(expr2) :: Nil)
-    assert(instance2.apply(null).getString(0) === "BQGHoM3lqYcsurCRq3PlUw==")
   }
 
   test("aesDecrypt") {
@@ -182,34 +179,33 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(AesDecrypt(UnBase64(Literal.create(null, StringType)),
       Literal.create(null, BinaryType)), null)
 
+    val expr3 = AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal("1234567890".getBytes))
+    val expr4 = AesDecrypt(UnBase64(Literal("y6Ss+zCsdYObpCbgfWfyNW3Twewr")),
+      Literal("1234567890123456".getBytes))
+    val expr5 = AesDecrypt(UnBase64(Literal("t6Ss+zCYObpCbgfWfyNWTw==")),
+      Literal("1234567890123456".getBytes))
+
     // key length (80 bits) is not one of the permitted values (128, 192 or 256 bits)
     intercept[java.security.InvalidKeyException] {
-      evaluate(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
-        Literal("1234567890".getBytes)))
+      evaluate(expr3)
+    }
+    intercept[java.security.InvalidKeyException] {
+      UnsafeProjection.create(expr3::Nil).apply(null)
     }
-
     // input can not be decrypted
     intercept[javax.crypto.IllegalBlockSizeException] {
-      evaluate(AesDecrypt(UnBase64(Literal("y6Ss+zCsdYObpCbgfWfyNW3Twewr")),
-        Literal("1234567890123456".getBytes)));
+      evaluate(expr4)
+    }
+    intercept[javax.crypto.IllegalBlockSizeException] {
+      UnsafeProjection.create(expr4::Nil).apply(null)
     }
-
     // input can not be decrypted
     intercept[javax.crypto.BadPaddingException] {
-      evaluate(AesDecrypt(UnBase64(Literal("t6Ss+zCYObpCbgfWfyNWTw==")),
-        Literal("1234567890123456".getBytes)));
+      evaluate(expr5)
     }
-
-    // check codegen
-    val instance1 = UnsafeProjection.create(expr1 :: Nil)
-    assert(instance1.apply(null).getString(0) === "ABC")
-
-    val instance2 = UnsafeProjection.create(expr2 :: Nil)
-    assert(instance2.apply(null).getString(0) === "")
-
-    intercept[java.security.InvalidKeyException] {
-      UnsafeProjection.create(AesDecrypt(UnBase64(Literal("y6Ss+zCYObpCbgfWfyNWTw==")),
-        Literal("1234567890".getBytes)) :: Nil).apply(null)
+    intercept[javax.crypto.BadPaddingException] {
+      UnsafeProjection.create(expr5::Nil).apply(null)
     }
   }
 

From 0856fb0cad6f19a1a1ae2411729b7ad644bd65c7 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Fri, 19 Feb 2016 02:07:32 -0800
Subject: [PATCH 10/10] nit

---
 .../sql/catalyst/expressions/MiscFunctionsSuite.scala     | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
index 4f091cc9eb345..67f2dc457d333 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscFunctionsSuite.scala
@@ -156,7 +156,7 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       evaluate(expr3)
     }
     intercept[java.security.InvalidKeyException] {
-      UnsafeProjection.create(expr3::Nil).apply(null)
+      UnsafeProjection.create(expr3 :: Nil).apply(null)
     }
   }
 
@@ -191,21 +191,21 @@ class MiscFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       evaluate(expr3)
     }
     intercept[java.security.InvalidKeyException] {
-      UnsafeProjection.create(expr3::Nil).apply(null)
+      UnsafeProjection.create(expr3 :: Nil).apply(null)
     }
     // input can not be decrypted
     intercept[javax.crypto.IllegalBlockSizeException] {
       evaluate(expr4)
     }
     intercept[javax.crypto.IllegalBlockSizeException] {
-      UnsafeProjection.create(expr4::Nil).apply(null)
+      UnsafeProjection.create(expr4 :: Nil).apply(null)
     }
     // input can not be decrypted
     intercept[javax.crypto.BadPaddingException] {
       evaluate(expr5)
     }
     intercept[javax.crypto.BadPaddingException] {
-      UnsafeProjection.create(expr5::Nil).apply(null)
+      UnsafeProjection.create(expr5 :: Nil).apply(null)
     }
   }