// SPDX-License-Identifier: GPL-2.0-only

/* Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. */

#include <linux/devcoredump.h>
#include <linux/firmware.h>
#include <linux/limits.h>
#include <linux/mhi.h>
#include <linux/minmax.h>
#include <linux/mod_devicetable.h>
#include <linux/overflow.h>
#include <linux/types.h>
#include <linux/vmalloc.h>
#include <linux/workqueue.h>

#include "sahara.h"

#define SAHARA_HELLO_CMD                0x1  /* Min protocol version 1.0 */
#define SAHARA_HELLO_RESP_CMD           0x2  /* Min protocol version 1.0 */
#define SAHARA_READ_DATA_CMD            0x3  /* Min protocol version 1.0 */
#define SAHARA_END_OF_IMAGE_CMD         0x4  /* Min protocol version 1.0 */
#define SAHARA_DONE_CMD                 0x5  /* Min protocol version 1.0 */
#define SAHARA_DONE_RESP_CMD            0x6  /* Min protocol version 1.0 */
#define SAHARA_RESET_CMD                0x7  /* Min protocol version 1.0 */
#define SAHARA_RESET_RESP_CMD           0x8  /* Min protocol version 1.0 */
#define SAHARA_MEM_DEBUG_CMD            0x9  /* Min protocol version 2.0 */
#define SAHARA_MEM_READ_CMD             0xa  /* Min protocol version 2.0 */
#define SAHARA_CMD_READY_CMD            0xb  /* Min protocol version 2.1 */
#define SAHARA_SWITCH_MODE_CMD          0xc  /* Min protocol version 2.1 */
#define SAHARA_EXECUTE_CMD              0xd  /* Min protocol version 2.1 */
#define SAHARA_EXECUTE_RESP_CMD         0xe  /* Min protocol version 2.1 */
#define SAHARA_EXECUTE_DATA_CMD         0xf  /* Min protocol version 2.1 */
#define SAHARA_MEM_DEBUG64_CMD          0x10 /* Min protocol version 2.5 */
#define SAHARA_MEM_READ64_CMD           0x11 /* Min protocol version 2.5 */
#define SAHARA_READ_DATA64_CMD          0x12 /* Min protocol version 2.8 */
#define SAHARA_RESET_STATE_CMD          0x13 /* Min protocol version 2.9 */
#define SAHARA_WRITE_DATA_CMD           0x14 /* Min protocol version 3.0 */

#define SAHARA_PACKET_MAX_SIZE          0xffffU /* MHI_MAX_MTU */
#define SAHARA_TRANSFER_MAX_SIZE        0x80000
#define SAHARA_READ_MAX_SIZE            0xfff0U /* Avoid unaligned requests */
#define SAHARA_NUM_TX_BUF               DIV_ROUND_UP(SAHARA_TRANSFER_MAX_SIZE,\
                                                        SAHARA_PACKET_MAX_SIZE)
#define SAHARA_IMAGE_ID_NONE            U32_MAX

#define SAHARA_VERSION                  2
#define SAHARA_SUCCESS                  0
#define SAHARA_TABLE_ENTRY_STR_LEN      20

#define SAHARA_MODE_IMAGE_TX_PENDING    0x0
#define SAHARA_MODE_IMAGE_TX_COMPLETE   0x1
#define SAHARA_MODE_MEMORY_DEBUG        0x2
#define SAHARA_MODE_COMMAND             0x3

#define SAHARA_HELLO_LENGTH             0x30
#define SAHARA_READ_DATA_LENGTH         0x14
#define SAHARA_END_OF_IMAGE_LENGTH      0x10
#define SAHARA_DONE_LENGTH              0x8
#define SAHARA_RESET_LENGTH             0x8
#define SAHARA_MEM_DEBUG64_LENGTH       0x18
#define SAHARA_MEM_READ64_LENGTH        0x18

struct sahara_packet {
        __le32 cmd;
        __le32 length;

        union {
                struct {
                        __le32 version;
                        __le32 version_compat;
                        __le32 max_length;
                        __le32 mode;
                } hello;
                struct {
                        __le32 version;
                        __le32 version_compat;
                        __le32 status;
                        __le32 mode;
                } hello_resp;
                struct {
                        __le32 image;
                        __le32 offset;
                        __le32 length;
                } read_data;
                struct {
                        __le32 image;
                        __le32 status;
                } end_of_image;
                struct {
                        __le64 table_address;
                        __le64 table_length;
                } memory_debug64;
                struct {
                        __le64 memory_address;
                        __le64 memory_length;
                } memory_read64;
        };
};

