/*
 * patch_io.c --- This is the "patch" io manager that writes the new data into
 * a separate sparse file to apply it later.
 *
 * Patch file format:
 * 1) sparse data blocks - same size as the filesystem, but only changed blocks are written
 * 2) updated block bitmap - fs_size/block_size/8 bytes
 * 3) 4 byte FS block size
 * 4) 8 byte FS size in blocks
 *
 * Copyright (c) Vitaliy Filippov <vitalif@mail.ru> 2014
 * License: GNU GPLv2 or later
 *
 * %Begin-Header%
 * This file may be redistributed under the terms of the GNU Library
 * General Public License, version 2 or later.
 * %End-Header%
 */

#define _LARGEFILE_SOURCE
#define _LARGEFILE64_SOURCE

#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>

#include <ext2fs/ext2_fs.h>
#include <ext2fs/ext2fs.h>

#ifdef __GNUC__
#define ATTR(x) __attribute__(x)
#else
#define ATTR(x)
#endif

#define EXT2_CHECK_MAGIC(struct, code) if ((struct)->magic != (code)) return (code)

struct patch_private_data {
	int	magic;
	char *patch_file;
	int patch_fd;
	int block_size;
	blk64_t size;
	ext2fs_generic_bitmap bmap;
	/* The backing io channel */
	io_channel real;
	/* to support offset in unix I/O manager */
	ext2_loff_t offset;
};

static errcode_t patch_open(const char *name, int flags, io_channel *channel);
static errcode_t patch_close(io_channel channel);
static errcode_t patch_set_blksize(io_channel channel, int blksize);
static errcode_t patch_read_blk64(io_channel channel, unsigned long long block, int count, void *data);
static errcode_t patch_write_blk64(io_channel channel, unsigned long long block, int count, const void *data);
static errcode_t patch_read_blk(io_channel channel, unsigned long block, int count, void *data);
static errcode_t patch_write_blk(io_channel channel, unsigned long block, int count, const void *data);
static errcode_t patch_flush(io_channel channel);
static errcode_t patch_write_byte(io_channel channel, unsigned long offset, int size, const void *data);
static errcode_t patch_set_option(io_channel channel, const char *option, const char *arg);
static errcode_t patch_get_stats(io_channel channel, io_stats *stats);

static struct struct_io_manager struct_patch_manager = {
	EXT2_ET_MAGIC_IO_MANAGER,
	"Patch I/O Manager",
	patch_open,
	patch_close,
	patch_set_blksize,
	patch_read_blk,
	patch_write_blk,
	patch_flush,
	patch_write_byte,
	patch_set_option,
	patch_get_stats,
	patch_read_blk64,
	patch_write_blk64,
};

io_manager patch_io_manager = &struct_patch_manager;
static char *patch_file;
static io_manager patch_io_backing_manager;

errcode_t set_patch_io_backing_manager(io_manager manager)
{
	patch_io_backing_manager = manager;
	return 0;
}

errcode_t set_patch_io_patch_file(char *file)
{
	patch_file = file;
	return 0;
}

static void patch_read_bmap(struct patch_private_data *data, int fd)
{
	int bufsize = 65536;
	blk64_t i, r;
	void *buf = malloc(bufsize);
	ext2fs_llseek(fd, data->size*data->block_size, SEEK_SET);
	for (i = 0; i < data->size/8; )
	{
		r = bufsize;
		if (data->size/8 - i < r)
			r = data->size/8 - i;
		r = read(fd, buf, r);
		if (r < 0)
		{
			r = 0;
			if (errno != EAGAIN)
				break;
		}
		ext2fs_set_generic_bmap_range(data->bmap, i*8, r*8, buf);
		i += r;
	}
	free(buf);
}

static void patch_write_bmap(struct patch_private_data *data, int fd)
{
	int bufsize = 65536;
	blk64_t i, r;
	void *buf = malloc(bufsize);
	ext2fs_llseek(fd, data->size*data->block_size, SEEK_SET);
	for (i = 0; i < data->size/8; )
	{
		r = bufsize;
		if (data->size/8 - i < r)
			r = data->size/8 - i;
		ext2fs_get_generic_bmap_range(data->bmap, i*8, r*8, buf);
		r = write(fd, buf, r);
		if (r < 0)
		{
			r = 0;
			if (errno != EAGAIN)
				break;
		}
		i += r;
	}
	free(buf);
	write(fd, &data->block_size, sizeof(int));
	write(fd, &data->size, sizeof(blk64_t));
}

