/*
 * Direct MTD block device access, readonly with NAND bad block skipping
 *
 * (C) 2007 Andre McCurdy, NXP Semiconductors
 * (C) 2000-2003 Nicolas Pitre <nico@cam.org>
 * (C) 1999-2003 David Woodhouse <dwmw2@infradead.org>
 */

#include <linux/fs.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <linux/types.h>
#include <linux/vmalloc.h>

#include <linux/mtd/mtd.h>
#include <linux/mtd/blktrans.h>
#include <linux/mutex.h>

#if defined (CONFIG_PNX8335_AES)
#include <linux/pnx8335_aes.h>
#define AIMAGE_DECRYPTION
#endif

#define DEV_BLOCKSIZE		(512)			/* fixed by MTD block driver infrastructure */
#define CACHE_SIZE		(4 * 1024)		/* must be a multiple of DEV_BLOCKSIZE and a fraction of NAND erase block size */
#define CACHE_BASE_INVALID	((unsigned long) -1)

/*
   Pick a major device number to 'borrow'... whichever one we pick, the
   driver which really owns that number MUST NOT ALSO BE INCLUDE IN THE KERNEL !!

   31 : The offical mtdblock major number. Picking this allows the driver to be
	tested as a drop in (readonly) replacement for the normal mtdblock driver.

   44 : The mtd 'ftl.c' major number. Note that the bootloader, fstab, etc,
	etc may need tweaking to use an unexpected device...
*/

#if 0
#define ROBBS_NAME	"mtdblock"
#define ROBBS_MAJOR	31
#else
#define ROBBS_NAME	"mtdblock_robbs"
#define ROBBS_MAJOR	44
#endif

#if defined (AIMAGE_DECRYPTION)
#define IMG_MAGIC                       0x41676d69
#define IMG_MAGIC2                      0x43634d52
#define IMG_TYPE_INITFS_CRAMFS          (0x0A)
#define IMG_FLAG_ENCMODE_NONE           (0 << 6)
#define IMG_FLAG_ENCMODE_CBC_ONEPASS    (1 << 6)
#define IMG_FLAG_ENCMODE_CBC_4KBLOCKS   (2 << 6)
#define IMG_FLAG_ENCMODE_CTR            (3 << 6)
#define IMG_FLAG_ENCMODE_MASK           (7 << 6)

typedef struct
{
	unsigned int _d0[2];
	unsigned int magic;
	unsigned int magic2;
	unsigned int _d1[2];
	unsigned int type;
	unsigned int _d2[5];
	unsigned int flags;
	unsigned int _d3[47];
	unsigned int iv[4];
}
aimage_v1_header_t;

static const unsigned char stb225_default_aes_key[16] =
{
    0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F
};
#endif

static struct mtdblk_dev {
	int count;
	struct mtd_info *mtd;
	unsigned int bblist_count;
	unsigned long *bblist_data;
	unsigned long cache_base;
	struct mutex cache_mutex;
#if defined (AIMAGE_DECRYPTION)
	int encrypted;
	unsigned char iv[16];
#endif
	unsigned char cache_buffer[CACHE_SIZE];
}
*mtdblks[MAX_MTD_DEVICES];

static int mtdblock_robbs_cache_load (struct mtdblk_dev *mtdblk, unsigned long pos_phys, unsigned long pos_virt)
{
	struct mtd_info *mtd = mtdblk->mtd;
	size_t retlen;
	int ret;

//	printk ("%s: %s at 0x%08lx\n", __FUNCTION__, mtd->name, pos);

	if (pos_phys & (CACHE_SIZE - 1)) {
		printk ("%s: pos_phys (0x%08lx) not CACHE_SIZE aligned ?!?\n", __FUNCTION__, pos_phys);
		return -EIO;
	}

	mutex_lock (&mtdblk->cache_mutex);

	/*
	   Fixme: is mtd->read() error handling correct ??
	*/
	ret = mtd->read (mtd, pos_phys, CACHE_SIZE, &retlen, mtdblk->cache_buffer);
	if (ret == 0)
		mtdblk->cache_base = (retlen == CACHE_SIZE) ? pos_phys : CACHE_BASE_INVALID;

#if defined (AIMAGE_DECRYPTION)
	if ((mtdblk->cache_base != CACHE_BASE_INVALID) && (mtdblk->encrypted)) {
		struct aes_op op;
		unsigned int enc_start_block;

		if (pos_virt == 0) {
			enc_start_block = (pos_virt + 256) / 16;
			op.src = mtdblk->cache_buffer + 256;
			op.blocks = (sizeof(mtdblk->cache_buffer) - 256) / 16;
		}
		else {
			enc_start_block = pos_virt / 16;
			op.src = mtdblk->cache_buffer;
			op.blocks = sizeof(mtdblk->cache_buffer) / 16;
		}
		op.dst = op.src;
		op.cmd = AESCMD_DECRYPT_CBC_KEY_USER;
		memcpy (op.key, stb225_default_aes_key, 16);
		memcpy (op.iv, mtdblk->iv, 12);
		op.iv[12] = (enc_start_block >> 24) & 0xFF;
		op.iv[13] = (enc_start_block >> 16) & 0xFF;
		op.iv[14] = (enc_start_block >> 8) & 0xFF;
		op.iv[15] = (enc_start_block >> 0) & 0xFF;
		if (aes_process_op_in_kernel (&op) != 0)
			printk ("aes operation failed !!\n");
	}
#endif

	mutex_unlock (&mtdblk->cache_mutex);

	if (ret)
		return ret;
	if (retlen != CACHE_SIZE)
		return -EIO;

	return 0;
}