struct sahara_debug_table_entry64 {
        __le64  type;
        __le64  address;
        __le64  length;
        char    description[SAHARA_TABLE_ENTRY_STR_LEN];
        char    filename[SAHARA_TABLE_ENTRY_STR_LEN];
};

struct sahara_dump_table_entry {
        u64     type;
        u64     address;
        u64     length;
        char    description[SAHARA_TABLE_ENTRY_STR_LEN];
        char    filename[SAHARA_TABLE_ENTRY_STR_LEN];
};

#define SAHARA_DUMP_V1_MAGIC 0x1234567890abcdef
#define SAHARA_DUMP_V1_VER   1
struct sahara_memory_dump_meta_v1 {
        u64     magic;
        u64     version;
        u64     dump_size;
        u64     table_size;
};

/*
 * Layout of crashdump provided to user via devcoredump
 *              +------------------------------------------+
 *              |         Crashdump Meta structure         |
 *              | type: struct sahara_memory_dump_meta_v1  |
 *              +------------------------------------------+
 *              |             Crashdump Table              |
 *              | type: array of struct                    |
 *              |       sahara_dump_table_entry            |
 *              |                                          |
 *              |                                          |
 *              +------------------------------------------+
 *              |                Crashdump                 |
 *              |                                          |
 *              |                                          |
 *              |                                          |
 *              |                                          |
 *              |                                          |
 *              +------------------------------------------+
 *
 * First is the metadata header. Userspace can use the magic number to verify
 * the content type, and then check the version for the rest of the format.
 * New versions should keep the magic number location/value, and version
 * location, but increment the version value.
 *
 * For v1, the metadata lists the size of the entire dump (header + table +
 * dump) and the size of the table. Then the dump image table, which describes
 * the contents of the dump. Finally all the images are listed in order, with
 * no deadspace in between. Userspace can use the sizes listed in the image
 * table to reconstruct the individual images.
 */

struct sahara_context {
        struct sahara_packet            *tx[SAHARA_NUM_TX_BUF];
        struct sahara_packet            *rx;
        struct work_struct              fw_work;
        struct work_struct              dump_work;
        struct work_struct              read_data_work;
        struct mhi_device               *mhi_dev;
        const char * const              *image_table;
        u32                             table_size;
        u32                             active_image_id;
        const struct firmware           *firmware;
        u64                             dump_table_address;
        u64                             dump_table_length;
        size_t                          rx_size;
        size_t                          rx_size_requested;
        void                            *mem_dump;
        size_t                          mem_dump_sz;
        struct sahara_dump_table_entry  *dump_image;
        u64                             dump_image_offset;
        void                            *mem_dump_freespace;
        u64                             dump_images_left;
        u32                             read_data_offset;
        u32                             read_data_length;
        bool                            is_mem_dump_mode;
        bool                            non_streaming;
};

static const char * const aic100_image_table[] = {
        [1]  = "qcom/aic100/fw1.bin",
        [2]  = "qcom/aic100/fw2.bin",
        [4]  = "qcom/aic100/fw4.bin",
        [5]  = "qcom/aic100/fw5.bin",
        [6]  = "qcom/aic100/fw6.bin",
        [8]  = "qcom/aic100/fw8.bin",
        [9]  = "qcom/aic100/fw9.bin",
        [10] = "qcom/aic100/fw10.bin",
};

static const char * const aic200_image_table[] = {
        [5]  = "qcom/aic200/uefi.elf",
        [12] = "qcom/aic200/aic200-nsp.bin",
        [23] = "qcom/aic200/aop.mbn",
        [32] = "qcom/aic200/tz.mbn",
        [33] = "qcom/aic200/hypvm.mbn",
        [38] = "qcom/aic200/xbl_config.elf",
        [39] = "qcom/aic200/aic200_abl.elf",
        [40] = "qcom/aic200/apdp.mbn",
        [41] = "qcom/aic200/devcfg.mbn",
        [42] = "qcom/aic200/sec.elf",
        [43] = "qcom/aic200/aic200-hlos.elf",
        [49] = "qcom/aic200/shrm.elf",
        [50] = "qcom/aic200/cpucp.elf",
        [51] = "qcom/aic200/aop_devcfg.mbn",
        [54] = "qcom/aic200/qupv3fw.elf",
        [57] = "qcom/aic200/cpucp_dtbs.elf",
        [62] = "qcom/aic200/uefi_dtbs.elf",
        [63] = "qcom/aic200/xbl_ac_config.mbn",
        [64] = "qcom/aic200/tz_ac_config.mbn",
        [65] = "qcom/aic200/hyp_ac_config.mbn",
        [66] = "qcom/aic200/pdp.elf",
        [67] = "qcom/aic200/pdp_cdb.elf",
        [68] = "qcom/aic200/sdi.mbn",
        [69] = "qcom/aic200/dcd.mbn",
        [73] = "qcom/aic200/gearvm.mbn",
        [74] = "qcom/aic200/sti.bin",
        [76] = "qcom/aic200/tz_qti_config.mbn",
        [78] = "qcom/aic200/pvs.bin",
};

