src/iceberg/util/base64.cc

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include "iceberg/util/base64.h"

#include <array>
#include <cstdint>

namespace iceberg {

namespace {

// Shared base64 encode logic. The encode alphabet maps a 6-bit value -> ASCII char.
// When `pad` is true, the output is padded with '=' to a multiple of 4 characters.
std::string Base64EncodeWithAlphabet(std::string_view input, std::string_view alphabet,
                                     bool pad) {
  const size_t full_blocks = input.size() / 3;
  const size_t remainder = input.size() % 3;
  const size_t output_size =
      full_blocks * 4 + (remainder == 0 ? 0 : (pad ? 4 : remainder + 1));
  std::string output(output_size, '\0');

  const auto* bytes = reinterpret_cast<const uint8_t*>(input.data());
  size_t out = 0;
  for (size_t i = 0; i < full_blocks * 3; i += 3) {
    const uint32_t block = (static_cast<uint32_t>(bytes[i]) << 16) |
                           (static_cast<uint32_t>(bytes[i + 1]) << 8) |
                           static_cast<uint32_t>(bytes[i + 2]);
    output[out++] = alphabet[(block >> 18) & 0x3F];
    output[out++] = alphabet[(block >> 12) & 0x3F];
    output[out++] = alphabet[(block >> 6) & 0x3F];
    output[out++] = alphabet[block & 0x3F];
  }

  if (remainder > 0) {
    const size_t i = full_blocks * 3;
    uint32_t block = static_cast<uint32_t>(bytes[i]) << 16;
    if (remainder == 2) {
      block |= static_cast<uint32_t>(bytes[i + 1]) << 8;
    }
    output[out++] = alphabet[(block >> 18) & 0x3F];
    output[out++] = alphabet[(block >> 12) & 0x3F];
    if (remainder == 2) {
      output[out++] = alphabet[(block >> 6) & 0x3F];
    }
    while (out < output_size) {
      output[out++] = '=';
    }
  }

  return output;
}

// Shared base64 decode logic. The decode table maps ASCII char -> 6-bit value.
// 0xFF means invalid character.
Result<std::string> Base64DecodeWithTable(std::string_view input,
                                          const std::array<uint8_t, 256>& table) {
  auto padded_size = input.size();
  if (auto padding_pos = input.find('='); padding_pos != std::string_view::npos) {
    auto padding = input.substr(padding_pos);
    if (padding.find_first_not_of('=') != std::string_view::npos) {
      return InvalidArgument("Invalid base64 padding");
    }
    auto padding_size = padding.size();
    if (padding_size > 2 || input.size() % 4 != 0) {
      return InvalidArgument("Invalid base64 padding");
    }
  }

  // Strip trailing padding after validating its count and placement.
  while (!input.empty() && input.back() == '=') {
    input.remove_suffix(1);
  }
  auto unpadded_size = input.size();
  if (unpadded_size % 4 == 1 ||
      (padded_size != unpadded_size && unpadded_size % 4 == 0)) {
    return InvalidArgument("Invalid base64 length");
  }
  if (input.empty()) {
    return std::string{};
  }

  std::string output;
  output.reserve((input.size() * 3) / 4);

  uint32_t buffer = 0;
  int bits_collected = 0;

  for (char c : input) {
    uint8_t val = table[static_cast<uint8_t>(c)];
    if (val == 0xFF) {
      return InvalidArgument("Invalid base64 character: '{}'", c);
    }
    buffer = (buffer << 6) | val;
    bits_collected += 6;
    if (bits_collected >= 8) {
      bits_collected -= 8;
      output.push_back(static_cast<char>((buffer >> bits_collected) & 0xFF));
    }
  }

  return output;
}

// Standard base64 alphabet: A-Z=0-25, a-z=26-51, 0-9=52-61, +=62, /=63
constexpr std::string_view kBase64Chars =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

// Base64url alphabet: same as standard but '-'=62, '_'=63 (RFC 4648 §5)
constexpr std::string_view kBase64UrlChars =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

// Standard base64 decode table: A-Z=0-25, a-z=26-51, 0-9=52-61, +=62, /=63
constexpr std::array<uint8_t, 256> kBase64DecodeTable = [] {
  std::array<uint8_t, 256> table{};
  table.fill(0xFF);
  for (int i = 0; i < 26; ++i) {
    table[static_cast<size_t>('A' + i)] = static_cast<uint8_t>(i);
    table[static_cast<size_t>('a' + i)] = static_cast<uint8_t>(26 + i);
  }
  for (int i = 0; i < 10; ++i) {
    table[static_cast<size_t>('0' + i)] = static_cast<uint8_t>(52 + i);
  }
  table[static_cast<size_t>('+')] = 62;
  table[static_cast<size_t>('/')] = 63;
  return table;
}();

// Base64url decode table: same as standard but '-'=62, '_'=63 (RFC 4648 §5)
constexpr std::array<uint8_t, 256> kBase64UrlDecodeTable = [] {
  auto table = kBase64DecodeTable;
  table[static_cast<size_t>('+')] = 0xFF;  // '+' is invalid in base64url
  table[static_cast<size_t>('/')] = 0xFF;  // '/' is invalid in base64url
  table[static_cast<size_t>('-')] = 62;
  table[static_cast<size_t>('_')] = 63;
  return table;
}();

}  // namespace

std::string Base64::Encode(std::string_view data) {
  return Base64EncodeWithAlphabet(data, kBase64Chars, /*pad=*/true);
}

Result<std::string> Base64::Decode(std::string_view encoded) {
  return Base64DecodeWithTable(encoded, kBase64DecodeTable);
}

std::string Base64::UrlEncode(std::string_view data) {
  return Base64EncodeWithAlphabet(data, kBase64UrlChars, /*pad=*/false);
}

Result<std::string> Base64::UrlDecode(std::string_view encoded) {
  return Base64DecodeWithTable(encoded, kBase64UrlDecodeTable);
}

}  // namespace iceberg