static int mtdblock_robbs_readsect (struct mtd_blktrans_dev *dev, unsigned long lba, char *buf)
{
	struct mtdblk_dev *mtdblk = mtdblks[dev->devnum];
	struct mtd_info *mtd = mtdblk->mtd;
	unsigned long pos_abs = lba * DEV_BLOCKSIZE;
	unsigned long pos_base = (pos_abs / CACHE_SIZE) * CACHE_SIZE;
	unsigned long pos_offset = pos_abs - pos_base;
	unsigned long pos_base_original = pos_base;
	int result;
	int i;

//	printk ("%s: %s at 0x%08lx (0x%08lx + %4ld)\n", __FUNCTION__, mtd->name, pos_abs, pos_base, pos_offset);

	if (mtdblk->bblist_count != 0) {
		for (i = 0; i < mtdblk->bblist_count; i++) {
			if (pos_base < mtdblk->bblist_data[i])
				break;
			pos_base += CACHE_SIZE;
			if (pos_base >= mtd->size) {
				printk ("%s: %s 0x%08lx slipped out of range (0x%08x)\n",
					__FUNCTION__, mtd->name, pos_abs, mtd->size);
				memset (buf, 0xFF, DEV_BLOCKSIZE);
				return 0;
			}
		}
#if 0
		if ((pos_base != pos_base_original) && (pos_offset == 0))
			printk ("%s: %s (0x%08lx + %4ld) slipped to (0x%08lx + %4ld)\n",
				__FUNCTION__, mtd->name, pos_base_original, pos_offset,
				pos_base, pos_offset);
#endif
	}

	if (pos_base != mtdblk->cache_base)
		if ((result = mtdblock_robbs_cache_load (mtdblk, pos_base, pos_base_original)) != 0)
			return result;

	memcpy (buf, mtdblk->cache_buffer + pos_offset, DEV_BLOCKSIZE);
	return 0;
}

