/*

 * Copyright 2020 Cypress Semiconductor Corporation.

 * Copyright 2020 Richard Hughes <richard@hughsie.com>

 *

 * SPDX-License-Identifier: LGPL-2.1-or-later

 */

#include "config.h"

#include <string.h>

#include "fu-ccgx-dmc-firmware.h"

#include "fu-ccgx-dmc-struct.h"

struct _FuCcgxDmcFirmware {

  FuFirmwareClass parent_instance;

  GPtrArray *image_records;

  GBytes *fwct_blob;

  GBytes *custom_meta_blob;

  guint32 row_data_offset_start;

  guint32 fw_data_size;

};

G_DEFINE_TYPE(FuCcgxDmcFirmware, fu_ccgx_dmc_firmware, FU_TYPE_FIRMWARE)

#define DMC_FWCT_MAX_SIZE     2048

#define DMC_HASH_SIZE       32

#define DMC_CUSTOM_META_LENGTH_FIELD_SIZE 2

static void

fu_ccgx_dmc_firmware_record_free(FuCcgxDmcFirmwareRecord *rcd)

{

  if (rcd->seg_records != NULL)

    g_ptr_array_unref(rcd->seg_records);

  g_free(rcd);

}

static void

fu_ccgx_dmc_firmware_segment_record_free(FuCcgxDmcFirmwareSegmentRecord *rcd)

{

  if (rcd->data_records != NULL)

    g_ptr_array_unref(rcd->data_records);

  g_free(rcd);

}

G_DEFINE_AUTOPTR_CLEANUP_FUNC(FuCcgxDmcFirmwareRecord, fu_ccgx_dmc_firmware_record_free)

G_DEFINE_AUTOPTR_CLEANUP_FUNC(FuCcgxDmcFirmwareSegmentRecord,

            fu_ccgx_dmc_firmware_segment_record_free)

GPtrArray *

fu_ccgx_dmc_firmware_get_image_records(FuCcgxDmcFirmware *self)

{

  g_return_val_if_fail(FU_IS_CCGX_DMC_FIRMWARE(self), NULL);

  return self->image_records;

}

GBytes *

fu_ccgx_dmc_firmware_get_fwct_record(FuCcgxDmcFirmware *self)

{

  g_return_val_if_fail(FU_IS_CCGX_DMC_FIRMWARE(self), NULL);

  return self->fwct_blob;

}

GBytes *

fu_ccgx_dmc_firmware_get_custom_meta_record(FuCcgxDmcFirmware *self)

{

  g_return_val_if_fail(FU_IS_CCGX_DMC_FIRMWARE(self), NULL);

  return self->custom_meta_blob;

}

guint32

fu_ccgx_dmc_firmware_get_fw_data_size(FuCcgxDmcFirmware *self)

{

  g_return_val_if_fail(FU_IS_CCGX_DMC_FIRMWARE(self), 0);

  return self->fw_data_size;

}

static void

fu_ccgx_dmc_firmware_export(FuFirmware *firmware, FuFirmwareExportFlags flags, XbBuilderNode *bn)

{

  FuCcgxDmcFirmware *self = FU_CCGX_DMC_FIRMWARE(firmware);

  if (flags & FU_FIRMWARE_EXPORT_FLAG_INCLUDE_DEBUG) {

    fu_xmlb_builder_insert_kx(bn, "fw_data_size", self->fw_data_size);

    fu_xmlb_builder_insert_kx(bn, "image_records", self->image_records->len);

  }

}

static gboolean

fu_ccgx_dmc_firmware_parse_segment(FuFirmware *firmware,

           GInputStream *stream,

           FuCcgxDmcFirmwareRecord *img_rcd,

           gsize *seg_off,

           FuFirmwareParseFlags flags,

           GError **error)

