erofs-utils: resize image to the correct size

In the end, it's necessary to resize image to
the proper size since buffers could be dropped.

Reviewed-by: Li Guifu <blucerlee@gmail.com>
Signed-off-by: Gao Xiang <gaoxiang25@huawei.com>

3 files changed, 40 insertions, 4 deletions

diff --git a/include/erofs/io.h b/include/erofs/io.h
index 4b574bd..9775047 100644
--- a/include/erofs/io.h
+++ b/include/erofs/io.h
@@ -21,6 +21,7 @@ void dev_close(void);
 int dev_write(const void *buf, u64 offset, size_t len);
 int dev_fillzero(u64 offset, size_t len, bool padding);
 int dev_fsync(void);
+int dev_resize(erofs_blk_t nblocks);
 u64 dev_length(void);
 
 static inline int blk_write(const void *buf, erofs_blk_t blkaddr,
diff --git a/lib/io.c b/lib/io.c
index 15c5a35..7f5f94d 100644
--- a/lib/io.c
+++ b/lib/io.c
@@ -79,6 +79,7 @@ int dev_open(const char *dev)
 			close(fd);
 			return ret;
 		}
+		erofs_devsz = round_down(erofs_devsz, EROFS_BLKSIZ);
 		break;
 	case S_IFREG:
 		ret = ftruncate(fd, 0);
@@ -98,7 +99,6 @@ int dev_open(const char *dev)
 
 	erofs_devname = dev;
 	erofs_devfd = fd;
-	erofs_devsz = round_down(erofs_devsz, EROFS_BLKSIZ);
 
 	erofs_info("successfully to open %s", dev);
 	return 0;
@@ -177,3 +177,33 @@ int dev_fsync(void)
 	}
 	return 0;
 }
+
+int dev_resize(unsigned int blocks)
+{
+	int ret;
+	struct stat st;
+	u64 length;
+
+	if (cfg.c_dry_run || erofs_devsz != INT64_MAX)
+		return 0;
+
+	ret = fstat(erofs_devfd, &st);
+	if (ret) {
+		erofs_err("failed to fstat.");
+		return -errno;
+	}
+
+	length = (u64)blocks * EROFS_BLKSIZ;
+	if (st.st_size == length)
+		return 0;
+	if (st.st_size > length)
+		return ftruncate(erofs_devfd, length);
+
+	length = length - st.st_size;
+#if defined(HAVE_FALLOCATE)
+	if (fallocate(erofs_devfd, 0, st.st_size, length) >= 0)
+		return 0;
+#endif
+	return dev_fillzero(st.st_size, length, true);
+}
+
diff --git a/mkfs/main.c b/mkfs/main.c
index 5efbf30..2dfd68e 100644
--- a/mkfs/main.c
+++ b/mkfs/main.c
@@ -145,7 +145,8 @@ static int mkfs_parse_options_cfg(int argc, char *argv[])
 }
 
 int erofs_mkfs_update_super_block(struct erofs_buffer_head *bh,
-				  erofs_nid_t root_nid)
+				  erofs_nid_t root_nid,
+				  erofs_blk_t *blocks)
 {
 	struct erofs_super_block sb = {
 		.magic     = cpu_to_le32(EROFS_SUPER_MAGIC_V1),
@@ -166,7 +167,8 @@ int erofs_mkfs_update_super_block(struct erofs_buffer_head *bh,
 		sb.build_time_nsec = cpu_to_le32(t.tv_usec);
 	}
 
-	sb.blocks       = cpu_to_le32(erofs_mapbh(NULL, true));
+	*blocks         = erofs_mapbh(NULL, true);
+	sb.blocks       = cpu_to_le32(*blocks);
 	sb.root_nid     = cpu_to_le16(root_nid);
 
 	buf = calloc(sb_blksize, 1);
@@ -189,6 +191,7 @@ int main(int argc, char **argv)
 	struct erofs_inode *root_inode;
 	erofs_nid_t root_nid;
 	struct stat64 st;
+	erofs_blk_t nblocks;
 
 	erofs_init_configure();
 	fprintf(stderr, "%s %s\n", basename(argv[0]), cfg.c_version);
@@ -250,13 +253,15 @@ int main(int argc, char **argv)
 	root_nid = erofs_lookupnid(root_inode);
 	erofs_iput(root_inode);
 
-	err = erofs_mkfs_update_super_block(sb_bh, root_nid);
+	err = erofs_mkfs_update_super_block(sb_bh, root_nid, &nblocks);
 	if (err)
 		goto exit;
 
 	/* flush all remaining buffers */
 	if (!erofs_bflush(NULL))
 		err = -EIO;
+	else
+		err = dev_resize(nblocks);
 exit:
 	z_erofs_compress_exit();
 	dev_close();