static int mtdblock_robbs_open (struct mtd_blktrans_dev *mbd)
{
	int dev = mbd->devnum;
	struct mtdblk_dev *mtdblk;
	struct mtd_info *mtd = mbd->mtd;
	unsigned long pos;
	unsigned int count;

	/*
	   Fixme: does mtdblock core code provide locking around calls to this function ?!?
	   If not, things look broken for SMP or preempt...
	*/

	if (mtdblks[dev]) {
		mtdblks[dev]->count++;
		printk ("mtdblock_open: %s, count: %d\n", mtd->name, mtdblks[dev]->count);
		return 0;
	}

	printk ("%s: %s, erasesize: %d, sizeof(struct mtdblk_dev): %d\n",
		__FUNCTION__, mtd->name, mtd->erasesize, sizeof(struct mtdblk_dev));

	mtdblk = kmalloc (sizeof(struct mtdblk_dev), GFP_KERNEL);
	if (! mtdblk)
		return -ENOMEM;

	memset (mtdblk, 0, sizeof(struct mtdblk_dev));
	mtdblk->count = 1;
	mtdblk->mtd = mtd;

	mtdblk->bblist_count = 0;
	for (pos = 0; pos < mtd->size; pos += CACHE_SIZE) {
		if (mtd->block_isbad (mtd, pos)) {
			mtdblk->bblist_count++;
			printk ("%s: badpos: 0x%08lx (count %d)\n",
				__FUNCTION__, pos, mtdblk->bblist_count);
		}
	}

	if (mtdblk->bblist_count != 0) {
		mtdblk->bblist_data = kmalloc (mtdblk->bblist_count * sizeof (*mtdblk->bblist_data), GFP_KERNEL);
		if (! mtdblk->bblist_data) {
			kfree (mtdblk);
			return -ENOMEM;
		}
		count = 0;
		for (pos = 0; pos < mtd->size; pos += CACHE_SIZE) {
			if (mtd->block_isbad (mtd, pos)) {
				count++;
//				printk ("%s: badpos: 0x%08lx (count %d)\n", __FUNCTION__, pos, count);
				if (count > mtdblk->bblist_count) {
					printk ("%s: increase in bad block count since first pass !?!\n", __FUNCTION__);
					break;
				}
				mtdblk->bblist_data[count - 1] = pos;
			}
		}
	}

	mtdblk->cache_base = CACHE_BASE_INVALID;
	mutex_init (&mtdblk->cache_mutex);
	mtdblks[dev] = mtdblk;

#if defined (AIMAGE_DECRYPTION)
	{
		int ret;
		unsigned char sector_buf[DEV_BLOCKSIZE];
		aimage_v1_header_t *header = (aimage_v1_header_t *) sector_buf;

		if ((ret = mtdblock_robbs_readsect (mbd, 0, sector_buf)) != 0)
			return ret;

		if ((le32_to_cpu (header->magic) == IMG_MAGIC) &&
		    (le32_to_cpu (header->magic2) == IMG_MAGIC2) &&
		    (le32_to_cpu (header->type) == IMG_TYPE_INITFS_CRAMFS) &&
		    ((le32_to_cpu (header->flags) & IMG_FLAG_ENCMODE_MASK) == IMG_FLAG_ENCMODE_CBC_4KBLOCKS))
		{
			mtdblk->encrypted = 1;
			memcpy (mtdblk->iv, header->iv, 16);
			mtdblk->cache_base = CACHE_BASE_INVALID;
			printk ("%s: Encrypted FS detected, iv[16] : %02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x\n", __FUNCTION__,
				mtdblk->iv[0], mtdblk->iv[1], mtdblk->iv[ 2], mtdblk->iv[ 3], mtdblk->iv[ 4], mtdblk->iv[ 5], mtdblk->iv[ 6], mtdblk->iv[ 7],
				mtdblk->iv[8], mtdblk->iv[9], mtdblk->iv[10], mtdblk->iv[11], mtdblk->iv[12], mtdblk->iv[13], mtdblk->iv[14], mtdblk->iv[15]);
		}
	}
#endif

	return 0;
}

static int mtdblock_robbs_release (struct mtd_blktrans_dev *mbd)
{
	int dev = mbd->devnum;
	struct mtdblk_dev *mtdblk = mtdblks[dev];
	struct mtd_info *mtd = mtdblk->mtd;

	printk ("%s: %s\n", __FUNCTION__, mtd->name);

	if (--mtdblk->count == 0) {
		mtdblks[dev] = NULL;
		kfree (mtdblk->bblist_data);
		kfree (mtdblk);
	}

	printk ("%s: %s OK\n", __FUNCTION__, mtd->name);

	return 0;
}

static void mtdblock_robbs_add_mtd (struct mtd_blktrans_ops *tr, struct mtd_info *mtd)
{
	struct mtd_blktrans_dev *dev = kmalloc (sizeof(*dev), GFP_KERNEL);

	if (! dev)
		return;

	memset (dev, 0, sizeof(*dev));

	dev->mtd = mtd;
	dev->devnum = mtd->index;
	dev->size = mtd->size / DEV_BLOCKSIZE;
	dev->tr = tr;
	dev->readonly = 1;

	add_mtd_blktrans_dev (dev);
}

static void mtdblock_robbs_remove_dev (struct mtd_blktrans_dev *dev)
{
	del_mtd_blktrans_dev (dev);
	kfree (dev);
}

static struct mtd_blktrans_ops mtdblock_robbs_tr =
{
	.name		= ROBBS_NAME,
	.major		= ROBBS_MAJOR,
	.part_bits	= 0,
	.blksize 	= 512,
	.open		= mtdblock_robbs_open,
//	.flush		= mtdblock_robbs_flush,
	.release	= mtdblock_robbs_release,
	.readsect	= mtdblock_robbs_readsect,
	.add_mtd	= mtdblock_robbs_add_mtd,
	.remove_dev	= mtdblock_robbs_remove_dev,
	.owner		= THIS_MODULE,
};

static int __init init_mtdblock_robbs (void)
{
	return register_mtd_blktrans (&mtdblock_robbs_tr);
}

static void __exit cleanup_mtdblock_robbs (void)
{
	deregister_mtd_blktrans (&mtdblock_robbs_tr);
}

module_init(init_mtdblock_robbs);
module_exit(cleanup_mtdblock_robbs);

MODULE_LICENSE("GPL");