static errcode_t patch_open(const char *name, int flags, io_channel *channel)
{
	io_channel io = NULL;
	struct patch_private_data *data = NULL;
	errcode_t retval;

	if (name == 0)
		return EXT2_ET_BAD_DEVICE_NAME;
	retval = ext2fs_get_mem(sizeof(struct struct_io_channel), &io);
	if (retval)
		goto cleanup;
	memset(io, 0, sizeof(struct struct_io_channel));
	io->magic = EXT2_ET_MAGIC_IO_CHANNEL;
	retval = ext2fs_get_mem(sizeof(struct patch_private_data), &data);
	if (retval)
		goto cleanup;

	io->manager = patch_io_manager;
	retval = ext2fs_get_mem(strlen(name)+1, &io->name);
	if (retval)
		goto cleanup;

	strcpy(io->name, name);
	io->private_data = data;
	io->block_size = 1024;
	io->read_error = 0;
	io->write_error = 0;
	io->refcount = 1;

	memset(data, 0, sizeof(struct patch_private_data));
	data->magic = EXT2_ET_MAGIC_UNIX_IO_CHANNEL;

	if (patch_io_backing_manager)
	{
		retval = patch_io_backing_manager->open(name, flags, &data->real);
		if (retval)
			goto cleanup;
	}

	if (patch_file)
	{
		long long size;
		data->patch_file = strdup(patch_file);
		data->patch_fd = open(data->patch_file, O_CREAT|O_RDWR, 0666);
		if (data->patch_fd < 0)
		{
			retval = errno;
			goto cleanup;
		}
		size = ext2fs_llseek(data->patch_fd, 0, SEEK_END);
		if (size < 0)
		{
			retval = errno;
			goto cleanup;
		}
		if (size > sizeof(int)+sizeof(blk64_t))
		{
			ext2fs_llseek(data->patch_fd, size-sizeof(int)-sizeof(blk64_t), 0);
			read(data->patch_fd, &data->block_size, sizeof(int));
			read(data->patch_fd, &data->size, sizeof(blk64_t));
			retval = ext2fs_make_generic_bitmap(EXT2_ET_MAGIC_BLOCK_BITMAP, NULL,
				0, data->size, data->size, "overwritten blocks", 0, &data->bmap);
			if (retval)
				goto cleanup;
			patch_read_bmap(data, data->patch_fd);
		}
	}

	*channel = io;
	return 0;

cleanup:
	if (data)
	{
		if (data->bmap)
			ext2fs_free_generic_bmap(data->bmap);
		if (data->patch_fd >= 0)
			close(data->patch_fd);
		if (data->patch_file)
			free(data->patch_file);
		if (data->real)
			io_channel_close(data->real);
		ext2fs_free_mem(&data);
	}
	if (io)
	{
		if (io->name)
			ext2fs_free_mem(&io->name);
		ext2fs_free_mem(&io);
	}
	return retval;
}

static errcode_t patch_close(io_channel channel)
{
	struct patch_private_data *data;
	errcode_t retval = 0;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if (--channel->refcount > 0)
		return 0;

	if (data->bmap)
	{
		if (data->patch_fd >= 0)
			patch_write_bmap(data, data->patch_fd);
		ext2fs_free_generic_bmap(data->bmap);
	}
	if (data->patch_fd >= 0)
		close(data->patch_fd);
	if (data->patch_file)
		free(data->patch_file);
	if (data->real)
		retval = io_channel_close(data->real);
	ext2fs_free_mem(&channel->private_data);
	if (channel->name)
		ext2fs_free_mem(&channel->name);
	ext2fs_free_mem(&channel);

	return retval;
}

static errcode_t patch_set_blksize(io_channel channel, int blksize)
{
	struct patch_private_data *data;
	errcode_t retval = 0;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if (data->block_size && data->block_size != blksize)
		return EINVAL;
	if (data->real)
		retval = io_channel_set_blksize(data->real, blksize);
	channel->block_size = blksize;
	return retval;
}

static errcode_t re_read(int fd, unsigned long long offset, ssize_t size, void *buf)
{
	ssize_t r, done = 0;
	if ((unsigned)ext2fs_llseek(fd, offset, SEEK_SET) != offset)
		return errno ? errno : EXT2_ET_LLSEEK_FAILED;
	while (done < size)
	{
		r = read(fd, buf+done, size-done);
		if (r < 0 && errno != EAGAIN)
			break;
		done += r;
	}
	if (done < size)
		return errno;
	return 0;
}

