/**
 * \file
 *
 * \brief AES block device
 *
 * This is an AES block device, which adds an encryption layer to another
 * underlying block device.
 *
 * - Compiler:           IAR EWAVR32 and GNU GCC for AVR32
 * - Supported devices:  All devices
 * - AppNote:
 *
 * \author               Atmel Corporation: http://www.atmel.com \n
 *                       Support and FAQ: http://support.atmel.no/
 *
 * \page License
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3. The name of Atmel may not be used to endorse or promote products derived
 * from this software without specific prior written permission.
 *
 * 4. This software may only be redistributed and used in connection with an
 * Atmel AVR product.
 *
 * THIS SOFTWARE IS PROVIDED BY ATMEL "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT ARE
 * EXPRESSLY AND SPECIFICALLY DISCLAIMED. IN NO EVENT SHALL ATMEL BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */

#include <buffer.h>
#include <debug.h>
#include <dmapool.h>
#include <interrupt.h>
#include <workqueue.h>
#include <timer_tc.h>
#include <board/init.h>
#include <chip/memory-map.h>
#include <chip/clk.h>
#include <block/device.h>
#include <status-codes.h>
#include <malloc.h>
#include <string.h>
#include <aes.h>
#include <atomic.h>
#include <block/aesblk.h>

static void aesblk_worker(void *data);

/**
 * \brief Initialize AES block device
 *
 * This function will initialize the AES block device with the underlying
 * block device (typically a SDMMC card).
 *
 * \param  abdev	Pointer to the AES block device
 * \param  card_bdev	Pointer to the underlying block device
 */
void aesblk_init(struct aes_block_device *abdev,
		struct block_device *card_bdev)
{
	struct block_device	*bdev;

	bdev = &abdev->bdev;
	abdev->card_bdev = card_bdev;

	bdev->prepare_req = aesblk_prepare_req;
	bdev->alloc_req = aesblk_alloc_req;
	bdev->free_req = aesblk_free_req;
}

/**
 * \brief Load block device descriptions from underlying device.
 *
 * \param  aes_bdev	AES block device pointer
 */
void aesblk_update(struct aes_block_device *abdev)
{
	struct block_device *bdev = &abdev->bdev;
	struct block_device *card_bdev = abdev->card_bdev;

	bdev->block_size = card_bdev->block_size;
	bdev->nr_blocks = card_bdev->nr_blocks;
	bdev->flags = card_bdev->flags;
	bdev->get_dev_id = card_bdev->get_dev_id;
}

static inline struct aes_block_device *bdev_to_abdev(
		struct block_device *bdev)
{
	return container_of(bdev, struct aes_block_device, bdev);
}

static inline struct aes_block_request *abreq_of(
		struct block_request *req)
{
	return container_of(req, struct aes_block_request, breq);
}

/* Callback called when data read to a buffer list has been decrypted. */
static void aesblk_buf_list_decrypt_done(struct aes_request *areq,
		struct slist *buf_list, void *context)
{
	struct aes_block_request	*abreq = context;
	struct block_request		*breq = &abreq->breq;
	unsigned int iflags;

	assert(abreq);

	iflags = cpu_irq_save();
	/* Move the buffers to the requests buffer list. */
	slist_move_to_tail(&breq->buf_list, buf_list);
	aes_free_request(areq);

	/* Schedule the buffer list complete callback to be called. */
	workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
	abreq->state |= AES_BLK_STATE_LIST_DECRYPTION_DONE;

	/* If all buffer lists are processed and a request done is pending,
	 * schedule it for completion now. */
	if (--abreq->buf_counter == 0 && (abreq->state
			& AES_BLK_STATE_DECRYPTION_DONE_PENDING)) {
		abreq->state &= ~AES_BLK_STATE_DECRYPTION_DONE_PENDING;
		abreq->state |= AES_BLK_STATE_REQ_DECRYPTION_DONE;
	}
	cpu_irq_restore(iflags);
}

/* Callback called when data has been read from the storage medium to a
 * buffer list. */