static bool is_streaming(struct sahara_context *context)
{
        return !context->non_streaming;
}

static int sahara_find_image(struct sahara_context *context, u32 image_id)
{
        int ret;

        if (image_id == context->active_image_id)
                return 0;

        if (context->active_image_id != SAHARA_IMAGE_ID_NONE) {
                dev_err(&context->mhi_dev->dev, "image id %d is not valid as %d is active\n",
                        image_id, context->active_image_id);
                return -EINVAL;
        }

        if (image_id >= context->table_size || !context->image_table[image_id]) {
                dev_err(&context->mhi_dev->dev, "request for unknown image: %d\n", image_id);
                return -EINVAL;
        }

        /*
         * This image might be optional. The device may continue without it.
         * Only the device knows. Suppress error messages that could suggest an
         * a problem when we were actually able to continue.
         */
        ret = firmware_request_nowarn(&context->firmware,
                                      context->image_table[image_id],
                                      &context->mhi_dev->dev);
        if (ret) {
                dev_dbg(&context->mhi_dev->dev, "request for image id %d / file %s failed %d\n",
                        image_id, context->image_table[image_id], ret);
                return ret;
        }

        context->active_image_id = image_id;

        return 0;
}

static void sahara_release_image(struct sahara_context *context)
{
        if (context->active_image_id != SAHARA_IMAGE_ID_NONE)
                release_firmware(context->firmware);
        context->active_image_id = SAHARA_IMAGE_ID_NONE;
}

static void sahara_send_reset(struct sahara_context *context)
{
        int ret;

        context->is_mem_dump_mode = false;
        context->read_data_offset = 0;
        context->read_data_length = 0;

        context->tx[0]->cmd = cpu_to_le32(SAHARA_RESET_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_RESET_LENGTH);

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_RESET_LENGTH, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to send reset response %d\n", ret);
}

static void sahara_hello(struct sahara_context *context)
{
        int ret;

        dev_dbg(&context->mhi_dev->dev,
                "HELLO cmd received. length:%d version:%d version_compat:%d max_length:%d mode:%d\n",
                le32_to_cpu(context->rx->length),
                le32_to_cpu(context->rx->hello.version),
                le32_to_cpu(context->rx->hello.version_compat),
                le32_to_cpu(context->rx->hello.max_length),
                le32_to_cpu(context->rx->hello.mode));

        if (le32_to_cpu(context->rx->length) != SAHARA_HELLO_LENGTH) {
                dev_err(&context->mhi_dev->dev, "Malformed hello packet - length %d\n",
                        le32_to_cpu(context->rx->length));
                return;
        }
        if (le32_to_cpu(context->rx->hello.version) != SAHARA_VERSION) {
                dev_err(&context->mhi_dev->dev, "Unsupported hello packet - version %d\n",
                        le32_to_cpu(context->rx->hello.version));
                return;
        }

        if (le32_to_cpu(context->rx->hello.mode) != SAHARA_MODE_IMAGE_TX_PENDING &&
            le32_to_cpu(context->rx->hello.mode) != SAHARA_MODE_IMAGE_TX_COMPLETE &&
            le32_to_cpu(context->rx->hello.mode) != SAHARA_MODE_MEMORY_DEBUG) {
                dev_err(&context->mhi_dev->dev, "Unsupported hello packet - mode %d\n",
                        le32_to_cpu(context->rx->hello.mode));
                return;
        }

        context->tx[0]->cmd = cpu_to_le32(SAHARA_HELLO_RESP_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_HELLO_LENGTH);
        context->tx[0]->hello_resp.version = cpu_to_le32(SAHARA_VERSION);
        context->tx[0]->hello_resp.version_compat = cpu_to_le32(SAHARA_VERSION);
        context->tx[0]->hello_resp.status = cpu_to_le32(SAHARA_SUCCESS);
        context->tx[0]->hello_resp.mode = context->rx->hello_resp.mode;

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_HELLO_LENGTH, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to send hello response %d\n", ret);
}

