handler: make NAND flashing streamable

Use copyimage() for both NOR and NAND flashes.
The idea is to use custom writeimage callback for to account for flash
erasing, bad blocks and rewinding during writing.

This change automatically adds encrypted and compressed NAND images
support. Previous implementation silently installs encrypted image into
NAND flash without decrypting it (an attempt to install encrypted uboot
image into NAND flash bricks the board).

Additional changes:

- Do not always return -1 on error: propagate specific negative error
  code (errno) from a flash handler instead.

- Ensure that start offset (img->seek) is erase-block-aligned
  (otherwise we were loosing some old data while erasing first block).

- Try to unlock MTD block even if mtd_is_locked() is not supported.
  Some flashes support mtd_lock()/mtd_unlock() but do not support
  mtd_is_locked().

- Ignore "not supported" errors for mtd_erase() and mtd_mark_bad().

- Mark block as bad if mtd_erase() returned EIO.

- Writing to NOR flash accounts for bad blocks.

- Avoid flash_write_nand() issue for the following scenario:

    1) flash_erase_sector(mtdnum, img->seek, img->size) is executed.
    2) mtd_write() failed for a last flash block occupied by the image.
       (let's call it block X).
    3) Block X is erased and marked as bad.
    4) mtdoffset is adjusted past block X and past flash memory region
       erased by flash_erase_sector() in step #1.
    5) Further mtd_write() to non-erased flash memory past block X
       fails (it should've been erased before writing into it).

Signed-off-by: Viacheslav Volkov <convivo@nibeconsultants.eu>
Signed-off-by: Måns Andersson <mans.andersson@nibe.se>
---
 .gitignore                |    1 +
 configs/test_defconfig    |    2 +
 handlers/flash_handler.c  |  638 ++++++++++------
 test/Makefile             |    1 +
 test/test_flash_handler.c | 1456 +++++++++++++++++++++++++++++++++++++
 5 files changed, 1878 insertions(+), 220 deletions(-)
 create mode 100644 test/test_flash_handler.c

Hi Viaceslav,

On 13.01.25 11:13, 'Viacheslav Volkov' via swupdate wrote:
> Use copyimage() for both NOR and NAND flashes.
> The idea is to use custom writeimage callback for to account for flash
> erasing, bad blocks and rewinding during writing.
>

Yes, issue is known. And yes, the solution is to add a callback to write
the buffer, as you did, that manages the blocks. I agree on this
solution. The issue was reported and stored in roadmap, but nobody was
interested on this. So thanks for fixing it.

I just ask you to split the patch, and move the test/ to a separate
patch for better review.

> This change automatically adds encrypted and compressed NAND images
> support. Previous implementation silently installs encrypted image into
> NAND flash without decrypting it (an attempt to install encrypted uboot
> image into NAND flash bricks the board).
>
> Additional changes:
>
> - Do not always return -1 on error: propagate specific negative error
>    code (errno) from a flash handler instead.
>
> - Ensure that start offset (img->seek) is erase-block-aligned
>    (otherwise we were loosing some old data while erasing first block).
>
> - Try to unlock MTD block even if mtd_is_locked() is not supported.
>    Some flashes support mtd_lock()/mtd_unlock() but do not support
>    mtd_is_locked().

I am unsure if this is a good idea. Theoretically yes, but I had also
the case where locking should be disabled at due to issues in kernel or
in U-Boot, or even due to incompatibility between both. That is similar
to the issue you report. But well, an attempt to mtd_unlock() should be
harmlos.

>
> - Ignore "not supported" errors for mtd_erase() and mtd_mark_bad().
>
> - Mark block as bad if mtd_erase() returned EIO.
>
> - Writing to NOR flash accounts for bad blocks.
>
> - Avoid flash_write_nand() issue for the following scenario:
>
>      1) flash_erase_sector(mtdnum, img->seek, img->size) is executed.

ok