{

  FuCcgxDmcFirmware *self = FU_CCGX_DMC_FIRMWARE(firmware);

  gsize row_off;

  g_autoptr(GChecksum) csum = g_checksum_new(G_CHECKSUM_SHA256);

  /* set row data offset in current image */

  row_off = self->row_data_offset_start + img_rcd->img_offset;

  /* parse segment in image */

  img_rcd->seg_records =

      g_ptr_array_new_with_free_func((GFreeFunc)fu_ccgx_dmc_firmware_segment_record_free);

  for (guint32 i = 0; i < img_rcd->num_img_segments; i++) {

    guint16 row_size_bytes = 0;

    g_autoptr(FuCcgxDmcFirmwareSegmentRecord) seg_rcd = NULL;

    g_autoptr(GByteArray) st_info = NULL;

    /* read segment info  */

    seg_rcd = g_new0(FuCcgxDmcFirmwareSegmentRecord, 1);

    st_info =

        fu_struct_ccgx_dmc_fwct_segmentation_info_parse_stream(stream, *seg_off, error);

    if (st_info == NULL)

      return FALSE;

    seg_rcd->start_row =

        fu_struct_ccgx_dmc_fwct_segmentation_info_get_start_row(st_info);

    seg_rcd->num_rows = fu_struct_ccgx_dmc_fwct_segmentation_info_get_num_rows(st_info);

    /* calculate actual row size */

    row_size_bytes = img_rcd->row_size * 64;

    /* create data record array in segment record */

    seg_rcd->data_records =

        g_ptr_array_new_with_free_func((GDestroyNotify)g_bytes_unref);

    /* read row data in segment */

    for (int row = 0; row < seg_rcd->num_rows; row++) {

      g_autoptr(GBytes) data_rcd = NULL;

      data_rcd = fu_input_stream_read_bytes(stream,

                    row_off,

                    row_size_bytes,

                    NULL,

                    error);

      if (data_rcd == NULL)

        return FALSE;

      /* update hash */

      g_checksum_update(csum,

            (guchar *)g_bytes_get_data(data_rcd, NULL),

            g_bytes_get_size(data_rcd));

      /* add row data to data record */

      g_ptr_array_add(seg_rcd->data_records, g_steal_pointer(&data_rcd));

      /* increment row data offset */

      row_off += row_size_bytes;

    }

    /* add segment record to segment array */

    g_ptr_array_add(img_rcd->seg_records, g_steal_pointer(&seg_rcd));

    /* increment segment info offset */

    *seg_off += st_info->len;

  }

  /* check checksum */

  if ((flags & FU_FIRMWARE_PARSE_FLAG_IGNORE_CHECKSUM) == 0) {

    guint8 csumbuf[DMC_HASH_SIZE] = {0x0};

    gsize csumbufsz = sizeof(csumbuf);

    g_checksum_get_digest(csum, csumbuf, &csumbufsz);

    if (memcmp(csumbuf, img_rcd->img_digest, DMC_HASH_SIZE) != 0) {

      g_set_error_literal(error,

              FWUPD_ERROR,

              FWUPD_ERROR_NOT_SUPPORTED,

              "invalid hash");

      return FALSE;

    }

  }

  /* success */

  return TRUE;

}

static gboolean

fu_ccgx_dmc_firmware_parse_image(FuFirmware *firmware,

         guint8 image_count,

         GInputStream *stream,

         FuFirmwareParseFlags flags,

         GError **error)