static void aesblk_read_buf_list_done(struct block_device *bdev,
		struct block_request *breq, struct slist *buf_list)
{
	struct aes_block_request	*abreq = breq->context;
	struct aes_block_device		*abdev = abreq->abdev;
	struct aes_module		*module = &abdev->module;
	struct aes_request		*areq;

	/* If the AES module is not in decrypt mode, something is very
	 * wrong. */
	assert(aes_get_cipher(module) == AES_DECRYPT);

	/* In order to keep a request from completing while some buffers are
	 * still being processed, we need to count all buffers submitted to
	 * the AES module for decryption. */
	abreq->buf_counter++;

	/* Submit the buffers to the AES module for decryption. */
	areq = aes_alloc_request(module);
	aes_prepare_request(areq, buf_list);
	areq->req_done = aesblk_buf_list_decrypt_done;
	areq->context = abreq;
	aes_submit_request(areq);
}

/* Function called when a buffer list is submitted to an active
 * read request. */
static int aesblk_submit_read_buf_list(struct block_device *bdev,
		struct block_request *breq, struct slist *buf_list)
{
	struct aes_block_request	*abreq = abreq_of(breq);
	struct aes_block_device		*abdev = bdev_to_abdev(bdev);
	struct block_request		*card_breq = abreq->card_breq;
	unsigned int iflags;

	assert(bdev);
	assert(bdev == breq->bdev);
	assert(breq);
	assert(buf_list);
	assert(!slist_is_empty(buf_list));

	iflags = cpu_irq_save();
	/* Move the buffers to the block requests buffer list. */
	slist_move_to_tail(&card_breq->buf_list, buf_list);
	cpu_irq_restore(iflags);

	return block_submit_buf_list(abdev->card_bdev, card_breq,
			&card_breq->buf_list);
}

/* Callback called when the data read by a read request has been
 * decrypted. */
static void aesblk_req_decrypt_done(struct aes_request *areq,
		struct slist *buf_list, void *context)
{
	struct aes_block_request	*abreq = context;
	struct block_request		*breq = &abreq->breq;
	unsigned int iflags;

	assert(areq);

	iflags = cpu_irq_save();
	/* Move the decrypted buffers to the block requests buffer list. */
	if (!slist_is_empty(buf_list))
		slist_move_to_tail(&breq->buf_list, buf_list);
	aes_free_request(areq);

	if (abreq->buf_counter > 0) {
		/* If we still have buffer lists being processed, the
		 * request can not be completed yet. */
		abreq->state |= AES_BLK_STATE_DECRYPTION_DONE_PENDING;
	} else {
		/* All buffer lists have been processed, schedule the request
		 * done callback to be called. */
		workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
		abreq->state |= AES_BLK_STATE_REQ_DECRYPTION_DONE;
	}
	cpu_irq_restore(iflags);
}

/* Callback called when data has been read from the storage medium by a
 * read request */
static void aesblk_req_data_read(struct block_device *bdev,
		struct block_request *card_breq)
{
	struct aes_block_request	*abreq = card_breq->context;
	struct aes_module		*module = &abreq->abdev->module;
	struct aes_request		*areq;

	assert(abreq);

	/* Submit the read buffers to the AES module for decryption. */
	areq = aes_alloc_request(module);
	aes_prepare_request(areq, &card_breq->buf_list);
	areq->req_done = aesblk_req_decrypt_done;
	areq->context = abreq;
	aes_submit_request(areq);
}

/* Function called when a read request is submitted */
static void aesblk_submit_read_req(struct block_device *bdev,
		struct block_request *breq)
{
	struct aes_block_request	*abreq = abreq_of(breq);
	struct aes_block_device		*abdev = bdev_to_abdev(bdev);
	struct aes_module		*module = &abreq->abdev->module;
	struct block_request		*card_breq = abreq->card_breq;
	unsigned int iflags;

