#include "wasm_abi/mini_table/container_decode.h"

#include <setjmp.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "upb/base/descriptor_constants.h"
#include "upb/base/status.h"
#include "upb/mem/alloc.h"
#include "upb/mem/arena.h"
#include "upb/mem/internal/arena.h"
#include "upb/mini_descriptor/build_enum.h"
#include "upb/mini_descriptor/decode.h"
#include "upb/mini_descriptor/link.h"
#include "upb/mini_table/enum.h"
#include "upb/mini_table/extension.h"
#include "upb/mini_table/extension_registry.h"
#include "upb/mini_table/field.h"
#include "upb/mini_table/message.h"
#include "upb/wire/eps_copy_input_stream.h"
#include "wasm_abi/mini_table/container.h"

// Must be last.
#include "upb/port/def.inc"
#include "wasm_abi/def.inc"

// Lengths of MiniDescriptor strings are explicit in the format and can be
// immediately bound checked before read. Varint length is not known. To prevent
// OOB reads, the buffer passed to varint decoding routines must have
// MAX_VARINT_LEN in-bound characters.
#define MAX_VARINT_LEN 5

// Set maximums for allocating an array of pointers.
// Since container format uses u32 varints, our absolute max size is capped at
// u32 max, but we also need to account for platform limits.
#define MAX_PTR_ARRAY_COUNT UPB_MIN(UINT32_MAX, (SIZE_MAX / sizeof(intptr_t)))

typedef struct {
  upb_EpsCopyInputStream input;
  upb_Arena* arena;
  upb_Status* status;
  void* msg_buf;
  jmp_buf err;
} upb_ContainerDecoder;

UPB_FORCEINLINE void _upb_ContainerDecoder_Init(upb_ContainerDecoder* d,
                                                const char** ptr, size_t len,
                                                upb_Arena* arena,
                                                upb_Status* status) {
  upb_EpsCopyInputStream_Init(&d->input, ptr, len);
  d->arena = arena;
  d->status = status;
  d->msg_buf = NULL;
}

UPB_FORCEINLINE void _upb_ContainerDecoder_TearDown(upb_ContainerDecoder* d) {
  upb_gfree(d->msg_buf);
}

UPB_NORETURN UPB_NOINLINE UPB_WASM_PRESERVE_ALL static void*
_upb_ContainerDecoder_ErrorJmp(upb_ContainerDecoder* d, const char* msg) {
  UPB_ASSERT(d->status->ok);
  upb_Status_SetErrorMessage(d->status, msg);
  UPB_LONGJMP(d->err, 1);
}

// Propagate errors from descriptor parsers.
UPB_NORETURN UPB_NOINLINE UPB_WASM_PRESERVE_ALL static void*
_upb_ContainerDecoder_PropagateErrorJmp(upb_ContainerDecoder* d) {
  UPB_ASSERT(!d->status->ok);
  UPB_LONGJMP(d->err, 1);
}

// Varint decoding routines borrowed from UPB wire decode.
// Impl eliminates a mask op on the first byte but is incompatible with
// "padding" bytes. This may be required in the future for Container "linking".

typedef struct {
  const char* ptr;
  uint32_t val;
} _upb_DecodeLongVarint32Return;

UPB_NOINLINE
static _upb_DecodeLongVarint32Return _upb_ContainerDecoder_DecodeLongVarint32(
    const char* ptr, uint32_t val) {
  _upb_DecodeLongVarint32Return ret = {NULL, 0};
  uint32_t byte;
  for (int i = 1; i < MAX_VARINT_LEN; i++) {
    byte = (uint8_t)ptr[i];
    val += (byte - 1) << (i * 7);
    if (!(byte & 0x80)) {
      ret.ptr = ptr + i + 1;
      ret.val = val;
      return ret;
    }
  }
  return ret;
}

// Assumed that MAX_VARINT_LEN after ptr are valid to read.
// Caller must validate bounds after call.
UPB_FORCEINLINE
const char* _upb_DecodeVarint32(upb_ContainerDecoder* d, const char* ptr,
                                uint32_t* val) {
  UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(&d->input, MAX_VARINT_LEN);
  uint32_t byte = (uint8_t)*ptr;
  if (UPB_LIKELY((byte & 0x80) == 0)) {
    *val = byte;
    return ptr + 1;
  } else {
    _upb_DecodeLongVarint32Return res =
        _upb_ContainerDecoder_DecodeLongVarint32(ptr, byte);
    if (!res.ptr) _upb_ContainerDecoder_ErrorJmp(d, "Malformed varint");
    *val = res.val;
    return res.ptr;
  }
}

