Fix padding issue in Base64

Joe-Abraham · Mar 25, 2024 · 97586b1 · 97586b1
1 parent c722866
commit 97586b1
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 18 deletions.
diff --git a/velox/common/encode/Base64.cpp b/velox/common/encode/Base64.cpp
@@ -80,13 +80,6 @@ constexpr const Base64::ReverseIndex kBase64UrlReverseIndexTable = {
     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
     255};
 
-constexpr bool checkForwardIndex(
-    uint8_t idx,
-    const Base64::Charset& charset,
-    const Base64::ReverseIndex& table) {
-  return (table[static_cast<uint8_t>(charset[idx])] == idx) &&
-      (idx > 0 ? checkForwardIndex(idx - 1, charset, table) : true);
-}
 // Verify that for every entry in kBase64Charset, the corresponding entry
 // in kBase64ReverseIndexTable is correct.
 static_assert(
@@ -343,7 +336,7 @@ size_t Base64::calculateDecodedSize(const char* data, size_t& size) {
     // If padded, ensure that the string length is a multiple of the encoded
     // block size
     if (size % kEncodedBlockSize != 0) {
-      throw Base64Exception(
+      throw EncoderException(
           "Base64::decode() - invalid input string: "
           "string length is not a multiple of the encoded block size.");
     }
@@ -386,7 +379,7 @@ size_t Base64::decodeImpl(
 
   auto needed = calculateDecodedSize(src, src_len);
   if (dst_len < needed) {
-    throw Base64Exception(
+    throw EncoderException(
         "Base64::decode() - invalid output string: "
         "output string is too small.");
   }
@@ -396,10 +389,10 @@ size_t Base64::decodeImpl(
     // Each character of the 4 encode 6 bits of the original, grab each with
     // the appropriate shifts to rebuild the original and then split that back
     // into the original 8 bit bytes.
-    uint32_t last = (Base64ReverseLookup(src[0], reverse_lookup) << 18) |
-        (Base64ReverseLookup(src[1], reverse_lookup) << 12) |
-        (Base64ReverseLookup(src[2], reverse_lookup) << 6) |
-        Base64ReverseLookup(src[3], reverse_lookup);
+    uint32_t last = (baseReverseLookup(kBase, src[0], reverse_lookup) << 18) |
+        (baseReverseLookup(kBase, src[1], reverse_lookup) << 12) |
+        (baseReverseLookup(kBase, src[2], reverse_lookup) << 6) |
+        baseReverseLookup(kBase, src[3], reverse_lookup);
     dst[0] = (last >> 16) & 0xff;
     dst[1] = (last >> 8) & 0xff;
     dst[2] = last & 0xff;
@@ -408,14 +401,14 @@ size_t Base64::decodeImpl(
   // Handle the last 2-4 characters.  This is similar to the above, but the
   // last 2 characters may or may not exist.
   DCHECK(src_len >= 2);
-  uint32_t last = (Base64ReverseLookup(src[0], reverse_lookup) << 18) |
-      (Base64ReverseLookup(src[1], reverse_lookup) << 12);
+  uint32_t last = (baseReverseLookup(kBase, src[0], reverse_lookup) << 18) |
+      (baseReverseLookup(kBase, src[1], reverse_lookup) << 12);
   dst[0] = (last >> 16) & 0xff;
   if (src_len > 2) {
-    last |= Base64ReverseLookup(src[2], reverse_lookup) << 6;
+    last |= baseReverseLookup(kBase, src[2], reverse_lookup) << 6;
     dst[1] = (last >> 8) & 0xff;
     if (src_len > 3) {
-      last |= Base64ReverseLookup(src[3], reverse_lookup);
+      last |= baseReverseLookup(kBase, src[3], reverse_lookup);
       dst[2] = last & 0xff;
     }
   }

diff --git a/velox/common/encode/Base64.h b/velox/common/encode/Base64.h
@@ -21,6 +21,7 @@
 
 #include <folly/Range.h>
 #include <folly/io/IOBuf.h>
+#include "velox/common/encode/EncoderUtils.h"
 
 namespace facebook::velox::encoding {
 

diff --git a/velox/functions/prestosql/BinaryFunctions.h b/velox/functions/prestosql/BinaryFunctions.h
@@ -293,7 +293,7 @@ struct FromBase64Function {
           encoding::Base64::calculateDecodedSize(input.data(), inputSize));
       encoding::Base64::decode(
           input.data(), inputSize, result.data(), result.size());
-    } catch (const encoding::Base64Exception& e) {
+    } catch (const encoding::EncoderException& e) {
       VELOX_USER_FAIL(e.what());
     }
   }

diff --git a/velox/functions/prestosql/tests/BinaryFunctionsTest.cpp b/velox/functions/prestosql/tests/BinaryFunctionsTest.cpp
@@ -430,6 +430,10 @@ TEST_F(BinaryFunctionsTest, fromBase64) {
   EXPECT_EQ(
       "Hello World from Velox!",
       fromBase64("SGVsbG8gV29ybGQgZnJvbSBWZWxveCE="));
+  // Check encoded strings without padding
+  EXPECT_EQ("a", fromBase64("YQ"));
+  EXPECT_EQ("ab", fromBase64("YWI"));
+  EXPECT_EQ("abcd", fromBase64("YWJjZA"));
 
   EXPECT_THROW(fromBase64("YQ="), VeloxUserError);
   EXPECT_THROW(fromBase64("YQ==="), VeloxUserError);