/*
   * Copyright (C) 2015-2016 Freescale Semiconductor, Inc.
   * Copyright 2018 NXP
   *
   * SPDX-License-Identifier:	GPL-2.0+
   */
  
  #include <asm/io.h>
  #include <asm/arch/mx6_bee.h>
  #include <linux/errno.h>
  #include <asm/system.h>
  #include <common.h>
  #include <command.h>
  #include <fuse.h>
  #include <asm/arch/sys_proto.h>

  #include <cpu_func.h>

  
  DECLARE_GLOBAL_DATA_PTR;
  
  #if (defined(CONFIG_SYS_DCACHE_OFF) || defined(CONFIG_SYS_ICACHE_OFF))
  #error "Bee needs Cache Open"
  #endif
  
  struct bee_parameters {
  	int key_method;
  	int mode;
  	u32 start1;
  	u32 size1;
  	u32 start2;
  	u32 size2;
  };
  
  #define SOFT_KEY 0
  #define SNVS_KEY 1
  
  #define ECB_MODE 0
  #define CTR_MODE 1
  
  #define AES_REGION0_ADDR	0x10000000
  #define AES_REGION1_ADDR	0x30000000
  
  static struct bee_parameters para;
  static int bee_inited;
  
  union key_soft {
  	u8 s_key[16];
  	u32 b_key[4];
  };
  
  union key_soft key_bad;
  
  /* software version */
  u8 hw_get_random_byte(void)
  {
  	static u32 lcg_state;
  	static u32 nb_soft = 9876543;
  #define MAX_SOFT_RNG 1024
  	static const u32 a = 1664525;
  	static const u32 c = 1013904223;
  	nb_soft = (nb_soft + 1) % MAX_SOFT_RNG;
  	lcg_state = (a * lcg_state + c);
  	return (u8) (lcg_state >> 24);
  }
  
  /*
   * Lock bee GPR0 bits
   * Only reset can release these bits.
   */
  static int bee_lock(void)
  {
  	int val;
  
  	val = readl(BEE_BASE_ADDR + GPR0);
  	val |= (GPR0_CTRL_CLK_EN_LOCK | GPR0_CTRL_SFTRST_N_LOCK |
  		GPR0_CTRL_AES_MODE_LOCK | GPR0_SEC_LEVEL_LOCK |
  		GPR0_AES_KEY_SEL_LOCK | GPR0_BEE_ENABLE_LOCK);
  	writel(val, BEE_BASE_ADDR + GPR0);
  
  	return 0;
  }
  
  /* Only check bee enable lock is enough */
  static int bee_locked(void)
  {
  	int val;
  
  	val = readl(BEE_BASE_ADDR + GPR0);
  
  	return val & GPR0_BEE_ENABLE_LOCK ? 1 : 0;
  }
  
  int bee_init(struct bee_parameters *p)
  {
  	int i;
  	union key_soft *key = &key_bad;
  	u32 value;
  
  	if (bee_locked()) {
  		printf("BEE already enabled and locked.
  ");
  		return CMD_RET_FAILURE;
  	}
  
  	/* CLKGATE, SFTRST */
  	writel(GPR0_CTRL_CLK_EN | GPR0_CTRL_SFTRST_N, BEE_BASE_ADDR + GPR0);
  	/* OFFSET_ADDR0 */
  	writel(p->start1 >> 16, BEE_BASE_ADDR + GPR1);
  	/*
  	 * OFFSET_ADDR1
  	 * Default protect IRAM region, if what you want to protect
  	 * bigger that 512M which is the max size that one AES region
  	 * can protect, we need AES region 1 to cover.
  	 */
  	writel(p->start2 >> 16, BEE_BASE_ADDR + GPR2);
  
  	if (p->key_method == SOFT_KEY) {
  		for (i = 0; i < 16; i++)
  			key->s_key[i] = hw_get_random_byte();
  		/* AES 128 key from software */
  		/* aes0_key0_w0 */
  		writel(key->b_key[0], BEE_BASE_ADDR + GPR3);
  		/* aes0_key0_w1 */
  		writel(key->b_key[1], BEE_BASE_ADDR + GPR4);
  		/* aes0_key0_w2 */
  		writel(key->b_key[2], BEE_BASE_ADDR + GPR5);
  		/* aes0_key0_w3 */
  		writel(key->b_key[3], BEE_BASE_ADDR + GPR6);
  	}
  
  	if (p->mode == ECB_MODE) {
  		value = GPR0_CTRL_CLK_EN | GPR0_CTRL_SFTRST_N |
  			GPR0_SEC_LEVEL_3 | GPR0_AES_KEY_SEL_SNVS |
  			GPR0_BEE_ENABLE | GPR0_CTRL_AES_MODE_ECB;
  		if (p->key_method == SOFT_KEY)
  			value = GPR0_CTRL_CLK_EN | GPR0_CTRL_SFTRST_N |
  				GPR0_SEC_LEVEL_3 | GPR0_AES_KEY_SEL_SOFT |
  				GPR0_BEE_ENABLE | GPR0_CTRL_AES_MODE_ECB;
  		writel(value, BEE_BASE_ADDR + GPR0);
  	} else {
  		for (i = 0; i < 16; i++)
  			key->s_key[i] = hw_get_random_byte();
  		/* aes_key1_w0 */
  		writel(key->b_key[0], BEE_BASE_ADDR + GPR8);
  		/* aes_key1_w1 */
  		writel(key->b_key[1], BEE_BASE_ADDR + GPR9);
  		/* aes_key1_w2 */
  		writel(key->b_key[2], BEE_BASE_ADDR + GPR10);
  		/* aes_key1_w3 */
  		writel(key->b_key[3], BEE_BASE_ADDR + GPR11);
  
  		value = GPR0_CTRL_CLK_EN | GPR0_CTRL_SFTRST_N |
  			GPR0_SEC_LEVEL_3 | GPR0_AES_KEY_SEL_SNVS |
  			GPR0_BEE_ENABLE | GPR0_CTRL_AES_MODE_CTR;
  		if (p->key_method == SOFT_KEY)
  			value = GPR0_CTRL_CLK_EN | GPR0_CTRL_SFTRST_N |
  				GPR0_SEC_LEVEL_3 | GPR0_AES_KEY_SEL_SOFT |
  				GPR0_BEE_ENABLE | GPR0_CTRL_AES_MODE_CTR;
  		writel(value, BEE_BASE_ADDR + GPR0);
  	}
  
  	bee_lock();
  
  	printf("BEE is settings as: %s mode, %s %d key
  ",
  	       (p->mode == ECB_MODE) ? "ECB" : "CTR",
  	       (p->key_method == SOFT_KEY) ? "SOFT" : "SNVS HW",
  	       (p->mode == ECB_MODE) ? 128 : 256);
  
  	return CMD_RET_SUCCESS;
  }
  
  int bee_test(struct bee_parameters *p, int region)
  {
  	u32 result = 0, range, address;
  	int i, val;
  	/*
  	 * Test instruction running in AES Region:
  	 * int test(void)
  	 * {
  	 *	return 0x55aa55aa;
  	 * }
  	 * Assemble:
  	 * 0xe59f0000: ldr r0, [pc]
  	 * 0xe12fff1e: bx lr
  	 * 0x55aa55aa: 0x55aa55aa
  	 */
  	u32 inst[3] = {0xe59f0000, 0xe12fff1e, 0x55aa55aa};
  
  	/* Cache enabled? */
  	if ((get_cr() & (CR_I | CR_C)) != (CR_I | CR_C)) {
  		printf("Enable dcache and icache first!
  ");
  		return CMD_RET_FAILURE;
  	}
  
  	printf("Test Region %d
  Begin Data test: Writing... ", region);
  
  	range = (region == 0) ? p->size1 : p->size2;
  	address = (region == 0) ? AES_REGION0_ADDR : AES_REGION1_ADDR;
  	for (i = 0; i < range; i = i + 4)
  		writel(i, address + i);
  
  	printf("Finshed Write!
  ");
  
  	flush_dcache_range(address, address + range);
  
  	printf("Reading... ");
  	for (i = 0; i < range; i = i + 4) {
  		val = readl(address + i);
  		if (val != i)
  			result++;
  	}
  	printf("Finshed Read!
  ");
  
  	if (result > 0)
  		printf("BEE Data Test check Failed!
  ");
  	else
  		printf("BEE Data Test Check Passed!
  ");
  
  	for (i = 0; i < ARRAY_SIZE(inst); i++)
  		writel(inst[i], address + (i * 4));
  
  	flush_dcache_range(address, address + sizeof(inst));
  
  	val = ((int (*)(void))address)();
  
  	printf("
  Bee Instruction test, Program:
  "
  	       "int test(void)
  "
  	       "{
  "
  	       "      return 0x55aa55aa;
  "
  	       "}
  "
  	       "Assemble:
  "
  	       "0xe59f0000: ldr r0, [pc]
  "
  	       "0xe12fff1e: bx lr
  "
  	       "0x55aa55aa: 0x55aa55aa
  "
  	       "Runnint at 0x%x
  ", address);
  	if (val == 0x55aa55aa)
  		printf("Bee Instruction Test Passed!
  ");
  	else
  		printf("Bee Instruction Test Failed!
  ");
  
  	return CMD_RET_SUCCESS;
  }
  
  static int region_valid(u32 start, u32 size)
  {
  	if ((start < PHYS_SDRAM) || (start >= (start + size - 1)) ||
  	    (start >= (PHYS_SDRAM + PHYS_SDRAM_SIZE - 1))) {
  		printf("Invalid start 0x%x, size 0x%x
  ", start, size);
  		return -EINVAL;
  	}
  
  	if (size > SZ_512M) {
  		printf("The region size exceeds SZ_512M
  ");
  		return -EINVAL;
  	}
  
  	if ((start & 0xFFFF) && (size & 0xFFFF)) {
  		printf("start or size not 64KB aligned!
  ");
  		return -EINVAL;
  	}
  
  	/* 128K for U-Boot Stack */
  	if ((start + size - 1) >= (gd->start_addr_sp - SZ_128K)) {
  		printf("Overlap with uboot execution environment!
  "
  		       "Decrease size or start
  ");
  		return -EINVAL;
  	}
  
  	return 0;
  }
  
  static int do_bee_init(cmd_tbl_t *cmdtp, int flag, int argc,
  		       char * const argv[])
  {

  	u32 start, size;

  	int ret;
  	struct bee_parameters *p = ¶
  
  #if defined(CONFIG_SYS_ARM_CACHE_WRITETHROUGH)

  	enum dcache_option option = DCACHE_WRITETHROUGH & ~TTB_SECT_XN_MASK;

  #else

  	enum dcache_option option = DCACHE_WRITEBACK & ~TTB_SECT_XN_MASK;

  #endif
  
  	if (argc > 5)
  		return CMD_RET_USAGE;

  #ifdef CONFIG_MX6
  	if (check_module_fused(MX6_MODULE_BEE)) {
  		printf("BEE is fused, disable it!
  ");
  		return CMD_RET_FAILURE;

  	}

  #endif

  
  	/* Cache enabled? */
  	if ((get_cr() & (CR_I | CR_C)) != (CR_I | CR_C)) {
  		/*
  		 * Here we need icache and dcache both enabled, because
  		 * we may take the protected region for instruction and
  		 * data usage. And icache and dcache both enabled are
  		 * better for performance.
  		 */
  		printf("Please enable dcache and icache first!
  ");
  		return CMD_RET_FAILURE;
  	}
  
  	p->key_method = SOFT_KEY;
  	p->mode = ECB_MODE;
  	p->start1 = PHYS_SDRAM;
  	p->size1 = SZ_512M;
  	p->start2 = IRAM_BASE_ADDR;
  	p->size2 = IRAM_SIZE;
  
  	if (argc == 2) {
  		p->key_method = (int)simple_strtoul(argv[1], NULL, 16);
  		p->mode = ECB_MODE;
  		p->start1 = PHYS_SDRAM;
  		p->size1 = SZ_512M;
  	} else if (argc == 3) {
  		p->key_method = (int)simple_strtoul(argv[1], NULL, 16);
  		p->mode = (int)simple_strtoul(argv[2], NULL, 10);
  		p->start1 = PHYS_SDRAM;
  		p->size1 = SZ_512M;
  	} else if ((argc == 4) || (argc == 5)) {
  		p->key_method = (int)simple_strtoul(argv[1], NULL, 16);
  		p->mode = (int)simple_strtoul(argv[2], NULL, 10);
  		start = (u32)simple_strtoul(argv[3], NULL, 16);
  		/* Default size that AES Region0 can protected */
  		size = SZ_512M;
  		if (argc == 5)
  			size = (u32)simple_strtoul(argv[4], NULL, 16);
  		p->start1 = start;
  		p->size1 = size;
  	}
  
  	if ((p->key_method != SOFT_KEY) && (p->key_method != SNVS_KEY))
  		return CMD_RET_USAGE;
  
  	if ((p->mode != ECB_MODE) && (p->mode != CTR_MODE))
  		return CMD_RET_USAGE;
  
  	/*
  	 * No need to check region valid for IRAM, since it is fixed.
  	 * Only check DRAM region here.
  	 */
  	if (region_valid(p->start1, p->size1))
  		return CMD_RET_FAILURE;
  
  	ret = bee_init(p);
  	if (ret)
  		return CMD_RET_FAILURE;
  
  	/*
  	 * Set DCACHE OFF to AES REGION0 and AES REGION1 first
  	 * to avoid possible unexcepted cache settings.
  	 */
  	mmu_set_region_dcache_behaviour(AES_REGION0_ADDR, SZ_1G, DCACHE_OFF);
  
  	mmu_set_region_dcache_behaviour(AES_REGION0_ADDR, p->size1, option);
  
  	mmu_set_region_dcache_behaviour(AES_REGION1_ADDR, p->size2, option);
  
  	printf("Access Region 0x%x - 0x%x to protect 0x%x - 0x%x
  "
  	       "Do not directly access 0x%x - 0x%x
  "
  	       "Access Region 0x%x - 0x%x to protect 0x%x - 0x%x
  "
  	       "Do not directly access 0x%x - 0x%x
  ",
  	       AES_REGION0_ADDR, AES_REGION0_ADDR + p->size1 - 1,
  	       p->start1, p->start1 + p->size1 - 1,
  	       p->start1, p->start1 + p->size1 - 1,
  	       AES_REGION1_ADDR, AES_REGION1_ADDR + p->size2 - 1,
  	       p->start2, p->start2 + p->size2 - 1,
  	       p->start2, p->start2 + p->size2 - 1);
  
  	bee_inited = 1;
  
  	return CMD_RET_SUCCESS;
  }
  
  static int do_bee_test(cmd_tbl_t *cmdtp, int flag, int argc,
  		       char * const argv[])
  {
  	int ret;
  	int region;
  
  	if (bee_inited == 0) {
  		printf("Bee not initialized, run bee init first!
  ");
  		return CMD_RET_FAILURE;
  	}
  	if (argc > 2)
  		return CMD_RET_USAGE;
  
  	region = 0;
  	if (argc == 2)
  		region = (int)simple_strtoul(argv[1], NULL, 16);
  	/* Only two regions are supported, 0 and 1 */
  	if (region >= 2)
  		return CMD_RET_USAGE;
  
  	ret = bee_test(&para, region);
  	if (ret)
  		return CMD_RET_FAILURE;
  
  	return CMD_RET_SUCCESS;
  }
  
  static cmd_tbl_t cmd_bmp_sub[] = {
  	U_BOOT_CMD_MKENT(init, 5, 0, do_bee_init, "", ""),
  	U_BOOT_CMD_MKENT(test, 2, 0, do_bee_test, "", ""),
  };
  
  static int do_bee_ops(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
  {
  	cmd_tbl_t *c;
  
  	c = find_cmd_tbl(argv[1], &cmd_bmp_sub[0], ARRAY_SIZE(cmd_bmp_sub));
  
  	/* Drop off the 'bee' command argument */
  	argc--;
  	argv++;
  
  	if (c)
  		return  c->cmd(cmdtp, flag, argc, argv);
  	else
  		return CMD_RET_USAGE;
  }
  
  U_BOOT_CMD(
  	bee, CONFIG_SYS_MAXARGS, 1, do_bee_ops,
  	"BEE function test",
  	"init [key] [mode] [start] [size]	- BEE block initial
  "
  	"  key: 0 | 1, 0 means software key, 1 means SNVS random key
  "
  	"  mode: 0 | 1, 0 means ECB mode, 1 means CTR mode
  "
  	"  start: start address that you want to protect
  "
  	"  size: The size of the area that you want to protect
  "
  	"  start and end(start + size) addr both should be 64KB aligned.
  "
  	"
  "
  	" After initialization, the mapping:
  "
  	" 1. [0x10000000 - (0x10000000 + size - 1)] <--->
  "
  	"    [start - (start + size - 1)]
  "
  	"    Here [start - (start + size -1)] is fixed mapping to
  "
  	"    [0x10000000 - (0x10000000 + size - 1)], whatever start is.
  "
  	" 2. [0x30000000 - (0x30000000 + IRAM_SIZE - 1)] <--->
  "
  	"    [IRAM_BASE_ADDR - (IRAM_BASE_ADDR + IRAM_SIZE - 1)]
  "
  	"
  "
  	" Note: Here we only use AES region 0 to protect the DRAM
  "
  	"       area that you specified, max size SZ_512M.
  "
  	"       AES region 1 is used to protect IRAM area.
  "
  	" Example:
  "
  	" 1. bee init 1 1 0xa0000000 0x10000
  "
  	"    Access 0x10000000 - 0x10010000 to protect 0xa0000000 - 0xa0010000
  "
  	" 2. bee init 1 1 0x80000000 0x20000
  "
  	"    Access 0x10000000 - 0x10020000 to protect 0x80000000 - 0x80020000
  "
  	"
  "
  	" Default configuration if only `bee init` without any args:
  "
  	" 1. software key
  "
  	" 2. ECB mode
  "
  	" 3. Address protected:
  "
  	"   Remapped Region0: PHYS_SDRAM - PHYS_SDRAM + SZ_512M
  "
  	"   Remapped Region1: IRAM_BASE_ADDR - IRAM_BASE_ADDR + IRAM_SIZE
  "
  	" 4. Default Mapping for 6UL:
  "
  	"   [0x10000000 - 0x2FFFFFFF] <-> [0x80000000 - 0x9FFFFFFF]
  "
  	"   [0x30000000 - 0x3001FFFF] <-> [0x00900000 - 0x0091FFFF]
  "
  	"
  "
  	"bee test [region]			- BEE function test
  "
  	"  region: 0 | 1, 0 means region0, 1 means regions1
  "
  );