static int read_data_helper(struct sahara_context *context, int buf_index)
{
        enum mhi_flags mhi_flag;
        u32 pkt_data_len;
        int ret;

        pkt_data_len = min(context->read_data_length, SAHARA_PACKET_MAX_SIZE);

        memcpy(context->tx[buf_index],
               &context->firmware->data[context->read_data_offset],
               pkt_data_len);

        context->read_data_offset += pkt_data_len;
        context->read_data_length -= pkt_data_len;

        if (is_streaming(context) || !context->read_data_length)
                mhi_flag = MHI_EOT;
        else
                mhi_flag = MHI_CHAIN;

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE,
                            context->tx[buf_index], pkt_data_len, mhi_flag);
        if (ret) {
                dev_err(&context->mhi_dev->dev, "Unable to send read_data response %d\n", ret);
                return ret;
        }

        return 0;
}

static void sahara_read_data(struct sahara_context *context)
{
        u32 image_id, data_offset, data_len;
        int ret;
        int i;

        dev_dbg(&context->mhi_dev->dev,
                "READ_DATA cmd received. length:%d image:%d offset:%d data_length:%d\n",
                le32_to_cpu(context->rx->length),
                le32_to_cpu(context->rx->read_data.image),
                le32_to_cpu(context->rx->read_data.offset),
                le32_to_cpu(context->rx->read_data.length));

        if (le32_to_cpu(context->rx->length) != SAHARA_READ_DATA_LENGTH) {
                dev_err(&context->mhi_dev->dev, "Malformed read_data packet - length %d\n",
                        le32_to_cpu(context->rx->length));
                return;
        }

        image_id = le32_to_cpu(context->rx->read_data.image);
        data_offset = le32_to_cpu(context->rx->read_data.offset);
        data_len = le32_to_cpu(context->rx->read_data.length);

        ret = sahara_find_image(context, image_id);
        if (ret) {
                sahara_send_reset(context);
                return;
        }

        /*
         * Image is released when the device is done with it via
         * SAHARA_END_OF_IMAGE_CMD. sahara_send_reset() will either cause the
         * device to retry the operation with a modification, or decide to be
         * done with the image and trigger SAHARA_END_OF_IMAGE_CMD.
         * release_image() is called from SAHARA_END_OF_IMAGE_CMD. processing
         * and is not needed here on error.
         */

        if (context->non_streaming && data_len > SAHARA_TRANSFER_MAX_SIZE) {
                dev_err(&context->mhi_dev->dev, "Malformed read_data packet - data len %d exceeds max xfer size %d\n",
                        data_len, SAHARA_TRANSFER_MAX_SIZE);
                sahara_send_reset(context);
                return;
        }

        if (data_offset >= context->firmware->size) {
                dev_err(&context->mhi_dev->dev, "Malformed read_data packet - data offset %d exceeds file size %zu\n",
                        data_offset, context->firmware->size);
                sahara_send_reset(context);
                return;
        }

        if (size_add(data_offset, data_len) > context->firmware->size) {
                dev_err(&context->mhi_dev->dev, "Malformed read_data packet - data offset %d and length %d exceeds file size %zu\n",
                        data_offset, data_len, context->firmware->size);
                sahara_send_reset(context);
                return;
        }

        context->read_data_offset = data_offset;
        context->read_data_length = data_len;

        if (is_streaming(context)) {
                schedule_work(&context->read_data_work);
                return;
        }

        for (i = 0; i < SAHARA_NUM_TX_BUF && context->read_data_length; ++i) {
                ret = read_data_helper(context, i);
                if (ret)
                        break;
        }
}

static void sahara_end_of_image(struct sahara_context *context)
{
        int ret;

        dev_dbg(&context->mhi_dev->dev,
                "END_OF_IMAGE cmd received. length:%d image:%d status:%d\n",
                le32_to_cpu(context->rx->length),
                le32_to_cpu(context->rx->end_of_image.image),
                le32_to_cpu(context->rx->end_of_image.status));

        if (le32_to_cpu(context->rx->length) != SAHARA_END_OF_IMAGE_LENGTH) {
                dev_err(&context->mhi_dev->dev, "Malformed end_of_image packet - length %d\n",
                        le32_to_cpu(context->rx->length));
                return;
        }

        if (context->active_image_id != SAHARA_IMAGE_ID_NONE &&
            le32_to_cpu(context->rx->end_of_image.image) != context->active_image_id) {
                dev_err(&context->mhi_dev->dev, "Malformed end_of_image packet - image %d is not the active image\n",
                        le32_to_cpu(context->rx->end_of_image.image));
                return;
        }

        sahara_release_image(context);

        if (le32_to_cpu(context->rx->end_of_image.status))
                return;

        context->tx[0]->cmd = cpu_to_le32(SAHARA_DONE_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_DONE_LENGTH);

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_DONE_LENGTH, MHI_EOT);
        if (ret)
                dev_dbg(&context->mhi_dev->dev, "Unable to send done response %d\n", ret);
}