	iflags = cpu_irq_save();
	/* Move the buffers to the block request used to interface the
	 * underlying storage medium. */
	if (!slist_is_empty(&breq->buf_list))
		slist_move_to_tail(&card_breq->buf_list, &breq->buf_list);
	cpu_irq_restore(iflags);

	/* Set the AES module to decrypt the stored data. */
	aes_set_cipher(module, AES_DECRYPT);
	/* Start reading from the storage medium. */
	block_submit_req(abdev->card_bdev, card_breq);
}

/* Callback called when the buffer lists submitted to an ongoing write request
 * has been written to the storage medium. */
static void aesblk_write_buf_list_done(struct block_device *bdev,
		struct block_request *card_breq, struct slist *buf_list)
{
	struct aes_block_request	*abreq = card_breq->context;
	struct block_request		*breq = &abreq->breq;
	unsigned int iflags;

	assert(!slist_is_empty(buf_list));
	/* We are not able to handle errors in the underlying storage medium
	 * yet... */
	assert(card_breq->status == -STATUS_IN_PROGRESS);

	iflags = cpu_irq_save();
	slist_move_to_tail(&breq->buf_list, buf_list);
	workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
	abreq->state |= AES_BLK_STATE_WRITE_LIST_DONE;
	cpu_irq_restore(iflags);
}

/* Callback called when the buffer lists submitted to an ongoing write request
 * has been encrypted. */
static void aesblk_buf_list_encrypt_done(struct aes_request *areq,
		struct slist *buf_list, void *context)
{
	struct aes_block_request	*abreq = context;
	unsigned int iflags;

	assert(abreq);
	assert(buf_list);
	assert(!slist_is_empty(buf_list));
	assert(abreq->card_breq->status == -STATUS_IN_PROGRESS);

	iflags = cpu_irq_save();
	/* Move the encrypted buffers to the requests buffer list. They will
	 * later be written to the storage medium. */
	slist_move_to_tail(&abreq->buf_list, buf_list);
	aes_free_request(areq);

	/* Schedule the encrytped buffers to be written to the storage
	 * medium. */
	workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
	abreq->state |= AES_BLK_STATE_LIST_ENCRYPTION_DONE;
	assert(abreq->buf_counter);
	cpu_irq_restore(iflags);
}

/* Function called when a buffer list is submitted to an ongoing write
 * request. */
static int aesblk_submit_write_buf_list(struct block_device *bdev,
		struct block_request *breq, struct slist *buf_list)
{
	struct aes_block_request	*abreq = abreq_of(breq);
	struct aes_block_device		*abdev = abreq->abdev;
	struct aes_module		*module = &abdev->module;
	struct aes_request		*areq;

	assert(bdev);
	assert(bdev == breq->bdev);
	assert(abreq);
	assert(buf_list);
	assert(!slist_is_empty(buf_list));
	/* If the AES module is not in encryption mode, something is very
	 * wrong. */
	assert(aes_get_cipher(module) == AES_ENCRYPT);

	/* Submit the buffer list for encryption. */
	areq = aes_alloc_request(&abdev->module);
	aes_prepare_request(areq, buf_list);
	areq->req_done = aesblk_buf_list_encrypt_done;
	areq->context = abreq;
	aes_submit_request(areq);

	return 0;
}

/* Callback called when the data submitted with a write request has been
 * written to the storage medium. */
static void aesblk_req_data_written(struct block_device *bdev,
		struct block_request *card_breq)
{
	struct aes_block_request	*abreq = card_breq->context;
	struct block_request		*breq = &abreq->breq;
	unsigned int iflags;

	assert(abreq);
	/* We are not able to handle errors in the underlying storage medium
	 * yet... */
	assert(!card_breq->status);

	iflags = cpu_irq_save();
	/* Move the written buffers to the block requests buffer list and
	 * schedule its callback to be called. */
	if (!slist_is_empty(&card_breq->buf_list))
		slist_move_to_tail(&breq->buf_list, &card_breq->buf_list);
	workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
	abreq->state |= AES_BLK_STATE_WRITE_REQ_DONE;
	cpu_irq_restore(iflags);
}