// Container format uses 3 types of (unsigned) ints: Counts, Indexes and Sizes.
// Each type requires different bounds checking, encapsulated in the below
// methods.
// Methods also validate that EOF was not reached while decoding int.

// Decode `count` and verify `count <= limit`.
UPB_FORCEINLINE const char* _upb_ContainerDecode_DecodeCount(
    upb_ContainerDecoder* d, const char* ptr, uint32_t* count, size_t limit) {
  if (UPB_UNLIKELY(upb_EpsCopyInputStream_IsDone(&d->input, &ptr))) {
    _upb_ContainerDecoder_ErrorJmp(d, "Unexpected EOF");
  }
  ptr = _upb_DecodeVarint32(d, ptr, count);
  if (UPB_UNLIKELY(!upb_EpsCopyInputStream_CheckSize(&d->input, ptr, 0) ||
                   *count > limit)) {
    _upb_ContainerDecoder_ErrorJmp(d, "Invalid section count");
  }
  return ptr;
}

// Decode `length` and ensure `length` characters remain in the input.
UPB_FORCEINLINE const char* _upb_ContainerDecoder_DecodeLength(
    upb_ContainerDecoder* d, const char* ptr, uint32_t* len) {
  if (UPB_UNLIKELY(upb_EpsCopyInputStream_IsDone(&d->input, &ptr))) {
    _upb_ContainerDecoder_ErrorJmp(d, "Unexpected EOF");
  }
  ptr = _upb_DecodeVarint32(d, ptr, len);
  if (UPB_UNLIKELY(*len >= INT32_MAX || !upb_EpsCopyInputStream_CheckSize(
                                            &d->input, ptr, (int)*len))) {
    _upb_ContainerDecoder_ErrorJmp(d, "Invalid Descriptor length");
  }
  return ptr;
}

// Decode `idx` and validate that `limit < idx`.
UPB_FORCEINLINE const char* _upb_ContainerDecoder_DecodeIndex(
    upb_ContainerDecoder* d, const char* ptr, uint32_t* idx, uint32_t limit) {
  if (UPB_UNLIKELY(upb_EpsCopyInputStream_IsDone(&d->input, &ptr))) {
    _upb_ContainerDecoder_ErrorJmp(d, "Unexpected EOF");
  }
  ptr = _upb_DecodeVarint32(d, ptr, idx);
  if (UPB_UNLIKELY(!upb_EpsCopyInputStream_CheckSize(&d->input, ptr, 0) ||
                   *idx >= limit)) {
    _upb_ContainerDecoder_ErrorJmp(d, "Invalid Container index");
  }
  return ptr;
}

UPB_FORCEINLINE const char* _upb_ContainerDecode_DecodeEnumSection(
    upb_ContainerDecoder* d, const char* ptr, upb_MiniTableContainer* result) {
  ptr = _upb_ContainerDecode_DecodeCount(d, ptr, &result->enum_count,
                                         MAX_PTR_ARRAY_COUNT);
  result->enums = upb_Arena_Malloc(
      d->arena, result->enum_count * sizeof(upb_MiniTableEnum*));
  if (UPB_UNLIKELY(result->enums == NULL)) {
    _upb_ContainerDecoder_ErrorJmp(d, "Out of memory");
  }

  for (uint32_t enum_idx = 0; enum_idx < result->enum_count; enum_idx++) {
    uint32_t enum_len;
    ptr = _upb_ContainerDecoder_DecodeLength(d, ptr, &enum_len);
    result->enums[enum_idx] =
        upb_MiniTableEnum_Build(ptr, enum_len, d->arena, d->status);
    if (UPB_UNLIKELY(d->status->ok == false)) {
      _upb_ContainerDecoder_PropagateErrorJmp(d);
    } else if (UPB_UNLIKELY(result->enums[enum_idx] == NULL)) {
      _upb_ContainerDecoder_ErrorJmp(
          d, "Null Enum returned from parser without error");
    }
    ptr += enum_len;
  }
  return ptr;
}