static void sahara_memory_debug64(struct sahara_context *context)
{
        int ret;

        dev_dbg(&context->mhi_dev->dev,
                "MEMORY DEBUG64 cmd received. length:%d table_address:%#llx table_length:%#llx\n",
                le32_to_cpu(context->rx->length),
                le64_to_cpu(context->rx->memory_debug64.table_address),
                le64_to_cpu(context->rx->memory_debug64.table_length));

        if (le32_to_cpu(context->rx->length) != SAHARA_MEM_DEBUG64_LENGTH) {
                dev_err(&context->mhi_dev->dev, "Malformed memory debug64 packet - length %d\n",
                        le32_to_cpu(context->rx->length));
                return;
        }

        context->dump_table_address = le64_to_cpu(context->rx->memory_debug64.table_address);
        context->dump_table_length = le64_to_cpu(context->rx->memory_debug64.table_length);

        if (context->dump_table_length % sizeof(struct sahara_debug_table_entry64) != 0 ||
            !context->dump_table_length) {
                dev_err(&context->mhi_dev->dev, "Malformed memory debug64 packet - table length %lld\n",
                        context->dump_table_length);
                return;
        }

        /*
         * From this point, the protocol flips. We make memory_read requests to
         * the device, and the device responds with the raw data. If the device
         * has an error, it will send an End of Image command. First we need to
         * request the memory dump table so that we know where all the pieces
         * of the dump are that we can consume.
         */

        context->is_mem_dump_mode = true;

        /*
         * Assume that the table is smaller than our MTU so that we can read it
         * in one shot. The spec does not put an upper limit on the table, but
         * no known device will exceed this.
         */
        if (context->dump_table_length > SAHARA_PACKET_MAX_SIZE) {
                dev_err(&context->mhi_dev->dev, "Memory dump table length %lld exceeds supported size. Discarding dump\n",
                        context->dump_table_length);
                sahara_send_reset(context);
                return;
        }

        context->tx[0]->cmd = cpu_to_le32(SAHARA_MEM_READ64_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_MEM_READ64_LENGTH);
        context->tx[0]->memory_read64.memory_address = cpu_to_le64(context->dump_table_address);
        context->tx[0]->memory_read64.memory_length = cpu_to_le64(context->dump_table_length);

        context->rx_size_requested = context->dump_table_length;

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_MEM_READ64_LENGTH, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to send read for dump table %d\n", ret);
}

