Bad block support for environment variables partition in OneNAND

index cd44781..0e7fe79 100644
@@ -56,23 +56,51 @@ uchar env_get_char_spec(int index)
 	return (*((uchar *) (gd->env_addr + index)));
 }
 
+static int readenv(ulong ofs)
+{
+	struct mtd_info *mtd = &onenand_mtd;
+	size_t retlen;
+
+	if (mtd->block_isbad(mtd, ofs)) {
+		printf("readenv: skip bad block at 0x%08x\n", (unsigned)ofs);
+		return 1;
+	}
+
+	if (mtd->read(mtd, ofs, mtd->writesize, &retlen, (u_char *) env_ptr)) {
+		printf("readenv: read failed at 0x%08x\n", (unsigned)ofs);
+		return 1;
+	}
+
+	printf("readenv: read params at 0x%08x\n", (unsigned) ofs);
+	return 0;
+}
+
 void env_relocate_spec(void)
 {
 	struct onenand_chip *this = &onenand_chip;
 	unsigned long env_addr;
 	int use_default = 0;
-	size_t retlen;
+	unsigned env_len;
+	unsigned slc = 0;
 
-	env_addr = CONFIG_ENV_ADDR;
-	if (FLEXONENAND(this))
-		env_addr = CONFIG_ENV_ADDR_FLEX;
+	env_addr = FLEXONENAND(this) ? CONFIG_ENV_ADDR_FLEX :
+					CONFIG_ENV_ADDR;
+
+	env_len = FLEXONENAND(this) ? CONFIG_ENV_SIZE_FLEX :
+					CONFIG_ENV_SIZE;
 
 	/* Check OneNAND exist */
-	if (onenand_mtd.writesize)
-		/* Ignore read fail */
-		onenand_read(&onenand_mtd, env_addr, onenand_mtd.writesize,
-			     &retlen, (u_char *) env_ptr);
-	else
+	if (onenand_mtd.writesize) {
+		while (env_addr < env_addr + env_len) {
+			if (!readenv(env_addr))
+				break;
+			onenand_get_block(this, env_addr, &slc);
+			if (slc)
+				env_addr += 1 << (this->erase_shift -1);
+			else
+				env_addr += 1 << this->erase_shift;
+		}
+	} else
 		onenand_mtd.writesize = MAX_ONENAND_PAGESIZE;
 
 	if (crc32(0, env_ptr->data, ONENAND_ENV_SIZE(onenand_mtd)) !=

@@ -90,37 +118,71 @@ void env_relocate_spec(void)
 	gd->env_valid = 1;
 }
 
-int saveenv(void)
+static int writeenv(ulong ofs)
 {
+	struct mtd_info *mtd = &onenand_mtd;
 	struct onenand_chip *this = &onenand_chip;
-	unsigned long env_addr = CONFIG_ENV_ADDR;
 	struct erase_info instr = {
 		.callback	= NULL,
 	};
 	size_t retlen;
+	unsigned slc = 0;
 
-	instr.len = CONFIG_ENV_SIZE;
-	if (FLEXONENAND(this)) {
-		env_addr = CONFIG_ENV_ADDR_FLEX;
-		instr.len = CONFIG_ENV_SIZE_FLEX;
-		instr.len <<= onenand_mtd.eraseregions[0].numblocks == 1 ?
-				1 : 0;
+	if (mtd->block_isbad(mtd, ofs)) {
+		printf("writeenv: skip bad block at 0x%08x\n", (unsigned)ofs);
+		return 1;
 	}
-	instr.addr = env_addr;
-	instr.mtd = &onenand_mtd;
-	if (onenand_erase(&onenand_mtd, &instr)) {
-		printf("OneNAND: erase failed at 0x%08lx\n", env_addr);
+
+	onenand_get_block(this, ofs, &slc);
+	if (slc)
+		instr.len += 1 << (this->erase_shift -1);
+	else
+		instr.len += 1 << this->erase_shift;
+	instr.addr = ofs;
+
+	if (mtd->erase(mtd, &instr)) {
+		printf("writeenv: erase failed at 0x%08lx\n", ofs);
+		if(mtd->block_markbad(mtd, ofs))
+			printf("writeenv: Mark bad failed at 0x%08lx\n", ofs);
+		return 1;
+	}
+
+	if (mtd->write(mtd, ofs, mtd->writesize, &retlen, (u_char *) env_ptr)) {
+		printf("writeenv: write failed at 0x%08x\n", (unsigned) ofs);
+		mtd->erase(mtd, &instr);
+		mtd->block_markbad(mtd, ofs);
 		return 1;
 	}
 
+	printf("writeenv: written params at 0x%08x\n", (unsigned) ofs);
+	return 0;
+}
+
+int saveenv(void)
+{
+	struct onenand_chip *this = &onenand_chip;
+	unsigned long env_addr;
+	unsigned env_len;
+	unsigned slc = 0;
+
+	env_addr = FLEXONENAND(this) ? CONFIG_ENV_ADDR_FLEX :
+					CONFIG_ENV_ADDR;
+
+	env_len = FLEXONENAND(this) ? CONFIG_ENV_SIZE_FLEX :
+					CONFIG_ENV_SIZE;
+
 	/* update crc */
 	env_ptr->crc =
 	    crc32(0, env_ptr->data, ONENAND_ENV_SIZE(onenand_mtd));
 
-	if (onenand_write(&onenand_mtd, env_addr, onenand_mtd.writesize, &retlen,
-	     (u_char *) env_ptr)) {
-		printf("OneNAND: write failed at 0x%08x\n", instr.addr);
-		return 2;
+	while (env_addr < env_addr + env_len) {
+		if(!writeenv(env_addr))
+			break;
+		onenand_get_block(this, env_addr, &slc);
+		if (slc)
+			env_addr += 1 << (this->erase_shift -1);
+		else
+			env_addr += 1 << this->erase_shift;
 	}
 
 	return 0;