{

  FuCcgxDmcFirmware *self = FU_CCGX_DMC_FIRMWARE(firmware);

  gsize img_off = FU_STRUCT_CCGX_DMC_FWCT_INFO_SIZE;

  gsize seg_off = FU_STRUCT_CCGX_DMC_FWCT_INFO_SIZE +

      image_count * FU_STRUCT_CCGX_DMC_FWCT_IMAGE_INFO_SIZE;

  /* set initial segment info offset */

  for (guint32 i = 0; i < image_count; i++) {

    g_autoptr(FuCcgxDmcFirmwareRecord) img_rcd = NULL;

    g_autoptr(GByteArray) st_img = NULL;

    /* read image info */

    img_rcd = g_new0(FuCcgxDmcFirmwareRecord, 1);

    st_img = fu_struct_ccgx_dmc_fwct_image_info_parse_stream(stream, img_off, error);

    if (st_img == NULL)

      return FALSE;

    img_rcd->row_size = fu_struct_ccgx_dmc_fwct_image_info_get_row_size(st_img);

    if (img_rcd->row_size == 0) {

      g_set_error(error,

            FWUPD_ERROR,

            FWUPD_ERROR_NOT_SUPPORTED,

            "invalid row size 0x%x",

            img_rcd->row_size);

      return FALSE;

    }

    img_rcd->img_offset = fu_struct_ccgx_dmc_fwct_image_info_get_img_offset(st_img);

    img_rcd->num_img_segments =

        fu_struct_ccgx_dmc_fwct_image_info_get_num_img_segments(st_img);

    /* segments are optional */

    if (img_rcd->num_img_segments > 0) {

      gsize img_digestsz = 0;

      const guint8 *img_digest;

      img_digest =

          fu_struct_ccgx_dmc_fwct_image_info_get_img_digest(st_img,

                        &img_digestsz);

      if (!fu_memcpy_safe((guint8 *)&img_rcd->img_digest,

              sizeof(img_rcd->img_digest),

              0x0, /* dst */

              img_digest,

              img_digestsz,

              0, /* src */

              img_digestsz,

              error))

        return FALSE;

      /* parse segment */

      if (!fu_ccgx_dmc_firmware_parse_segment(firmware,

                stream,

                img_rcd,

                &seg_off,

                flags,

                error))

        return FALSE;

      /* add image record to image record array */

      g_ptr_array_add(self->image_records, g_steal_pointer(&img_rcd));

    }

    /* increment image offset */

    img_off += FU_STRUCT_CCGX_DMC_FWCT_IMAGE_INFO_SIZE;

  }

  return TRUE;

}

static gboolean

fu_ccgx_dmc_firmware_validate(FuFirmware *firmware,

            GInputStream *stream,

            gsize offset,

            GError **error)

{

  return fu_struct_ccgx_dmc_fwct_info_validate_stream(stream, offset, error);

}

static gboolean

fu_ccgx_dmc_firmware_parse(FuFirmware *firmware,

         GInputStream *stream,

         FuFirmwareParseFlags flags,

         GError **error)

{

  FuCcgxDmcFirmware *self = FU_CCGX_DMC_FIRMWARE(firmware);

  gsize streamsz = 0;

  guint16 hdr_size = 0;

  guint16 mdbufsz = 0;

  guint32 hdr_composite_version = 0;

  guint8 hdr_image_count = 0;

  g_autoptr(GByteArray) st_hdr = NULL;

  /* parse */

  st_hdr = fu_struct_ccgx_dmc_fwct_info_parse_stream(stream, 0x0, error);

  if (st_hdr == NULL)

    return FALSE;

  /* check fwct size */

  hdr_size = fu_struct_ccgx_dmc_fwct_info_get_size(st_hdr);

  if (hdr_size > DMC_FWCT_MAX_SIZE || hdr_size == 0) {

    g_set_error(error,

          FWUPD_ERROR,

          FWUPD_ERROR_NOT_SUPPORTED,

          "invalid dmc fwct size, expected <= 0x%x, got 0x%x",

          (guint)DMC_FWCT_MAX_SIZE,

          (guint)hdr_size);

    return FALSE;

  }

  /* set version */

  hdr_composite_version = fu_struct_ccgx_dmc_fwct_info_get_composite_version(st_hdr);

  if (hdr_composite_version != 0)

    fu_firmware_set_version_raw(firmware, hdr_composite_version);

  /* read fwct data */

  self->fwct_blob = fu_input_stream_read_bytes(stream, 0x0, hdr_size, NULL, error);

  if (self->fwct_blob == NULL)

    return FALSE;

  /* create custom meta binary */

  if (!fu_input_stream_read_u16(stream, hdr_size, &mdbufsz, G_LITTLE_ENDIAN, error)) {

    g_prefix_error(error, "failed to read metadata size: ");

    return FALSE;

  }

  if (mdbufsz > 0) {

    self->custom_meta_blob =

        fu_input_stream_read_bytes(stream, hdr_size + 2, mdbufsz, NULL, error);

    if (self->custom_meta_blob == NULL)

      return FALSE;

  }

  /* set row data start offset */

  self->row_data_offset_start = hdr_size + DMC_CUSTOM_META_LENGTH_FIELD_SIZE + mdbufsz;

  if (!fu_input_stream_size(stream, &streamsz, error))

    return FALSE;

  self->fw_data_size = streamsz - self->row_data_offset_start;

  /* parse image */

  hdr_image_count = fu_struct_ccgx_dmc_fwct_info_get_image_count(st_hdr);

  if (!fu_ccgx_dmc_firmware_parse_image(firmware, hdr_image_count, stream, flags, error))

    return FALSE;

  /* success */

  return TRUE;

}