/* Callback called when the data submitted with a write request has been
 * encrypted */
static void aesblk_req_encrypt_done(struct aes_request *areq,
		struct slist *buf_list, void *context)
{
	struct aes_block_request *abreq = context;
	unsigned int iflags;

	assert(abreq);

	iflags = cpu_irq_save();
	/* Move buffers to the requests buffer list, these will later be
	 * written to the storage medium. */
	if (!slist_is_empty(buf_list))
		slist_move_to_tail(&abreq->buf_list, buf_list);
	aes_free_request(areq);

	/* Schedule the actual write operation. */
	workqueue_add_item_safe(&main_workqueue, &abreq->worker_item);
	abreq->state |= AES_BLK_STATE_REQ_ENCRYPTION_DONE;
	abreq->buf_counter++;
	cpu_irq_restore(iflags);
}

/* Function called when a write request is submitted */
static void aesblk_submit_write_req(struct block_device *bdev,
		struct block_request *breq)
{
	struct aes_block_request	*abreq = abreq_of(breq);
	struct aes_block_device		*abdev = bdev_to_abdev(bdev);
	struct aes_module		*module = &abdev->module;
	struct aes_request		*aes_req;

	/* Prepare the AES module and submit the requests buffer list
	 * for encryption. */
	aes_set_cipher(module, AES_ENCRYPT);

	aes_req = aes_alloc_request(module);
	aes_prepare_request(aes_req, &abreq->breq.buf_list);
	aes_req->req_done = aesblk_req_encrypt_done;
	aes_req->context = abreq;
	aes_submit_request(aes_req);
}

/* Worker function called from the workqueue when there is work pending
 * for the AES block device */
static void aesblk_worker(void *data)
{
	struct aes_block_request *abreq = data;
	struct aes_block_device *abdev = abreq->abdev;
	struct block_request *breq = &abreq->breq;
	struct block_request *card_breq = abreq->card_breq;
	struct block_device *card_bdev = abdev->card_bdev;

	cpu_irq_disable();
	if (abreq->state & AES_BLK_STATE_REQ_ENCRYPTION_DONE) {
		abreq->state &= ~AES_BLK_STATE_REQ_ENCRYPTION_DONE;
		cpu_irq_enable();
		block_submit_req(card_bdev, card_breq);
		cpu_irq_disable();
	}
	if (abreq->state & AES_BLK_STATE_LIST_ENCRYPTION_DONE) {
		abreq->state &= ~AES_BLK_STATE_LIST_ENCRYPTION_DONE;
		cpu_irq_enable();
		assert(abreq->card_breq->status == -STATUS_IN_PROGRESS);
		block_submit_buf_list(card_bdev, card_breq, &abreq->buf_list);
		cpu_irq_disable();
	}
	if (abreq->state & AES_BLK_STATE_WRITE_LIST_DONE) {
		abreq->state &= ~AES_BLK_STATE_WRITE_LIST_DONE;
		if (breq->buf_list_done) {
			cpu_irq_enable();
			breq->buf_list_done(breq->bdev, breq,
					&breq->buf_list);
			cpu_irq_disable();
		}
	}
	if (abreq->state & AES_BLK_STATE_WRITE_REQ_DONE) {
		breq->status = card_breq->status;
		breq->bytes_xfered = card_breq->bytes_xfered;
		cpu_irq_enable();
		if (breq->req_done)
			breq->req_done(&abdev->bdev, &abreq->breq);
		cpu_irq_disable();
	}

	if (abreq->state & AES_BLK_STATE_LIST_DECRYPTION_DONE) {
		abreq->state &= ~AES_BLK_STATE_LIST_DECRYPTION_DONE;
		cpu_irq_enable();
		if (breq->buf_list_done)
			breq->buf_list_done(breq->bdev, breq,
					&breq->buf_list);
		cpu_irq_disable();
	}
	if (abreq->state & AES_BLK_STATE_REQ_DECRYPTION_DONE) {
		abreq->state &= ~AES_BLK_STATE_REQ_DECRYPTION_DONE;
		cpu_irq_enable();
		breq->status = card_breq->status;
		breq->bytes_xfered = card_breq->bytes_xfered;
		if (breq->req_done)
			breq->req_done(&abdev->bdev, breq);
		cpu_irq_disable();
	}
	cpu_irq_enable();
}