>      2) mtd_write() failed for a last flash block occupied by the image.
>         (let's call it block X).
>      3) Block X is erased and marked as bad.
>      4) mtdoffset is adjusted past block X and past flash memory region
>         erased by flash_erase_sector() in step #1.
>      5) Further mtd_write() to non-erased flash memory past block X
>         fails (it should've been erased before writing into it).

Urgh....you are right. Thatnks for discovering this.

>
> Signed-off-by: Viacheslav Volkov <convivo@nibeconsultants.eu>
> Signed-off-by: Måns Andersson <mans.andersson@nibe.se>
> ---
>   .gitignore                |    1 +
>   configs/test_defconfig    |    2 +
>   handlers/flash_handler.c  |  638 ++++++++++------
>   test/Makefile             |    1 +
>   test/test_flash_handler.c | 1456 +++++++++++++++++++++++++++++++++++++
>   5 files changed, 1878 insertions(+), 220 deletions(-)
>   create mode 100644 test/test_flash_handler.c
>
> diff --git a/.gitignore b/.gitignore
> index 4755ff07..313a8b55 100644
> --- a/.gitignore
> +++ b/.gitignore
> @@ -52,6 +52,7 @@ tools/*.out
>   test/test_crypt
>   test/test_hash
>   test/test_json
> +test/test_flash_handler
>   test/test_server_hawkbit
>   test/test_util
>   test/test_verify
> diff --git a/configs/test_defconfig b/configs/test_defconfig
> index ac8fff6d..5fef117b 100644
> --- a/configs/test_defconfig
> +++ b/configs/test_defconfig
> @@ -14,3 +14,5 @@ CONFIG_LUASCRIPTHANDLER=y
>   CONFIG_RAW=y
>   CONFIG_REMOTE_HANDLER=y
>   CONFIG_SHELLSCRIPTHANDLER=y
> +CONFIG_MTD=y
> +CONFIG_CFI=y
> diff --git a/handlers/flash_handler.c b/handlers/flash_handler.c
> index a6b22cd5..02cd1b34 100644
> --- a/handlers/flash_handler.c
> +++ b/handlers/flash_handler.c
> @@ -21,6 +21,7 @@
>   #include <stdlib.h>
>   #include <stdbool.h>
>   #include <errno.h>
> +#include <assert.h>
>   #include <linux/version.h>
>   #include <sys/ioctl.h>
>
> @@ -33,6 +34,57 @@
>
>   #define PROCMTD	"/proc/mtd"
>   #define LINESIZE	80
> +#define EMPTY_BYTE	0xFF
> +#define ROUND_UP(N, S)	((((N) + (S) - 1) / (S)) * (S))
> +#define ROUND_DOWN(N, S)	(((N) / (S)) * (S))

Maybe these defines should be global, that is moved to util.h ?

> +
> +/* According to linux kernel ENOTSUPP is not a standard error code and should
> + * be avoided in new patches. EOPNOTSUPP should be used instead. Unfortunately
> + * some drivers still return ENOTSUPP. Do not confuse with ENOTSUP! See also:
> + * https://lists.gnu.org/archive/html/bug-glibc/2002-08/msg00017.html
> + * https://linux-fsdevel.vger.kernel.narkive.com/pERNbmWG/kernel-glibc-eopnotsupp-vs-enotsup-vs-enotsupp-also-rfc-posix-acl-kernel-infrastructure
> + */
> +#ifndef ENOTSUPP
> +#define ENOTSUPP 524 /* Operation is not supported */
> +static const char *strerror_enotsupp(int code)
> +{
> +	if (code == ENOTSUPP)
> +		return "Operation is not supported";
> +	return strerror(code);
> +}
> +#define STRERROR(code) strerror_enotsupp(code)
> +#else
> +#define STRERROR(code) strerror(code)
> +#endif
> +
> +static inline bool is_not_supported(int code)
> +{
> +	return (code == EOPNOTSUPP) || (code == ENOTSUP) || (code == ENOTSUPP);
> +}

These are global, too. It could be that other part of code will need
this. Managing ENOTSUP is unrelated to MTD / NAND, and can be set global.

> +
> +static inline int mtd_error(int mtdnum, const char *func, int eb)
> +{
> +	int code = errno;
> +	ERROR("mtd%d: mtd_%s() failed at block %d: %s (errno = %d)", mtdnum,
> +	      func, eb, STRERROR(code), code);
> +	return -code;
> +}
> +#define MTD_ERROR(func) mtd_error(priv->mtdnum, (func), priv->eb)
> +
> +#define MTD_TRACE(func, msg) do { \
> +		int _code = errno; \
> +		TRACE("mtd%d: mtd_%s() %s at block %d: %s (errno = %d)", \
> +		      priv->mtdnum, (func), (msg), priv->eb, \
> +		      STRERROR(_code), _code); \
> +	} while (0)
> +#define MTD_TRACE_FAILED(func) MTD_TRACE((func), "failed")
> +#define MTD_TRACE_NOT_SUPPORTED(func) MTD_TRACE((func), "not supported")
> +
> +static inline int too_many_bad_blocks(int mtdnum)
> +{
> +	ERROR("mtd%d: too many bad blocks", mtdnum);
> +	return -ENOSPC;
> +}
>
>   void flash_handler(void);
>
> @@ -69,293 +121,438 @@ static inline int buffer_check_pattern(unsigned char *buffer, size_t size,
>    * The function reassembles nandwrite from mtd-utils
>    * dropping all options that are not required here.
>    */
> -
>   static void erase_buffer(void *buffer, size_t size)
>   {
> -	const uint8_t kEraseByte = 0xff;
> -
>   	if (buffer != NULL && size > 0)
> -		memset(buffer, kEraseByte, size);
> +		memset(buffer, EMPTY_BYTE, size);
>   }
>
> -static int flash_write_nand(int mtdnum, struct img_type *img)
> +struct flash_priv {
> +	/* File descriptor of a flash device (/dev/mtdX) to write into: */
> +	int fdout;
> +	/* Index of a flash device (X in /dev/mtdX): */
> +	int mtdnum;
> +	/* Number of image bytes we still need to write into flash: */
> +	long long imglen;
> +	/* A buffer for caching data soon to be written to flash: */
> +	unsigned char *filebuf;
> +	/* Amount of bytes we've read from image file into filebuf: */
> +	int filebuf_len;
> +	/* An offset within filebuf to write to flash starting from: */
> +	int writebuf_offset;
> +	/* Current flash erase block: */
> +	int eb;
> +	/*
> +	 * The following data is kept here only to speed-up execution:
> +	 */
> +	int eb_end; /* First invalid erase block (after flash memory end). */
> +	struct mtd_dev_info *mtd;
> +	libmtd_t libmtd;
> +	bool is_nand; /* is it NAND or some other (e.g. NOR) flash type? */
> +	bool check_bad; /* do we need to check for bad blocks? */
> +	bool check_locked; /* do we need to check whether a block is locked? */
> +	bool do_unlock; /* do we need to try to unlock a block? */
> +	unsigned char *readout_buf; /* a buffer to read erase block into */
> +};
> +
> +static int erase_block(struct flash_priv *priv)
>   {
> -	char mtd_device[LINESIZE];
> -	struct flash_description *flash = get_flash_info();
> -	struct mtd_dev_info *mtd = &flash->mtd_info[mtdnum].mtd;
> -	int pagelen;
> -	bool baderaseblock = false;
> -	long long imglen = 0;
> -	long long blockstart = -1;
> -	long long offs;
> -	unsigned char *filebuf = NULL;
> -	size_t filebuf_max = 0;
> -	size_t filebuf_len = 0;
> -	long long mtdoffset = img->seek;
> -	int ifd = img->fdin;
> -	int fd = -1;
> -	bool failed = true;
>   	int ret;
> -	unsigned char *writebuf = NULL;
> +	ret = mtd_erase(priv->libmtd, priv->mtd, priv->fdout, priv->eb);
> +	if (ret) {
> +		if (is_not_supported(errno)) {
> +			/* Some MTD drivers in linux kernel may not initialize
> +			 * mtd->erasesize or master->_erase, hence we expect
> +			 * these errors in such cases. */
> +			MTD_TRACE_NOT_SUPPORTED("erase");
> +			return 0;
> +		} else if (errno == EIO) { /* May happen for a bad block. */
> +			MTD_TRACE_FAILED("erase");
> +			return -EIO;
> +		} else
> +			return MTD_ERROR("erase");
> +	}
> +	return 0;
> +}
>
> -	/*
> -	 * if nothing to do, returns without errors
> -	 */
> -	if (!img->size)
> -		return 0;
> +static int mark_bad_block(struct flash_priv *priv)
> +{
> +	int ret;
> +	TRACE("mtd%d: marking block %d (offset %lld) as bad", priv->mtdnum,
> +	      priv->eb, (long long)(priv->eb) * priv->mtd->eb_size);
> +	ret = mtd_mark_bad(priv->mtd, priv->fdout, priv->eb);
> +	if (ret) {
> +		if (is_not_supported(errno)) {
> +			/* Some MTD drivers in linux kernel may not initialize
> +			 * master->_block_markbad, hence we expect these errors
> +			 * in such cases. */
> +			MTD_TRACE_NOT_SUPPORTED("mark_bad");
> +		} else
> +			return MTD_ERROR("mark_bad");
> +	}
> +	return 0;
> +}
>
> -	if (mtdoffset & (mtd->min_io_size - 1)) {
> -		ERROR("The start address is not page-aligned !\n"
> -			   "The pagesize of this NAND Flash is 0x%x.\n",
> -			   mtd->min_io_size);
> -		return -EIO;
> +/*
> + * Read input data from (*p_in_buf) (size = (*p_in_len)), adjust both
> + * (*p_in_buf) and (*p_in_len) while reading. Write input data to a proper
> + * offset within priv->filebuf.
> + *
> + * On success return 0 and initialize output arguments:
> + * - (*p_wbuf) (a pointer within priv->filebuf to write data into flash
> + *   starting from).
> + * - (*p_to_write) (number of bytes from (*p_wbuf) to write into flash).
> + * If necessary append padding (EMPTY_BYTE) to the data to be written.
> + * (*p_to_write) is guaranteed to be multiple of priv->mtd->min_io_size.
> + *
> + * On failure return -1. In this case there is nothing to write to flash, all
> + * data from (*p_in_buf), has been read and stored in priv->filebuf.
> + * Need to wait for next flash_write() call to get more data.
> + */
> +static int read_data(struct flash_priv *priv, const unsigned char **p_in_buf,
> +		size_t *p_in_len, unsigned char **p_wbuf, int *p_to_write)
> +{
> +	int read_available, to_read, write_available, to_write;
> +	size_t in_len = *p_in_len;
> +
> +	assert(priv->filebuf_len <= priv->mtd->eb_size);
> +	assert(priv->writebuf_offset <= priv->filebuf_len);
> +	assert((priv->writebuf_offset % priv->mtd->min_io_size) == 0);
> +	assert(priv->imglen >= in_len);
> +
> +	/* Read as much as possible data from (*p_in_buf) to priv->filebuf: */
> +	read_available = priv->mtd->eb_size - priv->filebuf_len;
> +	assert(read_available >= 0);
> +	to_read = (in_len > read_available) ? read_available : in_len;
> +	assert(to_read <= read_available);
> +	assert(to_read <= in_len);
> +	memcpy(priv->filebuf + priv->filebuf_len, *p_in_buf, to_read);
> +	*p_in_buf += to_read;
> +	(*p_in_len) -= to_read;
> +	priv->filebuf_len += to_read;
> +	priv->imglen -= to_read;
> +
> +	if (priv->imglen == 0) {
> +		/* We've read all image data available.
> +		 * Add padding to priv->filebuf if necessary: */
> +		int len, pad_bytes;
> +		len = ROUND_UP(priv->filebuf_len, priv->mtd->min_io_size);
> +		assert(len >= priv->filebuf_len);
> +		assert(len <= priv->mtd->eb_size);
> +		pad_bytes = len - priv->filebuf_len;
> +		assert(pad_bytes >= 0);
> +		erase_buffer(priv->filebuf + priv->filebuf_len, pad_bytes);
> +		priv->filebuf_len = len;
>   	}
>
> -	pagelen = mtd->min_io_size;
> -	imglen = img->size;
> -	snprintf(mtd_device, sizeof(mtd_device), "/dev/mtd%d", mtdnum);
> +	write_available = priv->filebuf_len - priv->writebuf_offset;
> +	assert(write_available >= 0);

SWUpdate exits with error here. For design, SWUpdate shouldn't exit
during an update, but should report enough information in ERROR to allow
to follow the reason.

>
> -	if ((imglen / pagelen) * mtd->min_io_size > mtd->size - mtdoffset) {
> -		ERROR("Image %s does not fit into mtd%d", img->fname, mtdnum);
> -		return -EIO;
> +	to_write = ROUND_DOWN(write_available, priv->mtd->min_io_size);
> +	assert(to_write <= write_available);
> +	assert((to_write % priv->mtd->min_io_size) == 0);
> +
> +	if (to_write == 0) {
> +		/* Got not enough data to write. */
> +		return -1; /* Wait for more data in next flash_write() call. */
>   	}
>
> -	/* Flashing to NAND is currently not streamable */
> -	if (img->install_directly) {
> -		ERROR("Raw NAND not streamable");
> -		return -EINVAL;
> +	if (priv->is_nand) {
> +		/* For NAND flash limit amount of data to be written to a
> +		 * single page. This allows us to skip writing "empty" pages
> +		 * (filled with EMPTY_BYTE).
> +		 *
> +		 * Note: for NOR flash typical min_io_size is 1. Writing 1 byte
> +		 * at time is not practical. */
> +		to_write = priv->mtd->min_io_size;
>   	}
>
> -	if(flash_erase_sector(mtdnum, img->seek, img->size)) {
> -		ERROR("I cannot erasing %s",
> -			img->device);
> -		return -1;
> +	*p_wbuf = priv->filebuf + priv->writebuf_offset;
> +	*p_to_write = to_write;
> +	return 0; /* Need to write some data. */
> +}
> +
> +/*
> + * Check and process current erase block. Return:
> + * - 0 if proper erase block has been found.
> + * - 1 if current erase block is bad and need to continue the search.
> + * - Negative errno code on system error (in this case error reporting is
> + *   already done in this function).
> + */
> +static int process_new_erase_block(struct flash_priv *priv)
> +{
> +	int ret;
> +
> +	if (priv->check_bad) {
> +		int is_bad = mtd_is_bad(priv->mtd, priv->fdout, priv->eb);
> +		if (is_bad > 0)
> +			return 1;
> +		if (is_bad < 0) {
> +			if (is_not_supported(errno)) {
> +				/* I don't know whether such cases really
> +				 * exist.. Let's handle them just in case (as
> +				 * it is currently implemented in mtd-utils and
> +				 * in previous version of swupdate). */
> +				MTD_TRACE_NOT_SUPPORTED("is_bad");
> +				priv->check_bad = false;
> +			} else
> +				return MTD_ERROR("is_bad");
> +		}
>   	}
>
> -	if ((fd = open(mtd_device, O_RDWR)) < 0) {
> -		ERROR( "%s: %s: %s", __func__, mtd_device, strerror(errno));
> -		return -ENODEV;
> +	if (priv->check_locked) {
> +		int is_locked = mtd_is_locked(priv->mtd, priv->fdout,
> +		                              priv->eb);
> +		if (is_locked < 0) {
> +			if (is_not_supported(errno)) {
> +				/* Some MTD drivers in linux kernel may not
> +				 * initialize mtd->_is_locked, hence we expect
> +				 * these errors in such cases. At the same time
> +				 * the driver can initialize mtd->_unlock,
> +				 * hence we shall try to execute mtd_unlock()
> +				 * in such cases. */
> +				MTD_TRACE_NOT_SUPPORTED("is_locked");
> +				priv->check_locked = false;
> +				priv->do_unlock = true;
> +			} else
> +				return MTD_ERROR("is_locked");
> +		} else
> +			priv->do_unlock = (bool)is_locked;
>   	}
>
> -	filebuf_max = mtd->eb_size / mtd->min_io_size * pagelen;
> -	filebuf = calloc(1, filebuf_max);
> -	erase_buffer(filebuf, filebuf_max);
> +	if (priv->do_unlock) {
> +		ret = mtd_unlock(priv->mtd, priv->fdout, priv->eb);
> +		if (ret) {
> +			if (is_not_supported(errno)) {
> +				/* Some MTD drivers in linux kernel may not
> +				 * initialize mtd->_unlock, hence we expect
> +				 * these errors in such cases. */
> +				MTD_TRACE_NOT_SUPPORTED("unlock");
> +				priv->check_locked = false;
> +				priv->do_unlock = false;
> +			} else
> +				return MTD_ERROR("unlock");
> +		}
> +	}
>
>   	/*
> -	 * Get data from input and write to the device while there is
> -	 * still input to read and we are still within the device
> -	 * bounds. Note that in the case of standard input, the input
> -	 * length is simply a quasi-boolean flag whose values are page
> -	 * length or zero.
> +	 * NAND flash should always be erased (follow "write once rule").
> +	 * For other flash types check if the flash is already empty.
> +	 * In case of NOR flash, it can save a significant amount of time
> +	 * because erasing a NOR flash is very time expensive.
>   	 */
> -	while ((imglen > 0 || writebuf < filebuf + filebuf_len)
> -		&& mtdoffset < mtd->size) {
> -		/*
> -		 * New eraseblock, check for bad block(s)
> -		 * Stay in the loop to be sure that, if mtdoffset changes because
> -		 * of a bad block, the next block that will be written to
> -		 * is also checked. Thus, we avoid errors if the block(s) after the
> -		 * skipped block(s) is also bad
> -		 */
> -		while (blockstart != (mtdoffset & (~mtd->eb_size + 1))) {
> -			blockstart = mtdoffset & (~mtd->eb_size + 1);
> -			offs = blockstart;
> -
> -			/*
> -			 * if writebuf == filebuf, we are rewinding so we must
> -			 * not reset the buffer but just replay it
> -			 */
> -			if (writebuf != filebuf) {
> -				erase_buffer(filebuf, filebuf_len);
> -				filebuf_len = 0;
> -				writebuf = filebuf;
> -			}
> -
> -			baderaseblock = false;
> -
> -			do {
> -				ret = mtd_is_bad(mtd, fd, offs / mtd->eb_size);
> -				if (ret < 0) {
> -					ERROR("mtd%d: MTD get bad block failed", mtdnum);
> -					goto closeall;
> -				} else if (ret == 1) {
> -					baderaseblock = true;
> -				}
> -
> -				if (baderaseblock) {
> -					mtdoffset = blockstart + mtd->eb_size;
> -
> -					if (mtdoffset > mtd->size) {
> -						ERROR("too many bad blocks, cannot complete request");
> -						goto closeall;
> -					}
> -				}
> -
> -				offs +=  mtd->eb_size;
> -			} while (offs < blockstart + mtd->eb_size);
> +	if (!priv->is_nand) {
> +		ret = mtd_read(priv->mtd, priv->fdout, priv->eb, 0,
> +		               priv->readout_buf, priv->mtd->eb_size);
> +		if (ret)
> +			return MTD_ERROR("read");
> +		/* Check if already empty: */
> +		if (buffer_check_pattern(priv->readout_buf, priv->mtd->eb_size,
> +		                         EMPTY_BYTE)) {
> +			return 0;
>   		}
> +	}
>
> -		/* Read more data from the input if there isn't enough in the buffer */
> -		if (writebuf + mtd->min_io_size > filebuf + filebuf_len) {
> -			size_t readlen = mtd->min_io_size;
> -			size_t alreadyread = (filebuf + filebuf_len) - writebuf;
> -			size_t tinycnt = alreadyread;
> -			ssize_t cnt = 0;
> -
> -			while (tinycnt < readlen) {
> -				cnt = read(ifd, writebuf + tinycnt, readlen - tinycnt);
> -				if (cnt == 0) { /* EOF */
> -					break;
> -				} else if (cnt < 0) {
> -					ERROR("File I/O error on input");
> -					goto closeall;
> -				}
> -				tinycnt += cnt;
> -			}
> -
> -			/* No padding needed - we are done */
> -			if (tinycnt == 0) {
> -				imglen = 0;
> -				break;
> -			}
> -
> -			/* Padding */
> -			if (tinycnt < readlen) {
> -				erase_buffer(writebuf + tinycnt, readlen - tinycnt);
> -			}
> +	ret = erase_block(priv);
> +	if (ret) {
> +		switch (ret) {
> +		case -EIO:
> +			ret = mark_bad_block(priv);
> +			if (ret)
> +				return ret;
> +			return 1;
> +		default:
> +			return ret;
> +		}
> +	}
> +	return 0;
> +}
>
> -			filebuf_len += readlen - alreadyread;
> +/*
> + * Find a non-bad erase block (starting from priv->eb) we can write data into.
> + * Unlock the block and erase it if necessary.
> + * As a result update (if necessary) priv->eb to point to a block we can write
> + * data into.
> + */
> +static int prepare_new_erase_block(struct flash_priv *priv)
> +{
> +	for (; priv->eb < priv->eb_end; priv->eb++) {
> +		int ret = process_new_erase_block(priv);
> +		if (ret < 0)
> +			return ret;
> +		else if (ret == 0)
> +			return 0;
> +	}
> +	return too_many_bad_blocks(priv->mtdnum);
> +}
>
> -			imglen -= tinycnt - alreadyread;
> +/*
> + * A callback to be passed to copyimage().
> + * Write as much input data to flash as possible at this time.
> + * Skip existing and detect new bad blocks (mark them as bad) while writing.
> + * Store leftover data (if any) in private context (priv). The leftover data
> + * will be written into flash during next flash_write() executions.
> + * During last flash_write() execution all image data should be written to
> + * flash (the data should be padded with EMPTY_BYTE if necessary).
> + */
> +static int flash_write(void *out, const void *buf, size_t len)
> +{
> +	int ret;
> +	struct flash_priv *priv = (struct flash_priv *)out;
> +	const unsigned char **pbuf = (const unsigned char**)&buf;
> +
> +	while ((len > 0) || (priv->writebuf_offset < priv->filebuf_len)) {
> +		unsigned char *wbuf;
> +		int to_write;
> +
> +		assert(priv->eb <= priv->eb_end);
> +		if (priv->eb == priv->eb_end)
> +			return too_many_bad_blocks(priv->mtdnum);
> +
> +		ret = read_data(priv, pbuf, &len, &wbuf, &to_write);
> +		if (ret < 0) {
> +			/* Wait for more data to be written in next
> +			 * flash_write() call. */
> +			break;
> +		}
> +		/* We've got some data to be written to flash. */
>
> +		if (priv->writebuf_offset == 0) {
> +			/* Start of a new erase block. */
> +			ret = prepare_new_erase_block(priv);
> +			if (ret)
> +				return ret;
>   		}
> +		/* Now priv->eb points to a valid erased block we can write
> +		 * data into. */
>
> -		ret =0;
> -		if (!buffer_check_pattern(writebuf, mtd->min_io_size, 0xff)) {
> -			/* Write out data */
> -			ret = mtd_write(flash->libmtd, mtd, fd, mtdoffset / mtd->eb_size,
> -					mtdoffset % mtd->eb_size,
> -					writebuf,
> -					mtd->min_io_size,
> -					NULL,
> -					0,
> -					MTD_OPS_PLACE_OOB);
> +		ret = buffer_check_pattern(wbuf, to_write, EMPTY_BYTE);
> +		if (ret) {
> +			ret = 0; /* There is no need to write "empty" bytes. */
> +		} else {
> +			/* Write data to flash: */
> +			ret = mtd_write(priv->libmtd, priv->mtd, priv->fdout,
> +			                priv->eb, priv->writebuf_offset, wbuf,
> +			                to_write, NULL, 0, MTD_OPS_PLACE_OOB);
>   		}
>   		if (ret) {
> -			long long i;
> -			if (errno != EIO) {
> -				ERROR("mtd%d: MTD write failure", mtdnum);
> -				goto closeall;
> -			}
> -
> -			/* Must rewind to blockstart if we can */
> -			writebuf = filebuf;
> -
> -			for (i = blockstart; i < blockstart + mtd->eb_size; i += mtd->eb_size) {
> -				if (mtd_erase(flash->libmtd, mtd, fd, i / mtd->eb_size)) {
> -					int errno_tmp = errno;
> -					TRACE("mtd%d: MTD Erase failure", mtdnum);
> -					if (errno_tmp != EIO)
> -						goto closeall;
> -				}
> -			}
> +			if (errno != EIO)
> +				return MTD_ERROR("write");
>
> -			TRACE("Marking block at %08llx bad",
> -					mtdoffset & (~mtd->eb_size + 1));
> -			if (mtd_mark_bad(mtd, fd, mtdoffset / mtd->eb_size)) {
> -				ERROR("mtd%d: MTD Mark bad block failure", mtdnum);
> -				goto closeall;
> +			ret = erase_block(priv);
> +			if (ret) {
> +				if (ret != -EIO)
> +					return ret;
>   			}
> -			mtdoffset = blockstart + mtd->eb_size;
>
> +			ret = mark_bad_block(priv);
> +			if (ret)
> +				return ret;
> +			/* Rewind to erase block start. */
> +			priv->writebuf_offset = 0;
> +			priv->eb++;
>   			continue;
>   		}
>
> -		/*
> -		 * this handler does not use copyfile()
> -		 * and must update itself the progress bar
> -		 */
> -		swupdate_progress_update((img->size - imglen) * 100 / img->size);
> -
> -		mtdoffset += mtd->min_io_size;
> -		writebuf += pagelen;
> -	}
> -	failed = false;
> -
> -closeall:
> -	free(filebuf);
> -	close(fd);
> -
> -	if (failed) {
> -		ERROR("Installing image %s into mtd%d failed",
> -			img->fname,
> -			mtdnum);
> -		return -1;
> +		priv->writebuf_offset += to_write;
> +		assert(priv->writebuf_offset <= priv->filebuf_len);
> +		assert(priv->filebuf_len <= priv->mtd->eb_size);
> +		if (priv->writebuf_offset == priv->mtd->eb_size) {
> +			priv->eb++;
> +			priv->writebuf_offset = 0;
> +			priv->filebuf_len = 0;
> +		}
>   	}
>
>   	return 0;
>   }
>
> -static int flash_write_nor(int mtdnum, struct img_type *img)
> +static int flash_write_image(int mtdnum, struct img_type *img)
>   {
> -	int fdout;
> +	struct flash_priv priv;
>   	char mtd_device[LINESIZE];
>   	int ret;
>   	struct flash_description *flash = get_flash_info();
> +	priv.mtd = &flash->mtd_info[mtdnum].mtd;
> +	assert((priv.mtd->eb_size % priv.mtd->min_io_size) == 0);
>
> -	if  (!mtd_dev_present(flash->libmtd, mtdnum)) {
> +	if (!mtd_dev_present(flash->libmtd, mtdnum)) {
>   		ERROR("MTD %d does not exist", mtdnum);
>   		return -ENODEV;
>   	}
>
> -	long long size = get_output_size(img, true);
> -	if (size < 0) {
> -		size = get_mtd_size(mtdnum);
> -		if (size < 0) {
> -			ERROR("Could not get MTD %d device size", mtdnum);
> -			return -ENODEV;
> -		}
> +	if (img->seek & (priv.mtd->eb_size - 1)) {
> +		ERROR("The start address is not erase-block-aligned!\n"
> +		      "The erase block of this flash is 0x%x.\n",
> +		      priv.mtd->eb_size);
> +		return -EINVAL;
> +	}
>
> -		WARN("decompression-size not set, erasing flash device %s from %lld to %lld",
> -			img->device, img->seek, size);
> +	priv.imglen = get_output_size(img, true);
> +	if (priv.imglen < 0) {
> +		ERROR("Failed to determine output size, bailing out.");
> +		return (int)priv.imglen;
>   	}
> -	if (flash_erase_sector(mtdnum, img->seek, size)) {
> -		ERROR("Failed to erase sectors on /dev/mtd%d (start: %llu, size: %lld)",
> -			mtdnum, img->seek, size);
> -		return -1;
> +	if (!priv.imglen)
> +		return 0;
> +
> +	if (priv.imglen > priv.mtd->size - img->seek) {
> +		ERROR("Image %s does not fit into mtd%d", img->fname, mtdnum);
> +		return -ENOSPC;
>   	}
>
>   	snprintf(mtd_device, sizeof(mtd_device), "/dev/mtd%d", mtdnum);
> -	if ((fdout = open(mtd_device, O_RDWR)) < 0) {
> -		ERROR( "%s: %s: %s", __func__, mtd_device, strerror(errno));
> -		return -1;
> +	priv.fdout = open(mtd_device, O_RDWR);
> +	if (priv.fdout < 0) {
> +		ret = errno;
> +		ERROR( "%s: %s: %s", __func__, mtd_device, STRERROR(ret));
> +		return -ret;
>   	}
>
> -	ret = copyimage(&fdout, img, NULL);
> -	close(fdout);
> +	priv.mtdnum = mtdnum;
> +	priv.filebuf_len = 0;
> +	priv.writebuf_offset = 0;
> +	priv.eb = (int)(img->seek / priv.mtd->eb_size);
> +	priv.eb_end = (int)(priv.mtd->size / priv.mtd->eb_size);
> +	priv.libmtd = flash->libmtd;
> +	priv.is_nand = isNand(flash, mtdnum);
> +	priv.check_bad = true;
> +	priv.check_locked = true;
> +
> +	ret = priv.mtd->eb_size;
> +	if (!priv.is_nand)
> +		ret += priv.mtd->eb_size;
> +	priv.filebuf = malloc(ret);
> +	if (!priv.filebuf) {
> +		ERROR("Failed to allocate %d bytes of memory", ret);
> +		ret = -ENOMEM;
> +		goto end;
> +	}
> +	if (!priv.is_nand)
> +		priv.readout_buf = priv.filebuf + priv.mtd->eb_size;
> +
> +	ret = copyimage(&priv, img, flash_write);
> +	free(priv.filebuf);
> +
> +end:
> +	if (close(priv.fdout)) {
> +		if (!ret)
> +			ret = -errno;
> +		ERROR("close() failed: %s", STRERROR(errno));
> +	}
>
>   	/* tell 'nbytes == 0' (EOF) from 'nbytes < 0' (read error) */
>   	if (ret < 0) {
>   		ERROR("Failure installing into: %s", img->device);
> -		return -1;
> +		return ret;
>   	}
>   	return 0;
>   }
>
> -static int flash_write_image(int mtdnum, struct img_type *img)
> -{
> -	struct flash_description *flash = get_flash_info();
> -
> -	if (!isNand(flash, mtdnum))
> -		return flash_write_nor(mtdnum, img);
> -	else
> -		return flash_write_nand(mtdnum, img);
> -}
> -
>   static int install_flash_image(struct img_type *img,
>   	void __attribute__ ((__unused__)) *data)
>   {
> -	int mtdnum;
> +	int ret, mtdnum;
>
>   	if (strlen(img->mtdname))
>   		mtdnum = get_mtd_from_name(img->mtdname);
> @@ -364,15 +561,16 @@ static int install_flash_image(struct img_type *img,
>   	if (mtdnum < 0) {
>   		ERROR("Wrong MTD device in description: %s",
>   			strlen(img->mtdname) ? img->mtdname : img->device);
> -		return -1;
> +		return -EINVAL;
>   	}
>
>   	TRACE("Copying %s into /dev/mtd%d", img->fname, mtdnum);
> -	if (flash_write_image(mtdnum, img)) {
> +	ret = flash_write_image(mtdnum, img);
> +	if (ret) {
>   		ERROR("I cannot copy %s into %s partition",
>   			img->fname,
>   			img->device);
> -		return -1;
> +		return ret;
>   	}
>
>   	return 0;
> diff --git a/test/Makefile b/test/Makefile
> index ad26b40d..6f3da093 100644
> --- a/test/Makefile
> +++ b/test/Makefile
> @@ -27,6 +27,7 @@ endif
>   tests-$(CONFIG_SURICATTA_HAWKBIT) += test_json
>   tests-$(CONFIG_SURICATTA_HAWKBIT) += test_server_hawkbit
>   tests-y += test_util
> +tests-$(CONFIG_CFI) += test_flash_handler
>
>   ccflags-y += -I$(src)/../
>
> diff --git a/test/test_flash_handler.c b/test/test_flash_handler.c
> new file mode 100644
> index 00000000..dd4a0b20
> --- /dev/null
> +++ b/test/test_flash_handler.c
> @@ -0,0 +1,1456 @@
> +/*
> + * Author: Viacheslav Volkov
> + *
> + * SPDX-License-Identifier: GPL-2.0-or-later
> + */
> +
> +#include <stddef.h>
> +#include <stdarg.h>
> +#include <setjmp.h>
> +#include <cmocka.h>
> +
> +#include <stdint.h>
> +#include <libmtd.h>
> +#include <mtd/mtd-user.h>
> +
> +#include <stdbool.h>
> +#include <stdio.h>
> +#include <string.h>
> +#include <fcntl.h>
> +#include <unistd.h>
> +#include <errno.h>
> +#include <limits.h>
> +
> +#include "util.h"
> +#include "pctl.h"
> +#include "flash.h"
> +#include "handler.h"
> +#include "swupdate_image.h"
> +#include "swupdate_status.h"
> +
> +/* #define PRINT_LOGS */
> +#define RETURN_CODE_ENOSPC
> +#define EMPTY_BYTE 0xFF
> +#define DEV_MTD_PREFIX "/dev/mtd"
> +#define MTD_DEV_IDX 0
> +/* MTD_FD is a big dummy value: it is very unlikely for us to have so many open
> + * file descriptors: */
> +#define MTD_FD 999999
> +#define LIBMTD_T_VALUE ((libmtd_t)0x12345678) /* dummy value */
> +
> +#define FOREACH_PAGE_IN_EB(page, mtd, eb) \
> +	for (page = (eb) * ((mtd)->eb_size / (mtd)->min_io_size); \
> +	     page < (eb + 1) * ((mtd)->eb_size / (mtd)->min_io_size); \
> +	     page++)
> +
> +#define FOREACH_PAGE_WRITTEN(page, mtd, eb, offs, len) \
> +	for (page = ((eb) * (mtd)->eb_size + (offs)) / (mtd)->min_io_size; \
> +	page < ((eb) * (mtd)->eb_size + (offs) + (len)) / (mtd)->min_io_size; \
> +	     page++)
> +
> +/*
> + * Mock data
> + */
> +static struct img_type image = {
> +	.type = "flash",
> +};
> +static struct mtd_ubi_info mtd_ubi_info_s;
> +
> +/*
> + * Test-only data
> + */
> +static handler handler_func;
> +static bool is_nand;
> +static unsigned char *image_buf;
> +static void (*patch_image_buf)(unsigned char *buf);
> +static int eb_bytes;
> +static int pages_bytes;
> +/* Actual flash state before test run: */
> +static unsigned char *bad_blocks;
> +static unsigned char *locked_blocks;
> +static unsigned char *written_pages;
> +static unsigned char *flash_memory;
> +/* Expected flash state after test run: */
> +static unsigned char *expected_bad_blocks;
> +static unsigned char *expected_locked_blocks;
> +static unsigned char *expected_written_pages;
> +static unsigned char *expected_flash_memory;
> +
> +void *__real_malloc(size_t size);
> +static void *(*impl_malloc)(size_t size) = __real_malloc;
> +void *__wrap_malloc(size_t size);
> +void *__wrap_malloc(size_t size)
> +{
> +	return impl_malloc(size);
> +}
> +
> +int __real_open(const char *pathname, int flags);
> +static int default_open(const char *pathname, int flags)
> +{
> +	int ret;
> +	if (strcmp(pathname, image.device) == 0)
> +		return MTD_FD;
> +	ret = __real_open(pathname, flags);
> +	/* Usually file descriptors start from 0. Since MTD_FD is big enough,
> +	 * a collision is very unlikely. */
> +	assert_int_not_equal(ret, MTD_FD);
> +	return ret;
> +}
> +static int (*impl_open)(const char *pathname, int flags) = default_open;
> +int __wrap_open(const char *pathname, int flags);
> +int __wrap_open(const char *pathname, int flags)
> +{
> +	return impl_open(pathname, flags);
> +}
> +
> +int __real_close(int fd);
> +int __wrap_close(int fd);
> +int __wrap_close(int fd)
> +{
> +	if (fd == MTD_FD)
> +		return 0;
> +	return __real_close(fd);
> +}
> +
> +off_t __real_lseek(int fd, off_t offset, int whence);
> +off_t __wrap_lseek(int fd, off_t offset, int whence);
> +off_t __wrap_lseek(int fd, off_t offset, int whence)
> +{
> +	if (fd == MTD_FD) {
> +		/* __swupdate_copy() executes lseek() on /dev/mtdX. */
> +		assert_int_equal(image.seek, offset);
> +		assert_int_equal(SEEK_SET, whence);
> +		return offset;
> +	}
> +	return __real_lseek(fd, offset, whence);
> +}
> +
> +int __wrap_mtd_dev_present(libmtd_t desc, int mtd_num);
> +int __wrap_mtd_dev_present(libmtd_t UNUSED desc, int mtd_num)
> +{
> +	assert_int_equal(MTD_DEV_IDX, mtd_num);
> +	return 1;
> +}
> +
> +int __wrap_get_mtd_from_device(char *s);
> +int __wrap_get_mtd_from_device(char *s)
> +{
> +	int ret, mtdnum;
> +	assert_ptr_not_equal(NULL, s);
> +	if (strlen(s) < sizeof(DEV_MTD_PREFIX))
> +		return -1;
> +	ret = sscanf(s, DEV_MTD_PREFIX "%d", &mtdnum);
> +	if (ret <= 0)
> +		return -1;
> +	return mtdnum;
> +}
> +
> +static unsigned char get_byte_idx(int bit_idx)
> +{
> +	return bit_idx / CHAR_BIT;
> +}
> +
> +static unsigned char get_mask(int bit_idx)
> +{
> +	return 1 << (bit_idx % CHAR_BIT);
> +}
> +
> +static bool get_bit(unsigned char *data, int bit_idx)
> +{
> +	return (data[get_byte_idx(bit_idx)] & get_mask(bit_idx)) != 0;
> +}
> +
> +static void set_bit(unsigned char *data, int bit_idx)
> +{
> +	data[get_byte_idx(bit_idx)] |= get_mask(bit_idx);
> +}
> +
> +static void clear_bit(unsigned char *data, int bit_idx)
> +{
> +	data[get_byte_idx(bit_idx)] &= (~get_mask(bit_idx));
> +}
> +
> +static void set_multiple_bits(unsigned char *data, const int *bit_indices)
> +{
> +	int bit_idx = bit_indices[0];
> +	for (int i = 0; bit_idx >= 0; bit_idx = bit_indices[++i])
> +		set_bit(data, bit_idx);
> +}
> +
> +static void clear_multiple_bits(unsigned char *data, const int *bit_indices)
> +{
> +	int bit_idx = bit_indices[0];
> +	for (int i = 0; bit_idx >= 0; bit_idx = bit_indices[++i])
> +		clear_bit(data, bit_idx);
> +}
> +
> +static void check_args(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	assert_ptr_equal(&mtd_ubi_info_s.mtd, mtd);
> +	assert_true(
> +	            (mtd->type == MTD_NORFLASH) ||
> +	            (mtd->type == MTD_NANDFLASH));
> +	assert_int_equal(MTD_FD, fd);
> +	assert_true(eb >= 0);
> +	assert_true(eb < mtd->size / mtd->eb_size);
> +}
> +
> +static int default_mtd_get_bool(const struct mtd_dev_info UNUSED *mtd,
> +		int UNUSED fd, int eb, unsigned char *data)
> +{
> +	return get_bit(data, eb) ? 1 : 0;
> +}
> +
> +static int default_mtd_is_bad(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	return default_mtd_get_bool(mtd, fd, eb, bad_blocks);
> +}
> +static int (*impl_mtd_is_bad)(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_is_bad(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_is_bad(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	return impl_mtd_is_bad(mtd, fd, eb);
> +}
> +
> +static int default_mtd_mark_bad(const struct mtd_dev_info UNUSED *mtd,
> +		int UNUSED fd, int eb)
> +{
> +	set_bit(bad_blocks, eb);
> +	return 0;
> +}
> +static int (*impl_mtd_mark_bad)(const struct mtd_dev_info *mtd, int fd,
> +		int eb);
> +int __wrap_mtd_mark_bad(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_mark_bad(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	assert_false(get_bit(bad_blocks, eb));
> +	/* Note: we don't require erasing a block before marking it bad. */
> +	return impl_mtd_mark_bad(mtd, fd, eb);
> +}
> +
> +static int default_mtd_is_locked(const struct mtd_dev_info *mtd, int fd,
> +		int eb)
> +{
> +	return default_mtd_get_bool(mtd, fd, eb, locked_blocks);
> +}
> +static int (*impl_mtd_is_locked)(const struct mtd_dev_info *mtd, int fd,
> +		int eb);
> +int __wrap_mtd_is_locked(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_is_locked(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	return impl_mtd_is_locked(mtd, fd, eb);
> +}
> +
> +static int default_mtd_unlock(const struct mtd_dev_info UNUSED *mtd,
> +		int UNUSED fd, int eb)
> +{
> +	clear_bit(locked_blocks, eb);
> +	return 0;
> +}
> +static int (*impl_mtd_unlock)(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_unlock(const struct mtd_dev_info *mtd, int fd, int eb);
> +int __wrap_mtd_unlock(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	assert_false(get_bit(bad_blocks, eb));
> +	/* Unlocking already unlocked blocks is totaly fine:
> +	 * - not every flash supports mtd_is_locked() check;
> +	 * - some implementation might prefer to mtd_unlock() rigth away
> +	 *   without checking mtd_is_locked(). */
> +	return impl_mtd_unlock(mtd, fd, eb);
> +}
> +
> +static int default_mtd_erase(libmtd_t UNUSED desc,
> +		const struct mtd_dev_info *mtd, int UNUSED fd, int eb)
> +{
> +	unsigned int page;
> +	memset(flash_memory + eb * mtd->eb_size, EMPTY_BYTE, mtd->eb_size);
> +	FOREACH_PAGE_IN_EB(page, mtd, eb)
> +		clear_bit(written_pages, page);
> +	return 0;
> +}
> +static int (*impl_mtd_erase)(libmtd_t desc, const struct mtd_dev_info *mtd,
> +		int fd, int eb);
> +int __wrap_mtd_erase(libmtd_t desc, const struct mtd_dev_info *mtd, int fd,
> +		int eb);
> +int __wrap_mtd_erase(libmtd_t desc, const struct mtd_dev_info *mtd, int fd,
> +		int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	assert_false(get_bit(bad_blocks, eb));
> +	assert_false(get_bit(locked_blocks, eb));
> +	return impl_mtd_erase(desc, mtd, fd, eb);
> +}
> +
> +static int default_mtd_write(libmtd_t UNUSED desc,
> +		const struct mtd_dev_info *mtd, int UNUSED fd, int eb,
> +		int offs, void *data, int len, void UNUSED *oob,
> +		int UNUSED ooblen, uint8_t UNUSED mode)
> +{
> +	unsigned int page;
> +	unsigned int flash_offset = eb * mtd->eb_size + offs;
> +	const unsigned char *pdata = (const unsigned char*)data;
> +	FOREACH_PAGE_WRITTEN(page, mtd, eb, offs, len) {
> +		set_bit(written_pages, page);
> +		memcpy(flash_memory + flash_offset, pdata, mtd->min_io_size);
> +		flash_offset += mtd->min_io_size;
> +		pdata += mtd->min_io_size;
> +	}
> +	return 0;
> +}
> +
> +static int (*impl_mtd_write)(libmtd_t desc, const struct mtd_dev_info *mtd,
> +		int fd, int eb, int offs, void *data, int len, void *oob,
> +		int ooblen, uint8_t mode);
> +int __wrap_mtd_write(libmtd_t desc, const struct mtd_dev_info *mtd, int fd,
> +		int eb, int offs, void *data, int len, void *oob, int ooblen,
> +		uint8_t mode);
> +
> +int __wrap_mtd_write(libmtd_t desc, const struct mtd_dev_info *mtd, int fd,
> +		int eb, int offs, void *data, int len, void *oob, int ooblen,
> +		uint8_t mode)
> +{
> +	unsigned int page;
> +	check_args(mtd, fd, eb);
> +	assert_true(offs >= 0);
> +	assert_ptr_not_equal(NULL, data);
> +	assert_true(len > 0);
> +	assert_true(offs + len <= mtd->eb_size);
> +	assert_ptr_equal(NULL, oob);
> +	assert_int_equal(0, ooblen);
> +	assert_false(get_bit(bad_blocks, eb));
> +	assert_false(get_bit(locked_blocks, eb));
> +	assert_true(offs <= mtd->eb_size - mtd->min_io_size);
> +	assert_int_equal(0, offs % mtd->min_io_size);
> +	assert_int_equal(0, len % mtd->min_io_size);
> +
> +	if (mtd->type == MTD_NANDFLASH) {
> +		/* Follow "write once rule" for NAND flash. */
> +		FOREACH_PAGE_WRITTEN(page, mtd, eb, offs, len)
> +			assert_false(get_bit(written_pages, page));
> +	} else {
> +		/* Assume it is ok to write multiple times to NOR flash. */
> +		unsigned int flash_i = eb * mtd->eb_size + offs;
> +		for (int data_i = 0; data_i < len; data_i++) {
> +			uint8_t old_byte = flash_memory[flash_i + data_i];
> +			uint8_t new_byte = ((uint8_t*)data)[data_i];
> +			assert_true((old_byte & new_byte) == new_byte);
> +		}
> +	}
> +	return impl_mtd_write(desc, mtd, fd, eb, offs, data, len, oob, ooblen,
> +	                      mode);
> +}
> +
> +static int default_mtd_read(const struct mtd_dev_info *mtd, int fd, int eb,
> +		int offs, void *buf, int len)
> +{
> +	unsigned int flash_offset = eb * mtd->eb_size + offs;
> +	check_args(mtd, fd, eb);
> +	memcpy(buf, flash_memory + flash_offset, len);
> +	return 0;
> +}
> +static int (*impl_mtd_read)(const struct mtd_dev_info *mtd, int fd, int eb,
> +		int offs, void *buf, int len);
> +int __wrap_mtd_read(const struct mtd_dev_info *mtd, int fd, int eb, int offs,
> +		void *buf, int len);
> +int __wrap_mtd_read(const struct mtd_dev_info *mtd, int fd, int eb, int offs,
> +		void *buf, int len)
> +{
> +	check_args(mtd, fd, eb);
> +	assert_true(offs >= 0);
> +	assert_ptr_not_equal(NULL, buf);
> +	assert_true(len > 0);
> +	assert_true(offs + len <= mtd->eb_size);
> +	assert_false(get_bit(bad_blocks, eb));
> +	assert_false(get_bit(locked_blocks, eb));
> +	assert_true(offs <= mtd->eb_size - mtd->min_io_size);
> +	assert_int_equal(0, offs % mtd->min_io_size);
> +	assert_int_equal(0, len % mtd->min_io_size);
> +	return impl_mtd_read(mtd, fd, eb, offs, buf, len);
> +}
> +
> +static unsigned char *generate_image_file(const char *name, size_t size)
> +{
> +	int r, fd;
> +	unsigned char *buf = malloc(size);
> +	unsigned char *pbuf = buf;
> +	assert_ptr_not_equal(NULL, buf);
> +	fd = open(name, O_WRONLY);
> +	if (fd < 0)
> +		free(buf);
> +	assert_true(fd >= 0);
> +
> +	for (size_t i = 0; i < size; i++)
> +		buf[i] = (unsigned char)i;
> +
> +	if (patch_image_buf)
> +		patch_image_buf(buf);
> +
> +	while (size > 0) {
> +		ssize_t ret = write(fd, pbuf, size);
> +		if (ret < 0) {
> +			close(fd);
> +			free(buf);
> +		}
> +		assert_true(ret >= 0);
> +		pbuf += ret;
> +		size -= ret;
> +	}
> +
> +	r = close(fd);
> +	assert_int_equal(0, r);
> +	return buf;
> +}
> +
> +static int group_setup(void UNUSED **state)
> +{
> +	int ret, fd;
> +	struct flash_description *flash = get_flash_info();
> +	struct installer_handler *hnd = find_handler(&image);
> +
> +	assert_ptr_not_equal(NULL, hnd);
> +	handler_func = hnd->installer;
> +	assert_ptr_not_equal(NULL, handler_func);
> +	flash->mtd_info = &mtd_ubi_info_s;
> +	flash->libmtd = LIBMTD_T_VALUE;
> +
> +#ifdef PRINT_LOGS
> +	loglevel = DEBUGLEVEL;
> +	notify_init();
> +#endif
> +
> +	strcpy(image.fname, "swupdate_image_XXXXXX.bin");
> +	fd = mkstemps(image.fname, 4 /* = strlen(".bin") */);
> +	if (fd < 0)
> +		image.fname[0] = 0; /* required for group_teardown() */
> +	ret = close(fd);
> +	assert_int_equal(0, ret);
> +	return 0;
> +}
> +
> +static int group_teardown(void UNUSED **state)
> +{
> +	if (image.fname[0] != 0) {
> +		int ret = unlink(image.fname);
> +		assert_int_equal(0, ret);
> +	}
> +	return 0;
> +}
> +
> +static void copy_flash_state(void)
> +{
> +	struct mtd_dev_info *mtd = &mtd_ubi_info_s.mtd;
> +
> +	expected_bad_blocks = malloc(eb_bytes);
> +	assert_ptr_not_equal(NULL, expected_bad_blocks);
> +	memcpy(expected_bad_blocks, bad_blocks, eb_bytes);
> +
> +	expected_locked_blocks = malloc(eb_bytes);
> +	assert_ptr_not_equal(NULL, expected_locked_blocks);
> +	memcpy(expected_locked_blocks, locked_blocks, eb_bytes);
> +
> +	expected_written_pages = malloc(pages_bytes);
> +	assert_ptr_not_equal(NULL, expected_written_pages);
> +	memcpy(expected_written_pages, written_pages, pages_bytes);
> +
> +	expected_flash_memory = malloc(mtd->size);
> +	assert_ptr_not_equal(NULL, expected_flash_memory);
> +	memcpy(expected_flash_memory, flash_memory, mtd->size);
> +}
> +
> +static void init_flash_state(void)
> +{
> +	struct mtd_dev_info *mtd = &mtd_ubi_info_s.mtd;
> +	int num = (int)(mtd->size / mtd->eb_size);
> +	eb_bytes = num / CHAR_BIT;
> +	if (num % CHAR_BIT)
> +		eb_bytes++;
> +	/* By default there are no any bad blocks: */
> +	bad_blocks = calloc(eb_bytes, 1);
> +	assert_ptr_not_equal(NULL, bad_blocks);
> +	/* By default all blocks are locked: */
> +	locked_blocks = malloc(eb_bytes);
> +	assert_ptr_not_equal(NULL, locked_blocks);
> +	memset(locked_blocks, 0xFF, eb_bytes);
> +
> +	num = (int)(mtd->size / mtd->min_io_size);
> +	pages_bytes = num / CHAR_BIT;
> +	if (num % CHAR_BIT)
> +		pages_bytes++;
> +	/* By default all pages are written (require erase): */
> +	written_pages = malloc(pages_bytes);
> +	assert_ptr_not_equal(NULL, written_pages);
> +	memset(written_pages, 0xFF, pages_bytes);
> +
> +	flash_memory = malloc(mtd->size);
> +	assert_ptr_not_equal(NULL, flash_memory);
> +	/* Fill flash memory with initial pattern: */
> +	memset(flash_memory, 0xA5, mtd->size);
> +}
> +
> +static void test_init(void)
> +{
> +	struct mtd_dev_info *mtd = &mtd_ubi_info_s.mtd;
> +
> +	/* Unit test config (precondition) sanity check: */
> +	assert_true(mtd->size >= mtd->eb_size);
> +	assert_true(mtd->eb_size >= mtd->min_io_size);
> +	assert_int_equal(0, mtd->size % mtd->eb_size);
> +	assert_int_equal(0, mtd->eb_size % mtd->min_io_size);
> +
> +	is_nand = (mtd->type == MTD_NANDFLASH) ||
> +	          (mtd->type == MTD_MLCNANDFLASH);
> +	image_buf = generate_image_file(image.fname, image.size);
> +	image.fdin = open(image.fname, O_RDONLY);
> +	assert_true(image.fdin >= 0);
> +
> +	init_flash_state();
> +}
> +
> +static void verify_flash_state(void)
> +{
> +	struct mtd_dev_info *mtd = &mtd_ubi_info_s.mtd;
> +	assert_memory_equal(expected_bad_blocks, bad_blocks, eb_bytes);
> +	assert_memory_equal(expected_locked_blocks, locked_blocks, eb_bytes);
> +	assert_memory_equal(expected_written_pages, written_pages,
> +	                    pages_bytes);
> +	assert_memory_equal(expected_flash_memory, flash_memory, mtd->size);
> +}
> +
> +static void run_flash_test(void UNUSED **state, int expected_return_code)
> +{
> +	int ret = handler_func(&image, NULL);
> +#ifdef RETURN_CODE_ENOSPC
> +	/* Currently on callback() failure __swupdate_copy() always returns
> +	 * -ENOSPC (no matter what callback() returned). Hence some subset of
> +	 * flash handler return values are converted to -ENOSPC. Here we'll
> +	 * convert rest of return values to -ENOSPC to map expected and actual
> +	 * return values to a simple success/failure check. */
> +	if (ret < 0)
> +		ret = -ENOSPC;
> +	if (expected_return_code < 0)
> +		expected_return_code = -ENOSPC;
> +#endif
> +	assert_int_equal(expected_return_code, ret);
> +	verify_flash_state();
> +}
> +
> +static int test_setup(void UNUSED **state)
> +{
> +	struct mtd_dev_info *mtd = &mtd_ubi_info_s.mtd;
> +	image_buf = NULL;
> +	patch_image_buf = NULL;
> +	image.fdin = -1;
> +
> +	bad_blocks = NULL;
> +	locked_blocks = NULL;
> +	written_pages = NULL;
> +	flash_memory = NULL;
> +
> +	expected_bad_blocks = NULL;
> +	expected_locked_blocks = NULL;
> +	expected_written_pages = NULL;
> +	expected_flash_memory = NULL;
> +
> +	/* Tests can overwrite the following default implementations to inject
> +	 * errors: */
> +	impl_malloc = __real_malloc;
> +	impl_open = default_open;
> +	impl_mtd_is_bad = default_mtd_is_bad;
> +	impl_mtd_mark_bad = default_mtd_mark_bad;
> +	impl_mtd_is_locked = default_mtd_is_locked;
> +	impl_mtd_unlock = default_mtd_unlock;
> +	impl_mtd_erase = default_mtd_erase;
> +	impl_mtd_write = default_mtd_write;
> +	impl_mtd_read = default_mtd_read;
> +
> +	/* Some default values: */
> +	mtd->type = MTD_NANDFLASH;
> +	strcpy(image.device, DEV_MTD_PREFIX STR(MTD_DEV_IDX));
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	return 0;
> +}
> +
> +static int test_teardown(void UNUSED **state)
> +{
> +	unsigned char *pointer[] = {
> +		bad_blocks,
> +		locked_blocks,
> +		written_pages,
> +		flash_memory,
> +
> +		expected_bad_blocks,
> +		expected_locked_blocks,
> +		expected_written_pages,
> +		expected_flash_memory,
> +
> +		image_buf,
> +	};
> +	impl_malloc = __real_malloc;
> +	for (int i = 0; i < ARRAY_SIZE(pointer); i++)
> +		free(pointer[i]);
> +	if (image.fdin >= 0) {
> +		int ret = close(image.fdin);
> +		assert_int_equal(0, ret);
> +	}
> +	return 0;
> +}
> +
> +static void test_simple(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_simple_NOR(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 1;
> +	mtd_ubi_info_s.mtd.type = MTD_NORFLASH;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_padding_less_than_page(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 42;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	memset(expected_flash_memory + 42, EMPTY_BYTE, 6);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_padding_page(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 40;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	memset(expected_flash_memory + 40, EMPTY_BYTE, 8);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 5);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void patch_image_buf_empty_bytes_1(unsigned char *buf)
> +{
> +	memset(buf + 8, EMPTY_BYTE, 24);
> +}
> +
> +static void test_skip_write_page_empty_bytes(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 40;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	patch_image_buf = patch_image_buf_empty_bytes_1;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	memset(expected_flash_memory + 40, EMPTY_BYTE, 8);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	if (is_nand)
> +		clear_bit(expected_written_pages, 1);
> +	clear_bit(expected_written_pages, 2);
> +	clear_bit(expected_written_pages, 3);
> +	clear_bit(expected_written_pages, 5);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_padding_more_than_page(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 37;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	memset(expected_flash_memory + 37, EMPTY_BYTE, 11);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 5);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_seek(void **state)
> +{
> +	image.seek = 16;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory + 16, image_buf, image.size);
> +	for (int i = 1; i <= 3; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_seek_not_multiple_of_eb_size(void **state)
> +{
> +	image.seek = 8;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	test_init();
> +	copy_flash_state();
> +	run_flash_test(state, -EINVAL);
> +}
> +
> +static void test_not_enough_flash(void **state)
> +{
> +	image.seek = 16;
> +	image.size = 1020;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	test_init();
> +	copy_flash_state();
> +	run_flash_test(state, -ENOSPC);
> +}
> +
> +static void test_invalid_mtd_device(void **state)
> +{
> +	test_init();
> +	strcpy(image.device, "/dev/mtdX");
> +	copy_flash_state();
> +	run_flash_test(state, -EINVAL);
> +}
> +
> +static void test_invalid_image_size(void **state)
> +{
> +	test_init();
> +	image.size = -42;
> +	copy_flash_state();
> +	run_flash_test(state, -42);
> +}
> +
> +static void test_empty_image(void **state)
> +{
> +	test_init();
> +	image.size = 0;
> +	copy_flash_state();
> +	run_flash_test(state, 0);
> +}
> +
> +static int open_mtd_dev_failure_errno;
> +static int open_mtd_dev_failure(const char *pathname, int flags)
> +{
> +	if (strcmp(pathname, image.device) == 0) {
> +		errno = open_mtd_dev_failure_errno;
> +		return -1;
> +	}
> +	return default_open(pathname, flags);
> +}
> +static void test_mtd_dev_open_failure(void **state)
> +{
> +	test_init();
> +	impl_open = open_mtd_dev_failure;
> +	open_mtd_dev_failure_errno = EPERM;
> +	copy_flash_state();
> +	run_flash_test(state, -open_mtd_dev_failure_errno);
> +}
> +
> +static int malloc_filebuf_allocation_failure_size;
> +static void *malloc_filebuf_allocation_failure(size_t size)
> +{
> +	if (size == malloc_filebuf_allocation_failure_size) {
> +		errno = ENOMEM;
> +		return NULL;
> +	}
> +	return __real_malloc(size);
> +}
> +static void test_malloc_failure(void **state)
> +{
> +	test_init();
> +	copy_flash_state();
> +	impl_malloc = malloc_filebuf_allocation_failure;
> +	malloc_filebuf_allocation_failure_size = mtd_ubi_info_s.mtd.eb_size;
> +	if (!is_nand) {
> +		/* filebuf allocation includes space for readout buffer */
> +		malloc_filebuf_allocation_failure_size *= 2;
> +	}
> +	run_flash_test(state, -ENOMEM);
> +}
> +
> +static void test_skip_bad_blocks(void **state)
> +{
> +	image.seek = 896;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	/* Initial flash state:                    | flash memory ends here
> +	 * eb      56  57  58  59  60  61  62   63 |   64
> +	 * offset 896 912 928 944 960 976 992 1008 | 1024
> +	 * status bad  ok bad  ok bad bad  ok  bad | */
> +
> +	{
> +		int indices[] = {56, 58, 60, 61, 63, -1};
> +		set_multiple_bits(bad_blocks, indices);
> +	}
> +	copy_flash_state();
> +
> +	{
> +		int indices[] = {57, 59, 62, -1};
> +		clear_multiple_bits(expected_locked_blocks, indices);
> +	}
> +	memcpy(expected_flash_memory + 912, image_buf     , 16);
> +	memcpy(expected_flash_memory + 944, image_buf + 16, 16);
> +	memcpy(expected_flash_memory + 992, image_buf + 32, 16);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_too_many_known_bad_blocks(void **state)
> +{
> +	image.seek = 896;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	/* Initial flash state:                    | flash memory ends here
> +	 * eb      56  57  58  59  60  61  62   63 |   64
> +	 * offset 896 912 928 944 960 976 992 1008 | 1024
> +	 * status bad  ok bad  ok bad bad bad  bad | */
> +
> +	{
> +		int indices[] = {56, 58, 60, 61, 62, 63, -1};
> +		set_multiple_bits(bad_blocks, indices);
> +	}
> +	copy_flash_state();
> +
> +	{
> +		int indices[] = {57, 59, -1};
> +		clear_multiple_bits(expected_locked_blocks, indices);
> +	}
> +	memcpy(expected_flash_memory + 912, image_buf     , 16);
> +	memcpy(expected_flash_memory + 944, image_buf + 16, 16);
> +
> +	run_flash_test(state, -ENOSPC);
> +}
> +
> +static int mtd_write_failure_1_eb;
> +static int mtd_write_failure_1_offs;
> +static int mtd_write_failure_1_errno;
> +static int mtd_write_failure_1(libmtd_t desc, const struct mtd_dev_info *mtd,
> +		int fd, int eb, int offs, void *data, int len, void *oob,
> +		int ooblen, uint8_t mode)
> +{
> +	if (eb == mtd_write_failure_1_eb) {
> +		unsigned char *p_data = (unsigned char *)data;
> +		for (int o = offs; o < offs + len; o += mtd->min_io_size) {
> +			int ret;
> +			if (o == mtd_write_failure_1_offs) {
> +				errno = mtd_write_failure_1_errno;
> +				return -1;
> +			}
> +			ret = default_mtd_write(desc, mtd, fd, eb, o,
> +			                        p_data, mtd->min_io_size, oob,
> +			                        ooblen, mode);
> +			assert_int_equal(0, ret);
> +			p_data += mtd->min_io_size;
> +		}
> +	}
> +	return default_mtd_write(desc, mtd, fd, eb, offs, data, len, oob,
> +	                         ooblen, mode);
> +}
> +
> +static void test_too_many_unknown_bad_blocks(void **state)
> +{
> +	image.seek = 896;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_offs = 0;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	/* Initial flash state:                       | flash memory ends here
> +	 * eb      56  57  58  59  60  61  62      63 |   64
> +	 * offset 896 912 928 944 960 976 992    1008 | 1024
> +	 * status bad  ok bad  ok bad bad bad ok->bad | */
> +	{
> +		int indices[] = {56, 58, 60, 61, 62, -1};
> +		set_multiple_bits(bad_blocks, indices);
> +	}
> +	mtd_write_failure_1_eb = 63;
> +	copy_flash_state();
> +
> +	{
> +		int indices[] = {57, 59, 63, -1};
> +		clear_multiple_bits(expected_locked_blocks, indices);
> +	}
> +	memcpy(expected_flash_memory +  912, image_buf     , 16);
> +	memcpy(expected_flash_memory +  944, image_buf + 16, 16);
> +	memset(expected_flash_memory + 1008, EMPTY_BYTE, 16);
> +	set_bit(expected_bad_blocks, 63);
> +	clear_bit(expected_written_pages, 63 * 2);
> +	clear_bit(expected_written_pages, 63 * 2 + 1);
> +
> +	run_flash_test(state, -ENOSPC);
> +}
> +
> +static int mtd_is_bad_failure_1_errno;
> +static int mtd_is_bad_failure_1(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	errno = mtd_is_bad_failure_1_errno;
> +	return -1;
> +}
> +
> +static void test_mtd_is_bad_not_supported(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_is_bad = mtd_is_bad_failure_1;
> +	mtd_is_bad_failure_1_errno = EOPNOTSUPP;
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_is_bad_failure(void **state)
> +{
> +	test_init();
> +	impl_mtd_is_bad = mtd_is_bad_failure_1;
> +	mtd_is_bad_failure_1_errno = ERANGE; /* dummy value */
> +	copy_flash_state();
> +	run_flash_test(state, -mtd_is_bad_failure_1_errno);
> +}
> +
> +static int mtd_is_locked_failure_1_errno;
> +static int mtd_is_locked_failure_1(const struct mtd_dev_info *mtd, int fd,
> +		int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	errno = mtd_is_locked_failure_1_errno;
> +	return -1;
> +}
> +
> +static void test_mtd_is_locked_not_supported(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_is_locked = mtd_is_locked_failure_1;
> +	mtd_is_locked_failure_1_errno = EOPNOTSUPP;
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_is_locked_failure(void **state)
> +{
> +	test_init();
> +	impl_mtd_is_locked = mtd_is_locked_failure_1;
> +	mtd_is_locked_failure_1_errno = ERANGE; /* dummy value */
> +	copy_flash_state();
> +	run_flash_test(state, -mtd_is_locked_failure_1_errno);
> +}
> +
> +static int mtd_unlock_failure_1_errno;
> +static int mtd_unlock_failure_1(const struct mtd_dev_info *mtd, int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	errno = mtd_unlock_failure_1_errno;
> +	return -1;
> +}
> +
> +static void test_mtd_unlock_not_supported(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_unlock = mtd_unlock_failure_1;
> +	mtd_unlock_failure_1_errno = EOPNOTSUPP;
> +	memset(locked_blocks, 0, eb_bytes);
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_unlock_failure(void **state)
> +{
> +	test_init();
> +	impl_mtd_unlock = mtd_unlock_failure_1;
> +	mtd_unlock_failure_1_errno = ERANGE; /* dummy value */
> +	copy_flash_state();
> +	run_flash_test(state, -mtd_unlock_failure_1_errno);
> +}
> +
> +static int mtd_read_failure_1_errno;
> +static int mtd_read_failure_1(const struct mtd_dev_info *mtd, int fd, int eb,
> +		int UNUSED offs, void UNUSED *buf, int UNUSED len)
> +{
> +	check_args(mtd, fd, eb);
> +	errno = mtd_read_failure_1_errno;
> +	return -1;
> +}
> +
> +static void test_mtd_read_failure(void **state)
> +{
> +	mtd_ubi_info_s.mtd.type = MTD_NORFLASH;
> +	test_init();
> +	impl_mtd_read = mtd_read_failure_1;
> +	mtd_read_failure_1_errno = ERANGE; /* dummy value */
> +	copy_flash_state();
> +	clear_bit(expected_locked_blocks, 0);
> +	run_flash_test(state, -mtd_read_failure_1_errno);
> +}
> +
> +static int mtd_erase_failure_1_eb;
> +static int mtd_erase_failure_1_errno;
> +static int mtd_erase_failure_1(libmtd_t desc, const struct mtd_dev_info *mtd,
> +		int fd, int eb)
> +{
> +	if (eb == mtd_erase_failure_1_eb) {
> +		errno = mtd_erase_failure_1_errno;
> +		if (errno == EOPNOTSUPP)
> +			default_mtd_erase(desc, mtd, fd, eb);
> +		return -1;
> +	}
> +	return default_mtd_erase(desc, mtd, fd, eb);
> +}
> +
> +static void test_mtd_read_no_erase_empty_flash_bytes(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.type = MTD_NORFLASH;
> +	test_init();
> +	memset(flash_memory, EMPTY_BYTE, 16);
> +	impl_mtd_erase = mtd_erase_failure_1;
> +	mtd_erase_failure_1_eb = 0;
> +	mtd_erase_failure_1_errno = ERANGE; /* dummy value */
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_erase_not_supported(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_erase = mtd_erase_failure_1;
> +	mtd_erase_failure_1_eb = 0;
> +	mtd_erase_failure_1_errno = EOPNOTSUPP;
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_erase_failure(void **state)
> +{
> +	test_init();
> +	impl_mtd_erase = mtd_erase_failure_1;
> +	mtd_erase_failure_1_eb = 0;
> +	mtd_erase_failure_1_errno = ERANGE; /* dummy value */
> +	copy_flash_state();
> +	clear_bit(expected_locked_blocks, 0);
> +	run_flash_test(state, -mtd_erase_failure_1_errno);
> +}
> +
> +static void test_mtd_erase_failure_EIO(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_erase = mtd_erase_failure_1;
> +	mtd_erase_failure_1_eb = 0;
> +	mtd_erase_failure_1_errno = EIO;
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory + 16, image_buf, image.size);
> +	for (int i = 0; i <= 3; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	set_bit(expected_bad_blocks, 0);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void patch_image_buf_empty_bytes_2(unsigned char *buf)
> +{
> +	memset(buf, EMPTY_BYTE, 16);
> +}
> +
> +static void test_no_mtd_write_empty_image_bytes(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	patch_image_buf = patch_image_buf_empty_bytes_2;
> +	test_init();
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 0;
> +	mtd_write_failure_1_offs = 0;
> +	mtd_write_failure_1_errno = ERANGE; /* dummy value */
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	for (int i = 0; i <= 2; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 0);
> +	clear_bit(expected_written_pages, 1);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_write_failure(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = ERANGE; /* dummy value */
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, 24);
> +	memset(expected_flash_memory + 24, EMPTY_BYTE, 8);
> +	for (int i = 0; i <= 1; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 3);
> +
> +	run_flash_test(state, -mtd_write_failure_1_errno);
> +}
> +
> +static void test_mtd_write_bad_block(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	copy_flash_state();
> +	set_bit(expected_bad_blocks, 1);
> +	clear_bit(expected_written_pages, 2);
> +	clear_bit(expected_written_pages, 3);
> +	memcpy(expected_flash_memory, image_buf, 16);
> +	memset(expected_flash_memory + 16, EMPTY_BYTE, 16);
> +	memcpy(expected_flash_memory + 32, image_buf + 16, 32);
> +	for (int i = 0; i <= 3; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static int mtd_erase_failure_2_eb;
> +static int mtd_erase_failure_2_idx;
> +static int mtd_erase_failure_2_errno;
> +static int mtd_erase_failure_2(libmtd_t desc, const struct mtd_dev_info *mtd,
> +		int fd, int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	if (eb == mtd_erase_failure_2_eb) {
> +		if (mtd_erase_failure_2_idx == 0) {
> +			errno = mtd_erase_failure_2_errno;
> +			return -1;
> +		}
> +		--mtd_erase_failure_2_idx;
> +	}
> +	return default_mtd_erase(desc, mtd, fd, eb);
> +}
> +
> +static void test_mtd_write_bad_block_erase_failure(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	impl_mtd_erase = mtd_erase_failure_2;
> +	mtd_erase_failure_2_eb = 1;
> +	mtd_erase_failure_2_idx = 1;
> +	mtd_erase_failure_2_errno = ERANGE; /* dummy value */
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, 24);
> +	memset(expected_flash_memory + 24, EMPTY_BYTE, 8);
> +	for (int i = 0; i <= 1; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 3);
> +
> +	run_flash_test(state, -mtd_erase_failure_2_errno);
> +}
> +
> +static void test_mtd_write_bad_block_erase_failure_EIO(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	impl_mtd_erase = mtd_erase_failure_2;
> +	mtd_erase_failure_2_eb = 1;
> +	mtd_erase_failure_2_idx = 1;
> +	mtd_erase_failure_2_errno = EIO;
> +
> +	copy_flash_state();
> +	set_bit(expected_bad_blocks, 1);
> +	clear_bit(expected_written_pages, 3);
> +	memcpy(expected_flash_memory, image_buf, 24);
> +	memset(expected_flash_memory + 24, EMPTY_BYTE, 8);
> +	memcpy(expected_flash_memory + 32, image_buf + 16, 32);
> +	for (int i = 0; i <= 3; i++)
> +		clear_bit(expected_locked_blocks, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static int mtd_mark_bad_failure_1_errno;
> +static int mtd_mark_bad_failure_1(const struct mtd_dev_info *mtd, int fd,
> +		int eb)
> +{
> +	check_args(mtd, fd, eb);
> +	errno = mtd_mark_bad_failure_1_errno;
> +	return -1;
> +}
> +
> +static void test_mtd_write_bad_block_mark_not_supported(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	impl_mtd_mark_bad = mtd_mark_bad_failure_1;
> +	mtd_mark_bad_failure_1_errno = EOPNOTSUPP;
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, 16);
> +	memset(expected_flash_memory + 16, EMPTY_BYTE, 16);
> +	memcpy(expected_flash_memory + 32, image_buf + 16, 32);
> +
> +	for (int i = 0; i <= 3; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	for (int i = 2; i <= 3; i++)
> +		clear_bit(expected_written_pages, i);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +static void test_mtd_write_bad_block_mark_failure(void **state)
> +{
> +	image.seek = 0;
> +	image.size = 48;
> +	mtd_ubi_info_s.mtd.size = 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 16;
> +	mtd_ubi_info_s.mtd.min_io_size = 8;
> +	test_init();
> +
> +	impl_mtd_write = mtd_write_failure_1;
> +	mtd_write_failure_1_eb = 1;
> +	mtd_write_failure_1_offs = 8;
> +	mtd_write_failure_1_errno = EIO;
> +
> +	impl_mtd_mark_bad = mtd_mark_bad_failure_1;
> +	mtd_mark_bad_failure_1_errno = ERANGE; /* dummy value */
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, 16);
> +	memset(expected_flash_memory + 16, EMPTY_BYTE, 16);
> +	for (int i = 0; i <= 1; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	for (int i = 2; i <= 3; i++)
> +		clear_bit(expected_written_pages, i);
> +
> +	run_flash_test(state, -mtd_mark_bad_failure_1_errno);
> +}
> +
> +static void test_multiple_callbacks(void **state)
> +{
> +	/* flash_write() is typically executed with (len <= 16 * 1024) */
> +	image.seek = 0;
> +	image.size = 63 * 1024;
> +	mtd_ubi_info_s.mtd.size = 64 * 1024;
> +	mtd_ubi_info_s.mtd.eb_size = 8 * 1024;
> +	mtd_ubi_info_s.mtd.min_io_size = 1024;
> +	test_init();
> +
> +	copy_flash_state();
> +	memcpy(expected_flash_memory, image_buf, image.size);
> +	memset(expected_flash_memory + 63 * 1024, EMPTY_BYTE, 1024);
> +	for (int i = 0; i <= 7; i++)
> +		clear_bit(expected_locked_blocks, i);
> +	clear_bit(expected_written_pages, 63);
> +
> +	run_flash_test(state, 0);
> +}
> +
> +#define TEST(name) \
> +	cmocka_unit_test_setup_teardown(test_##name, test_setup, test_teardown)
> +
> +int main(void)
> +{
> +	static const struct CMUnitTest tests[] = {
> +		TEST(simple),
> +		TEST(simple_NOR),
> +		TEST(padding_less_than_page),
> +		TEST(padding_page),
> +		TEST(skip_write_page_empty_bytes),
> +		TEST(padding_more_than_page),
> +		TEST(seek),
> +		TEST(seek_not_multiple_of_eb_size),
> +		TEST(not_enough_flash),
> +		TEST(invalid_mtd_device),
> +		TEST(invalid_image_size),
> +		TEST(empty_image),
> +		TEST(mtd_dev_open_failure),
> +		TEST(malloc_failure),
> +		TEST(skip_bad_blocks),
> +		TEST(too_many_known_bad_blocks),
> +		TEST(too_many_unknown_bad_blocks),
> +		TEST(mtd_is_bad_not_supported),
> +		TEST(mtd_is_bad_failure),
> +		TEST(mtd_is_locked_not_supported),
> +		TEST(mtd_is_locked_failure),
> +		TEST(mtd_unlock_not_supported),
> +		TEST(mtd_unlock_failure),
> +		TEST(mtd_read_failure),
> +		TEST(mtd_read_no_erase_empty_flash_bytes),
> +		TEST(mtd_erase_not_supported),
> +		TEST(mtd_erase_failure),
> +		TEST(mtd_erase_failure_EIO),
> +		TEST(no_mtd_write_empty_image_bytes),
> +		TEST(mtd_write_failure),
> +		TEST(mtd_write_bad_block),
> +		TEST(mtd_write_bad_block_erase_failure),
> +		TEST(mtd_write_bad_block_erase_failure_EIO),
> +		TEST(mtd_write_bad_block_mark_not_supported),
> +		TEST(mtd_write_bad_block_mark_failure),
> +		TEST(multiple_callbacks),
> +	};
> +	return cmocka_run_group_tests_name("flash_handler", tests, group_setup,
> +	                                   group_teardown);
> +}

Best regards,
Stefano Babic