UPB_FORCEINLINE const char* _upb_ContainerDecode_DecodeMessageSection(
    upb_ContainerDecoder* d, const char* ptr, upb_MiniTableContainer* result) {
  ptr = _upb_ContainerDecode_DecodeCount(d, ptr, &result->message_count,
                                         MAX_PTR_ARRAY_COUNT);
  result->messages = upb_Arena_Malloc(
      d->arena, result->message_count * sizeof(upb_MiniTable*));
  if (UPB_UNLIKELY(result->messages == NULL)) {
    _upb_ContainerDecoder_ErrorJmp(d, "Out of memory");
  }

  size_t msg_buf_len = 0;
  for (uint32_t message_idx = 0; message_idx < result->message_count;
       message_idx++) {
    uint32_t message_len;
    ptr = _upb_ContainerDecoder_DecodeLength(d, ptr, &message_len);
    result->messages[message_idx] = upb_MiniTable_BuildWithBuf(
        ptr, message_len, kUpb_MiniTablePlatform_Native, d->arena, &d->msg_buf,
        &msg_buf_len, d->status);
    if (UPB_UNLIKELY(!d->status->ok)) {
      _upb_ContainerDecoder_PropagateErrorJmp(d);
    } else if (UPB_UNLIKELY(result->messages[message_idx] == NULL)) {
      _upb_ContainerDecoder_ErrorJmp(
          d, "Null Message returned from parser without error");
    }
    ptr += message_len;
  }
  return ptr;
}

UPB_FORCEINLINE const char* _upb_ContainerDecode_DecodeExtensionSection(
    upb_ContainerDecoder* d, const char* ptr, upb_MiniTableContainer* result) {
  ptr = _upb_ContainerDecode_DecodeCount(d, ptr, &result->extension_count,
                                         MAX_PTR_ARRAY_COUNT);
  result->extensions = upb_Arena_Malloc(
      d->arena, result->extension_count * sizeof(upb_MiniTableExtension*));
  result->registry = upb_ExtensionRegistry_New(d->arena);
  if (UPB_UNLIKELY(result->extensions == NULL || result->registry == NULL)) {
    _upb_ContainerDecoder_ErrorJmp(d, "Out of memory");
  }

  uint32_t ext_idx = 0;
  uint32_t remaining_exts = result->extension_count;
  while (remaining_exts > 0) {
    uint32_t parent_idx;
    ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &parent_idx,
                                            result->message_count);
    upb_MiniTable* parent = result->messages[parent_idx];

    uint32_t num_packed_exts;
    ptr = _upb_ContainerDecode_DecodeCount(d, ptr, &num_packed_exts,
                                           remaining_exts);
    remaining_exts -= num_packed_exts;

    for (uint32_t packed_idx = 0; packed_idx < num_packed_exts; packed_idx++) {
      uint32_t ext_len;
      ptr = _upb_ContainerDecoder_DecodeLength(d, ptr, &ext_len);
      upb_MiniTableExtension* ext = upb_MiniTableExtension_Build(
          ptr, ext_len, parent, d->arena, d->status);
      if (UPB_UNLIKELY(!d->status->ok)) {
        _upb_ContainerDecoder_PropagateErrorJmp(d);
      } else if (UPB_UNLIKELY(ext == NULL)) {
        _upb_ContainerDecoder_ErrorJmp(
            d, "Null Extension returned from parser without error");
      }
      ptr += ext_len;
      result->extensions[ext_idx++] = ext;
      upb_ExtensionRegistryStatus reg_status =
          upb_ExtensionRegistry_Add(result->registry, ext);
      if (UPB_UNLIKELY(reg_status != kUpb_ExtensionRegistryStatus_Ok)) {
        _upb_ContainerDecoder_ErrorJmp(d, "Extension registry error");
      }
      switch (upb_MiniTableExtension_CType(ext)) {
        case kUpb_CType_Message: {
          uint32_t sub_idx;
          ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &sub_idx,
                                                  result->message_count);
          upb_MiniTable* sub = result->messages[sub_idx];
          if (!upb_MiniTableExtension_SetSubMessage(ext, sub)) {
            _upb_ContainerDecoder_ErrorJmp(
                d, "Failed to set extension submessage");
          }
          break;
        }
        case kUpb_CType_Enum: {
          const upb_MiniTableField* field = upb_MiniTableExtension_ToField(ext);
          if (!upb_MiniTableField_IsClosedEnum(field)) {
            break;
          }
          uint32_t sub_idx;
          ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &sub_idx,
                                                  result->enum_count);
          upb_MiniTableEnum* sub = result->enums[sub_idx];
          if (!upb_MiniTableExtension_SetSubEnum(ext, sub)) {
            _upb_ContainerDecoder_ErrorJmp(d,
                                           "Failed to set extension subenum");
          }
          break;
        }
        default:;
          // Passthrough.
      }
    }
  }
  return ptr;
}

