From 3ce41a0e15eee79683c44d7da2ed8507ff3c6454 Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Thu, 19 Jan 2017 09:31:45 -0500
Subject: [PATCH] Painless: Add augmentation to string for base 64 (#22665)

We don't want to expose `String#getBytes` which is required for
`Base64.getEncoder.encode` to work because we're worried about
character sets. This adds `encodeBase64` and `decodeBase64`
methods to `String` in Painless that are duals of one another
such that:
`someString == someString.encodeBase64().decodeBase64()`.

Both methods work with the UTF-8 encoding of the string.

Closes #22648
---
 .../elasticsearch/painless/Augmentation.java    | 17 +++++++++++++++++
 .../org/elasticsearch/painless/java.lang.txt    |  3 ++-
 .../org/elasticsearch/painless/StringTests.java | 15 +++++++++++++--
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/Augmentation.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/Augmentation.java
index 9302f3c899cc2..28f26249707aa 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/Augmentation.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/Augmentation.java
@@ -19,7 +19,9 @@
 
 package org.elasticsearch.painless;
 
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -486,4 +488,19 @@ public static String replaceFirst(CharSequence receiver, Pattern pattern, Functi
     private static int initialBufferForReplaceWith(CharSequence seq) {
         return seq.length() + 16;
     }
+
+    /**
+     * Encode a String in Base64. Use {@link Base64.Encoder#encodeToString(byte[])} if you have to encode bytes rather than a string.
+     */
+    public static String encodeBase64(String receiver) {
+        return Base64.getEncoder().encodeToString(receiver.getBytes(StandardCharsets.UTF_8));
+    }
+
+    /**
+     * Decode some Base64 bytes and build a UTF-8 encoded string. Use {@link Base64.Decoder#decode(String)} if you'd prefer bytes to work
+     * with bytes.
+     */
+    public static String decodeBase64(String receiver) {
+        return new String(Base64.getDecoder().decode(receiver.getBytes(StandardCharsets.UTF_8)), StandardCharsets.UTF_8);
+    }
 }
diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
index 13f28d3ebeb61..a1cde1711bc51 100644
--- a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
+++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
@@ -756,6 +756,8 @@ class String -> java.lang.String extends CharSequence,Comparable,Object {
   boolean contentEquals(CharSequence)
   String copyValueOf(char[])
   String copyValueOf(char[],int,int)
+  String decodeBase64*()
+  String encodeBase64*()
   boolean endsWith(String)
   boolean equalsIgnoreCase(String)
   String format(Locale,String,def[])
@@ -1101,4 +1103,3 @@ class UnsupportedOperationException -> java.lang.UnsupportedOperationException e
   UnsupportedOperationException <init>()
   UnsupportedOperationException <init>(String)
 }
-
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/StringTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/StringTests.java
index b5b3e2cfbf6ac..da4558a693a0d 100644
--- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/StringTests.java
+++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/StringTests.java
@@ -19,12 +19,13 @@
 
 package org.elasticsearch.painless;
 
-import static org.elasticsearch.painless.WriterConstants.MAX_INDY_STRING_CONCAT_ARGS;
-
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
 
+import static java.util.Collections.singletonMap;
+import static org.elasticsearch.painless.WriterConstants.MAX_INDY_STRING_CONCAT_ARGS;
+
 public class StringTests extends ScriptTestCase {
 
     public void testAppend() {
@@ -238,4 +239,14 @@ public void testComplexCompoundAssignment() {
     public void testAppendStringIntoMap() {
         assertEquals("nullcat", exec("def a = new HashMap(); a.cat += 'cat'"));
     }
+
+    public void testBase64Augmentations() {
+        assertEquals("Y2F0", exec("'cat'.encodeBase64()"));
+        assertEquals("cat", exec("'Y2F0'.decodeBase64()"));
+        assertEquals("6KiA6Kqe", exec("'\u8A00\u8A9E'.encodeBase64()"));
+        assertEquals("\u8A00\u8A9E", exec("'6KiA6Kqe'.decodeBase64()"));
+
+        String rando = randomRealisticUnicodeOfLength(between(5, 1000));
+        assertEquals(rando, exec("params.rando.encodeBase64().decodeBase64()", singletonMap("rando", rando), true));
+    }
 }