/**
 * \brief Prepare an AES block request for use.
 *
 * This function will prepare an AES block request for submittion to the
 * AES block driver.
 *
 * \param  bdev		Pointer to the block device part of an AES block device.
 * \param  breq		Pointer to the block request to prepare.
 * \param  lba		Location of the blocks to operate on.
 * \param  nr_blocks	Number of blocks to read/write.
 * \param  operation	Which block operation to perform (read or write).
 */
void aesblk_prepare_req(struct block_device *bdev,
		struct block_request *breq,
		uint32_t lba, uint32_t nr_blocks,
		enum block_operation operation)
{
	struct aes_block_device		*abdev = bdev_to_abdev(bdev);
	struct aes_block_request	*abreq = abreq_of(breq);
	struct block_device		*card_bdev = abdev->card_bdev;
	struct block_request		*card_breq = abreq->card_breq;

	dbg_verbose("aes prep req %p: LBA %08x blocks %u op %d\n",
			abreq, lba, nr_blocks, operation);

	slist_init(&abreq->breq.buf_list);
	breq->status = -STATUS_IN_PROGRESS;
	breq->bytes_xfered = 0;

	breq->bdev = bdev;
	abreq->abdev = abdev;
	abreq->state = 0;
	abreq->buf_counter = 0;
	slist_init(&abreq->buf_list);

	workqueue_init_item(&abreq->worker_item, aesblk_worker, abreq);

	switch (operation) {
	case BLK_OP_READ:
		block_prepare_req(card_bdev, card_breq,
				lba, nr_blocks, operation);
		breq->req_submit = aesblk_submit_read_req;
		breq->req_submit_buf_list = aesblk_submit_read_buf_list;
		card_breq->req_done = aesblk_req_data_read;
		card_breq->buf_list_done = aesblk_read_buf_list_done;
		break;

	case BLK_OP_WRITE:
		block_prepare_req(card_bdev, card_breq,
				lba, nr_blocks, operation);
		breq->req_submit = aesblk_submit_write_req;
		breq->req_submit_buf_list = aesblk_submit_write_buf_list;
		card_breq->req_done = aesblk_req_data_written;
		card_breq->buf_list_done = aesblk_write_buf_list_done;
		break;

	default:
		unhandled_case(operation);
	}

	card_breq->context = abreq;
}

/**
 * \brief Allocate memory for one AES block request.
 *
 * This function will allocate all the memory needed by one AES block
 * request and return a pointer to this request.
 *
 * \param  bdev		Pointer to the block device interface of an AES
 *			block device.
 * \return		Pointer to the allocated block request.
 */
struct block_request *aesblk_alloc_req(struct block_device *bdev)
{
	struct aes_block_request	*abreq;
	struct aes_block_device		*abdev = bdev_to_abdev(bdev);

	abreq = malloc(sizeof(struct aes_block_request));
	if (unlikely(!abreq))
		goto err_aes;

	memset(abreq, 0, sizeof(struct aes_block_request));
	abreq->card_breq = block_alloc_request(abdev->card_bdev);
	if (unlikely(!abreq->card_breq))
		goto err_card;

	return &abreq->breq;

err_card:
	free(abreq);
err_aes:
	return NULL;
}

/**
 * \brief Free the memory allocated by a block request.
 *
 * \param  bdev		Pointer to the block device the block request was
 *			allocated for.
 * \param  breq		Pointer to the block request that is to be freed.
 */
void aesblk_free_req(struct block_device *bdev,
		struct block_request *breq)
{
	struct aes_block_request	*abreq = abreq_of(breq);

	block_free_request(abreq->abdev->card_bdev, abreq->card_breq);
	free(abreq);
}