UPB_FORCEINLINE const char* _upb_ContainerDecode_DecodeLinkageSection(
    upb_ContainerDecoder* d, const char* ptr, upb_MiniTableContainer* result) {
  while (!upb_EpsCopyInputStream_IsDone(&d->input, &ptr)) {
    // Since we assume all Messages should be fully linked, parent index is only
    // necessary due to cost of checking if a MiniTable requires linking.
    // Returning a bool from MiniTable_Build would solve this problem.
    uint32_t parent_idx;
    ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &parent_idx,
                                            result->message_count);
    upb_MiniTable* parent = result->messages[parent_idx];

    for (int i = 0; i < upb_MiniTable_FieldCount(parent); i++) {
      upb_MiniTableField* field =
          (upb_MiniTableField*)upb_MiniTable_GetFieldByIndex(parent, i);
      switch (upb_MiniTableField_CType(field)) {
        case kUpb_CType_Message: {
          uint32_t sub_idx;
          ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &sub_idx,
                                                  result->message_count);
          upb_MiniTable* sub = result->messages[sub_idx];
          if (UPB_UNLIKELY(!upb_MiniTable_SetSubMessage(parent, field, sub))) {
            _upb_ContainerDecoder_ErrorJmp(d, "Failure during linking");
          }
          break;
        }
        case kUpb_CType_Enum: {
          // Skip open enums, which are not represented with MiniTableEnums and
          // don't need linkage.
          if (!upb_MiniTableField_IsClosedEnum(field)) {
            break;
          }
          uint32_t sub_idx;
          ptr = _upb_ContainerDecoder_DecodeIndex(d, ptr, &sub_idx,
                                                  result->enum_count);
          upb_MiniTableEnum* sub = result->enums[sub_idx];
          if (UPB_UNLIKELY(!upb_MiniTable_SetSubEnum(parent, field, sub))) {
            _upb_ContainerDecoder_ErrorJmp(d, "Failure during linking");
          }
          break;
        }
        default:;
      }
    }
  }

  return ptr;
}

static upb_MiniTableContainer* upb_MiniTableContainer_BuildWithSetJmp(
    upb_ContainerDecoder* d, const char* ptr, upb_Arena* arena) {
  if (UPB_SETJMP(d->err) != 0) {
    _upb_ContainerDecoder_TearDown(d);
    return NULL;
  }

  upb_MiniTableContainer* result = (upb_MiniTableContainer*)upb_Arena_Malloc(
      arena, sizeof(upb_MiniTableContainer));

  ptr = _upb_ContainerDecode_DecodeEnumSection(d, ptr, result);
  ptr = _upb_ContainerDecode_DecodeMessageSection(d, ptr, result);
  ptr = _upb_ContainerDecode_DecodeExtensionSection(d, ptr, result);
  ptr = _upb_ContainerDecode_DecodeLinkageSection(d, ptr, result);

  _upb_ContainerDecoder_TearDown(d);

  return result;
}

upb_MiniTableContainer* upb_MiniTableContainer_Build(const char* data,
                                                     size_t len,
                                                     upb_Arena* arena,
                                                     upb_Status* status) {
  // Minimum format has length 4: Version byte, then 3 empty section counts.
  if (UPB_UNLIKELY(len < 4)) {
    upb_Status_SetErrorMessage(status, "No data provided");
    return NULL;
  }
  if (UPB_UNLIKELY(data[0] != kUpb_EncodedVersion_ContainerV1)) {
    upb_Status_SetErrorMessage(status, "Container has malformed header");
    return NULL;
  }

  upb_ContainerDecoder d;
  const char* ptr = data + 1;
  len = len - 1;
  _upb_ContainerDecoder_Init(&d, &ptr, len, arena, status);

  return upb_MiniTableContainer_BuildWithSetJmp(&d, ptr, arena);
}