static GByteArray *

fu_ccgx_dmc_firmware_write(FuFirmware *firmware, GError **error)

{

  g_autoptr(GByteArray) buf = g_byte_array_new();

  g_autoptr(GByteArray) st_hdr = fu_struct_ccgx_dmc_fwct_info_new();

  g_autoptr(GPtrArray) images = fu_firmware_get_images(firmware);

  /* add header */

  fu_struct_ccgx_dmc_fwct_info_set_size(

      st_hdr,

      FU_STRUCT_CCGX_DMC_FWCT_INFO_SIZE +

    (images->len * (FU_STRUCT_CCGX_DMC_FWCT_IMAGE_INFO_SIZE +

        FU_STRUCT_CCGX_DMC_FWCT_SEGMENTATION_INFO_SIZE)));

  fu_struct_ccgx_dmc_fwct_info_set_version(st_hdr, 0x2);

  fu_struct_ccgx_dmc_fwct_info_set_custom_meta_type(st_hdr, 0x3);

  fu_struct_ccgx_dmc_fwct_info_set_cdtt_version(st_hdr, 0x1);

  fu_struct_ccgx_dmc_fwct_info_set_device_id(st_hdr, 0x1);

  fu_struct_ccgx_dmc_fwct_info_set_composite_version(st_hdr,

                 fu_firmware_get_version_raw(firmware));

  fu_struct_ccgx_dmc_fwct_info_set_image_count(st_hdr, images->len);

  g_byte_array_append(buf, st_hdr->data, st_hdr->len);

  /* add image headers */

  for (guint i = 0; i < images->len; i++) {

    g_autoptr(GByteArray) st_img = fu_struct_ccgx_dmc_fwct_image_info_new();

    fu_struct_ccgx_dmc_fwct_image_info_set_device_type(st_img, 0x2);

    fu_struct_ccgx_dmc_fwct_image_info_set_img_type(st_img, 0x1);

    fu_struct_ccgx_dmc_fwct_image_info_set_row_size(st_img, 0x1);

    fu_struct_ccgx_dmc_fwct_image_info_set_fw_version(st_img, 0x330006d2);

    fu_struct_ccgx_dmc_fwct_image_info_set_app_version(st_img, 0x14136161);

    fu_struct_ccgx_dmc_fwct_image_info_set_num_img_segments(st_img, 0x1);

    g_byte_array_append(buf, st_img->data, st_img->len);

  }

  /* add segments */

  for (guint i = 0; i < images->len; i++) {

    FuFirmware *img = g_ptr_array_index(images, i);

    g_autoptr(GByteArray) st_info = fu_struct_ccgx_dmc_fwct_segmentation_info_new();

    g_autoptr(FuChunkArray) chunks = NULL;

    g_autoptr(GBytes) img_bytes = fu_firmware_get_bytes(img, error);

    if (img_bytes == NULL)

      return NULL;

    chunks = fu_chunk_array_new_from_bytes(img_bytes,

                   FU_CHUNK_ADDR_OFFSET_NONE,

                   FU_CHUNK_PAGESZ_NONE,

                   64);

    fu_struct_ccgx_dmc_fwct_segmentation_info_set_num_rows(

        st_info,

        MAX(fu_chunk_array_length(chunks), 1));

    g_byte_array_append(buf, st_info->data, st_info->len);

  }

  /* metadata */

  fu_byte_array_append_uint16(buf, 0x1, G_LITTLE_ENDIAN);

  fu_byte_array_append_uint8(buf, 0xff);

  /* add image headers */

  for (guint i = 0; i < images->len; i++) {

    FuFirmware *img = g_ptr_array_index(images, i);

    gsize csumbufsz = DMC_HASH_SIZE;

    gsize img_offset = FU_STRUCT_CCGX_DMC_FWCT_INFO_SIZE +

           (i * FU_STRUCT_CCGX_DMC_FWCT_IMAGE_INFO_SIZE);

    guint8 csumbuf[DMC_HASH_SIZE] = {0x0};

    g_autoptr(GChecksum) csum = g_checksum_new(G_CHECKSUM_SHA256);

    g_autoptr(GBytes) img_bytes = NULL;

    g_autoptr(GBytes) img_padded = NULL;

    g_autoptr(FuChunkArray) chunks = NULL;

    img_bytes = fu_firmware_get_bytes(img, error);

    if (img_bytes == NULL)

      return NULL;

    chunks = fu_chunk_array_new_from_bytes(img_bytes,

                   FU_CHUNK_ADDR_OFFSET_NONE,

                   FU_CHUNK_PAGESZ_NONE,

                   64);

    img_padded =

        fu_bytes_pad(img_bytes, MAX(fu_chunk_array_length(chunks), 1) * 64, 0xFF);

    fu_byte_array_append_bytes(buf, img_padded);

    g_checksum_update(csum,

          (const guchar *)g_bytes_get_data(img_padded, NULL),

          g_bytes_get_size(img_padded));

    g_checksum_get_digest(csum, csumbuf, &csumbufsz);

    /* update checksum */

    if (!fu_memcpy_safe(buf->data,

            buf->len, /* dst */

            img_offset +

          FU_STRUCT_CCGX_DMC_FWCT_IMAGE_INFO_OFFSET_IMG_DIGEST,

            csumbuf,

            sizeof(csumbuf),

            0x0, /* src */

            sizeof(csumbuf),

            error))

      return NULL;

  }

  return g_steal_pointer(&buf);

}