static void sahara_processing(struct work_struct *work)
{
        struct sahara_context *context = container_of(work, struct sahara_context, fw_work);
        int ret;

        switch (le32_to_cpu(context->rx->cmd)) {
        case SAHARA_HELLO_CMD:
                sahara_hello(context);
                break;
        case SAHARA_READ_DATA_CMD:
                sahara_read_data(context);
                break;
        case SAHARA_END_OF_IMAGE_CMD:
                sahara_end_of_image(context);
                break;
        case SAHARA_DONE_RESP_CMD:
                /* Intentional do nothing as we don't need to exit an app */
                break;
        case SAHARA_RESET_RESP_CMD:
                /* Intentional do nothing as we don't need to exit an app */
                break;
        case SAHARA_MEM_DEBUG64_CMD:
                sahara_memory_debug64(context);
                break;
        default:
                dev_err(&context->mhi_dev->dev, "Unknown command %d\n",
                        le32_to_cpu(context->rx->cmd));
                break;
        }

        ret = mhi_queue_buf(context->mhi_dev, DMA_FROM_DEVICE, context->rx,
                            SAHARA_PACKET_MAX_SIZE, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to requeue rx buf %d\n", ret);
}

static void sahara_parse_dump_table(struct sahara_context *context)
{
        struct sahara_dump_table_entry *image_out_table;
        struct sahara_debug_table_entry64 *dev_table;
        struct sahara_memory_dump_meta_v1 *dump_meta;
        u64 table_nents;
        u64 dump_length;
        u64 mul_bytes;
        int ret;
        u64 i;

        table_nents = context->dump_table_length / sizeof(*dev_table);
        context->dump_images_left = table_nents;
        dump_length = 0;

        dev_table = (struct sahara_debug_table_entry64 *)(context->rx);
        for (i = 0; i < table_nents; ++i) {
                /* Do not trust the device, ensure the strings are terminated */
                dev_table[i].description[SAHARA_TABLE_ENTRY_STR_LEN - 1] = 0;
                dev_table[i].filename[SAHARA_TABLE_ENTRY_STR_LEN - 1] = 0;

                if (check_add_overflow(dump_length,
                                       le64_to_cpu(dev_table[i].length),
                                       &dump_length)) {
                        /* Discard the dump */
                        sahara_send_reset(context);
                        return;
                }

                dev_dbg(&context->mhi_dev->dev,
                        "Memory dump table entry %lld type: %lld address: %#llx length: %#llx description: \"%s\" filename \"%s\"\n",
                        i,
                        le64_to_cpu(dev_table[i].type),
                        le64_to_cpu(dev_table[i].address),
                        le64_to_cpu(dev_table[i].length),
                        dev_table[i].description,
                        dev_table[i].filename);
        }

        if (check_add_overflow(dump_length, (u64)sizeof(*dump_meta), &dump_length)) {
                /* Discard the dump */
                sahara_send_reset(context);
                return;
        }
        if (check_mul_overflow((u64)sizeof(*image_out_table), table_nents, &mul_bytes)) {
                /* Discard the dump */
                sahara_send_reset(context);
                return;
        }
        if (check_add_overflow(dump_length, mul_bytes, &dump_length)) {
                /* Discard the dump */
                sahara_send_reset(context);
                return;
        }

        context->mem_dump_sz = dump_length;
        context->mem_dump = vzalloc(dump_length);
        if (!context->mem_dump) {
                /* Discard the dump */
                sahara_send_reset(context);
                return;
        }

        /* Populate the dump metadata and table for userspace */
        dump_meta = context->mem_dump;
        dump_meta->magic = SAHARA_DUMP_V1_MAGIC;
        dump_meta->version = SAHARA_DUMP_V1_VER;
        dump_meta->dump_size = dump_length;
        dump_meta->table_size = context->dump_table_length;

        image_out_table = context->mem_dump + sizeof(*dump_meta);
        for (i = 0; i < table_nents; ++i) {
                image_out_table[i].type = le64_to_cpu(dev_table[i].type);
                image_out_table[i].address = le64_to_cpu(dev_table[i].address);
                image_out_table[i].length = le64_to_cpu(dev_table[i].length);
                strscpy(image_out_table[i].description, dev_table[i].description,
                        SAHARA_TABLE_ENTRY_STR_LEN);
                strscpy(image_out_table[i].filename,
                        dev_table[i].filename,
                        SAHARA_TABLE_ENTRY_STR_LEN);
        }

        context->mem_dump_freespace = &image_out_table[i];

        /* Done parsing the table, switch to image dump mode */
        context->dump_table_length = 0;

        /* Request the first chunk of the first image */
        context->dump_image = &image_out_table[0];
        dump_length = min_t(u64, context->dump_image->length, SAHARA_READ_MAX_SIZE);
        /* Avoid requesting EOI sized data so that we can identify errors */
        if (dump_length == SAHARA_END_OF_IMAGE_LENGTH)
                dump_length = SAHARA_END_OF_IMAGE_LENGTH / 2;

        context->dump_image_offset = dump_length;

        context->tx[0]->cmd = cpu_to_le32(SAHARA_MEM_READ64_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_MEM_READ64_LENGTH);
        context->tx[0]->memory_read64.memory_address = cpu_to_le64(context->dump_image->address);
        context->tx[0]->memory_read64.memory_length = cpu_to_le64(dump_length);

        context->rx_size_requested = dump_length;

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_MEM_READ64_LENGTH, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to send read for dump content %d\n", ret);
}

static void sahara_parse_dump_image(struct sahara_context *context)
{
        u64 dump_length;
        int ret;

        memcpy(context->mem_dump_freespace, context->rx, context->rx_size);
        context->mem_dump_freespace += context->rx_size;

        if (context->dump_image_offset >= context->dump_image->length) {
                /* Need to move to next image */
                context->dump_image++;
                context->dump_images_left--;
                context->dump_image_offset = 0;

                if (!context->dump_images_left) {
                        /* Dump done */
                        dev_coredumpv(context->mhi_dev->mhi_cntrl->cntrl_dev,
                                      context->mem_dump,
                                      context->mem_dump_sz,
                                      GFP_KERNEL);
                        context->mem_dump = NULL;
                        sahara_send_reset(context);
                        return;
                }
        }

        /* Get next image chunk */
        dump_length = context->dump_image->length - context->dump_image_offset;
        dump_length = min_t(u64, dump_length, SAHARA_READ_MAX_SIZE);
        /* Avoid requesting EOI sized data so that we can identify errors */
        if (dump_length == SAHARA_END_OF_IMAGE_LENGTH)
                dump_length = SAHARA_END_OF_IMAGE_LENGTH / 2;

        context->tx[0]->cmd = cpu_to_le32(SAHARA_MEM_READ64_CMD);
        context->tx[0]->length = cpu_to_le32(SAHARA_MEM_READ64_LENGTH);
        context->tx[0]->memory_read64.memory_address =
                cpu_to_le64(context->dump_image->address + context->dump_image_offset);
        context->tx[0]->memory_read64.memory_length = cpu_to_le64(dump_length);

        context->dump_image_offset += dump_length;
        context->rx_size_requested = dump_length;

        ret = mhi_queue_buf(context->mhi_dev, DMA_TO_DEVICE, context->tx[0],
                            SAHARA_MEM_READ64_LENGTH, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev,
                        "Unable to send read for dump content %d\n", ret);
}

static void sahara_dump_processing(struct work_struct *work)
{
        struct sahara_context *context = container_of(work, struct sahara_context, dump_work);
        int ret;

        /*
         * We should get the expected raw data, but if the device has an error
         * it is supposed to send EOI with an error code.
         */
        if (context->rx_size != context->rx_size_requested &&
            context->rx_size != SAHARA_END_OF_IMAGE_LENGTH) {
                dev_err(&context->mhi_dev->dev,
                        "Unexpected response to read_data. Expected size: %#zx got: %#zx\n",
                        context->rx_size_requested,
                        context->rx_size);
                goto error;
        }

        if (context->rx_size == SAHARA_END_OF_IMAGE_LENGTH &&
            le32_to_cpu(context->rx->cmd) == SAHARA_END_OF_IMAGE_CMD) {
                dev_err(&context->mhi_dev->dev,
                        "Unexpected EOI response to read_data. Status: %d\n",
                        le32_to_cpu(context->rx->end_of_image.status));
                goto error;
        }

        if (context->rx_size == SAHARA_END_OF_IMAGE_LENGTH &&
            le32_to_cpu(context->rx->cmd) != SAHARA_END_OF_IMAGE_CMD) {
                dev_err(&context->mhi_dev->dev,
                        "Invalid EOI response to read_data. CMD: %d\n",
                        le32_to_cpu(context->rx->cmd));
                goto error;
        }

        /*
         * Need to know if we received the dump table, or part of a dump image.
         * Since we get raw data, we cannot tell from the data itself. Instead,
         * we use the stored dump_table_length, which we zero after we read and
         * process the entire table.
         */
        if (context->dump_table_length)
                sahara_parse_dump_table(context);
        else
                sahara_parse_dump_image(context);

        ret = mhi_queue_buf(context->mhi_dev, DMA_FROM_DEVICE, context->rx,
                            SAHARA_PACKET_MAX_SIZE, MHI_EOT);
        if (ret)
                dev_err(&context->mhi_dev->dev, "Unable to requeue rx buf %d\n", ret);

        return;

error:
        vfree(context->mem_dump);
        context->mem_dump = NULL;
        sahara_send_reset(context);
}

static void sahara_read_data_processing(struct work_struct *work)
{
        struct sahara_context *context = container_of(work, struct sahara_context, read_data_work);

        read_data_helper(context, 0);
}

static int sahara_mhi_probe(struct mhi_device *mhi_dev, const struct mhi_device_id *id)
{
        struct sahara_context *context;
        int ret;
        int i;

        context = devm_kzalloc(&mhi_dev->dev, sizeof(*context), GFP_KERNEL);
        if (!context)
                return -ENOMEM;

        context->rx = devm_kzalloc(&mhi_dev->dev, SAHARA_PACKET_MAX_SIZE, GFP_KERNEL);
        if (!context->rx)
                return -ENOMEM;

        if (!strcmp(mhi_dev->mhi_cntrl->name, "AIC200")) {
                context->image_table = aic200_image_table;
                context->table_size = ARRAY_SIZE(aic200_image_table);
        } else {
                context->image_table = aic100_image_table;
                context->table_size = ARRAY_SIZE(aic100_image_table);
                context->non_streaming = true;
        }

        /*
         * There are two firmware implementations for READ_DATA handling.
         * The older "SBL" implementation defines a Sahara transfer size, and
         * expects that the response is a single transport transfer. If the
         * FW wants to transfer a file that is larger than the transfer size,
         * the FW will issue multiple READ_DATA commands. For this
         * implementation, we need to allocate enough buffers to contain the
         * entire Sahara transfer size.
         *
         * The newer "XBL" implementation does not define a maximum transfer
         * size and instead expects the data to be streamed over using the
         * transport level MTU. The FW will issue a single READ_DATA command
         * of whatever size, and consume multiple transport level transfers
         * until the expected amount of data is consumed. For this
         * implementation we only need a single buffer of the transport MTU
         * but we'll need to be able to use it multiple times for a single
         * READ_DATA request.
         *
         * AIC100 is the SBL implementation and defines SAHARA_TRANSFER_MAX_SIZE
         * and we need 9x SAHARA_PACKET_MAX_SIZE to cover that. We can use
         * MHI_CHAIN to link multiple buffers into a single transfer but the
         * remote side will not consume the buffers until it sees an EOT, thus
         * we need to allocate enough buffers to put in the tx fifo to cover an
         * entire READ_DATA request of the max size.
         *
         * AIC200 is the XBL implementation, and so a single buffer will work.
         */
        for (i = 0; i < SAHARA_NUM_TX_BUF; ++i) {
                context->tx[i] = devm_kzalloc(&mhi_dev->dev,
                                              SAHARA_PACKET_MAX_SIZE,
                                              GFP_KERNEL);
                if (!context->tx[i])
                        return -ENOMEM;
                if (is_streaming(context))
                        break;
        }

        context->mhi_dev = mhi_dev;
        INIT_WORK(&context->fw_work, sahara_processing);
        INIT_WORK(&context->dump_work, sahara_dump_processing);
        INIT_WORK(&context->read_data_work, sahara_read_data_processing);

        context->active_image_id = SAHARA_IMAGE_ID_NONE;
        dev_set_drvdata(&mhi_dev->dev, context);

        ret = mhi_prepare_for_transfer(mhi_dev);
        if (ret)
                return ret;

        ret = mhi_queue_buf(mhi_dev, DMA_FROM_DEVICE, context->rx, SAHARA_PACKET_MAX_SIZE, MHI_EOT);
        if (ret) {
                mhi_unprepare_from_transfer(mhi_dev);
                return ret;
        }

        return 0;
}

static void sahara_mhi_remove(struct mhi_device *mhi_dev)
{
        struct sahara_context *context = dev_get_drvdata(&mhi_dev->dev);

        cancel_work_sync(&context->fw_work);
        cancel_work_sync(&context->dump_work);
        vfree(context->mem_dump);
        sahara_release_image(context);
        mhi_unprepare_from_transfer(mhi_dev);
}

static void sahara_mhi_ul_xfer_cb(struct mhi_device *mhi_dev, struct mhi_result *mhi_result)
{
        struct sahara_context *context = dev_get_drvdata(&mhi_dev->dev);

        if (!mhi_result->transaction_status && context->read_data_length && is_streaming(context))
                schedule_work(&context->read_data_work);
}

static void sahara_mhi_dl_xfer_cb(struct mhi_device *mhi_dev, struct mhi_result *mhi_result)
{
        struct sahara_context *context = dev_get_drvdata(&mhi_dev->dev);

        if (!mhi_result->transaction_status) {
                context->rx_size = mhi_result->bytes_xferd;
                if (context->is_mem_dump_mode)
                        schedule_work(&context->dump_work);
                else
                        schedule_work(&context->fw_work);
        }

}

static const struct mhi_device_id sahara_mhi_match_table[] = {
        { .chan = "QAIC_SAHARA", },
        {},
};

static struct mhi_driver sahara_mhi_driver = {
        .id_table = sahara_mhi_match_table,
        .remove = sahara_mhi_remove,
        .probe = sahara_mhi_probe,
        .ul_xfer_cb = sahara_mhi_ul_xfer_cb,
        .dl_xfer_cb = sahara_mhi_dl_xfer_cb,
        .driver = {
                .name = "sahara",
        },
};

int sahara_register(void)
{
        return mhi_driver_register(&sahara_mhi_driver);
}

void sahara_unregister(void)
{
        mhi_driver_unregister(&sahara_mhi_driver);
}