static errcode_t patch_read_blk64(io_channel channel, unsigned long long block, int count, void *buf)
{
	errcode_t retval = 0;
	struct patch_private_data *data;
	int b, n;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if (count < 0)
	{
		if (-count <= channel->block_size)
		{
			if (data->bmap && ext2fs_test_generic_bitmap(data->bmap, block))
				retval = re_read(data->patch_fd, block*channel->block_size, -count, buf);
			else
				retval = io_channel_read_blk64(data->real, block, count, buf);
			return retval;
		}
		else
			return EINVAL;
	}
	for (b = 0; b < count; )
	{
		for (n = 0; (b+n < count) && data->bmap && ext2fs_test_generic_bitmap(data->bmap, block+b+n); n++) {}
		if (n > 0)
		{
			retval = re_read(data->patch_fd, (block+b)*channel->block_size, n*channel->block_size, buf+b*channel->block_size);
			if (retval)
				break;
			b += n;
		}
		for (n = 0; (b+n < count) && (!data->bmap || !ext2fs_test_generic_bitmap(data->bmap, block+b+n)); n++) {}
		if (n > 0 && data->real)
		{
			retval = io_channel_read_blk64(data->real, block+b, n, buf+(block+b)*channel->block_size);
			if (retval)
				break;
			b += n;
		}
	}

	return retval;
}

static errcode_t patch_read_blk(io_channel channel, unsigned long block, int count, void *buf)
{
	return patch_read_blk64(channel, block, count, buf);
}

static errcode_t patch_check_bmap_init(struct patch_private_data *data, io_channel channel)
{
	errcode_t retval = 0;
	if (!data->bmap)
	{
		data->block_size = channel->block_size;
		retval = ext2fs_get_device_size2(channel->name, data->block_size, &data->size);
		if (retval)
			return retval;
		retval = ext2fs_make_generic_bitmap(EXT2_ET_MAGIC_BLOCK_BITMAP, NULL,
			0, data->size, data->size, "overwritten blocks", 0, &data->bmap);
	}
	return retval;
}

static errcode_t patch_write_blk64(io_channel channel, unsigned long long block, int count, const void *buf)
{
	struct patch_private_data *data;
	errcode_t retval = 0;
	ssize_t r, left = count*channel->block_size;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if ((unsigned)ext2fs_llseek(data->patch_fd, block*channel->block_size, SEEK_SET) != block*channel->block_size)
		return errno ? errno : EXT2_ET_LLSEEK_FAILED;
	retval = patch_check_bmap_init(data, channel);
	if (retval)
		return retval;
	if (count < 0)
	{
		if (-count > channel->block_size)
			return EINVAL;
		left = -count;
		count = 1;
	}
	ext2fs_mark_block_bitmap_range2(data->bmap, block, count);
	buf += left;
	while (left > 0)
	{
		r = write(data->patch_fd, buf-left, left);
		if (r < 0 && errno != EAGAIN)
			break;
		left -= r;
	}
	return left ? errno : 0;
}

static errcode_t patch_write_blk(io_channel channel, unsigned long block, int count, const void *buf)
{
	return patch_write_blk64(channel, block, count, buf);
}

static errcode_t patch_write_byte(io_channel channel, unsigned long offset, int size, const void *buf)
{
	return EXT2_ET_UNIMPLEMENTED;
}

/*
 * Flush data buffers to disk.
 */
static errcode_t patch_flush(io_channel channel)
{
	errcode_t retval = 0;
	struct patch_private_data *data;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if (data->real)
		retval = io_channel_flush(data->real);
	if (data->patch_fd)
		fsync(data->patch_fd);

	return retval;
}

static errcode_t patch_set_option(io_channel channel, const char *option, const char *arg)
{
	errcode_t retval = 0;
	struct patch_private_data *data;
	unsigned long tmp;
	char *end;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	/*
	 * Need to support offset option to work with
	 * Unix I/O manager
	 */
	if (data->real && data->real->manager->set_option)
		retval = data->real->manager->set_option(data->real, option, arg);
	if (!retval && !strcmp(option, "offset"))
	{
		if (!arg)
			return EXT2_ET_INVALID_ARGUMENT;

		tmp = strtoul(arg, &end, 0);
		if (*end)
			return EXT2_ET_INVALID_ARGUMENT;
		data->offset = tmp;
	}
	return retval;
}

static errcode_t patch_get_stats(io_channel channel, io_stats *stats)
{
	errcode_t retval = 0;
	struct patch_private_data *data;

	EXT2_CHECK_MAGIC(channel, EXT2_ET_MAGIC_IO_CHANNEL);
	data = (struct patch_private_data *) channel->private_data;
	EXT2_CHECK_MAGIC(data, EXT2_ET_MAGIC_UNIX_IO_CHANNEL);

	if (data->real)
		retval = (data->real->manager->get_stats)(data->real, stats);

	return retval;
}