static gchar *

fu_ccgx_dmc_firmware_convert_version(FuFirmware *firmware, guint64 version_raw)

{

  return fu_version_from_uint32(version_raw, fu_firmware_get_version_format(firmware));

}

static void

fu_ccgx_dmc_firmware_init(FuCcgxDmcFirmware *self)

{

  self->image_records =

      g_ptr_array_new_with_free_func((GFreeFunc)fu_ccgx_dmc_firmware_record_free);

  fu_firmware_add_flag(FU_FIRMWARE(self), FU_FIRMWARE_FLAG_HAS_CHECKSUM);

  fu_firmware_set_version_format(FU_FIRMWARE(self), FWUPD_VERSION_FORMAT_QUAD);

}

static void

fu_ccgx_dmc_firmware_finalize(GObject *object)

{

  FuCcgxDmcFirmware *self = FU_CCGX_DMC_FIRMWARE(object);

  if (self->fwct_blob != NULL)

    g_bytes_unref(self->fwct_blob);

  if (self->custom_meta_blob != NULL)

    g_bytes_unref(self->custom_meta_blob);

  if (self->image_records != NULL)

    g_ptr_array_unref(self->image_records);

  G_OBJECT_CLASS(fu_ccgx_dmc_firmware_parent_class)->finalize(object);

}

static void

fu_ccgx_dmc_firmware_class_init(FuCcgxDmcFirmwareClass *klass)

{

  GObjectClass *object_class = G_OBJECT_CLASS(klass);

  FuFirmwareClass *firmware_class = FU_FIRMWARE_CLASS(klass);

  firmware_class->convert_version = fu_ccgx_dmc_firmware_convert_version;

  object_class->finalize = fu_ccgx_dmc_firmware_finalize;

  firmware_class->validate = fu_ccgx_dmc_firmware_validate;

  firmware_class->parse = fu_ccgx_dmc_firmware_parse;

  firmware_class->write = fu_ccgx_dmc_firmware_write;

  firmware_class->export = fu_ccgx_dmc_firmware_export;

}

FuFirmware *

fu_ccgx_dmc_firmware_new(void)

{

  return FU_FIRMWARE(g_object_new(FU_TYPE_CCGX_DMC_FIRMWARE, NULL));

}