/*
  Copyright (c) 2008-2012 Red Hat, Inc. <http://www.redhat.com>
  This file is part of GlusterFS.

  This file is licensed to you under your choice of the GNU Lesser
  General Public License, version 3 or any later version (LGPLv3 or
  later), or the GNU General Public License, version 2 (GPLv2), in all
  cases as published by the Free Software Foundation.
*/

/*
 * stripe-merge.c
 *
 * This program recovers an original file based on the striped files stored on
 * the individual bricks of a striped volume. The file format and stripe
 * geometry is validated through the extended attributes stored in the file.
 *
 * TODO: Support optional xattr recovery (i.e., user xattrs). Perhaps provide a
 * 	 command-line flag to toggle this behavior.
 */

#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <errno.h>
#include <string.h>
#include <attr/xattr.h>
#include <fnmatch.h>

#define ATTRNAME_STRIPE_INDEX "trusted.*.stripe-index"
#define ATTRNAME_STRIPE_COUNT "trusted.*.stripe-count"
#define ATTRNAME_STRIPE_SIZE "trusted.*.stripe-size"
#define ATTRNAME_STRIPE_COALESCE "trusted.*.stripe-coalesce"

#define INVALID_FD -1
#define INVALID_MODE UINT32_MAX

struct file_stripe_info {
	int stripe_count;
	int stripe_size;
	int coalesce;
	mode_t mode;
	int fd[0];
};

static int close_files(struct file_stripe_info *);

static struct
file_stripe_info *alloc_file_stripe_info(int count)
{
	int i;
	struct file_stripe_info *finfo;

	finfo = calloc(1, sizeof(struct file_stripe_info) +
		(sizeof(int) * count));
	if (!finfo)
		return NULL;

	for (i = 0; i < count; i++)
		finfo->fd[i] = INVALID_FD;

	finfo->mode = INVALID_MODE;
	finfo->coalesce = INVALID_FD;

	return finfo;
}

/*
 * Search for an attribute matching the provided pattern. Return a count for
 * the total number of matching entries (including 0). Allocate a buffer for
 * the first matching entry found.
 */
static int
get_stripe_attr_name(const char *path, const char *pattern, char **attrname)
{
	char attrbuf[4096];
	char *ptr, *match = NULL;
	int len, r, match_count = 0;

	if (!path || !pattern || !attrname)
		return -1;

	len = listxattr(path, attrbuf, sizeof(attrbuf));
	if (len < 0)
		return len;

	ptr = attrbuf;
	while (ptr) {
		r = fnmatch(pattern, ptr, 0);
		if (!r) {
			if (!match)
				match = ptr;
			match_count++;
		} else if (r != FNM_NOMATCH) {
			return -1;
		}

		len -= strlen(ptr) + 1;
		if (len > 0)
			ptr += strlen(ptr) + 1;
		else
			ptr = NULL;
	}

	if (match)
		*attrname = strdup(match);

	return match_count;
}

/*
 * Get the integer representation of a named attribute.
 */
static int
get_stripe_attr_val(const char *path, const char *attr, int *val)
{
	char attrbuf[4096];
	int len;

	if (!path || !attr || !val)
		return -1;

	len = getxattr(path, attr, attrbuf, sizeof(attrbuf));
	if (len < 0)
		return len;

	*val = atoi(attrbuf);

	return 0;
}

/*
 * Get an attribute name/value (assumed to be an integer) pair based on a
 * specified search pattern. A buffer is allocated for the exact attr name
 * returned. Optionally, skip the pattern search if a buffer is provided
 * (which should contain an attribute name).
 *
 * Returns the attribute count or -1 on error. The value parameter is set only
 * when a single attribute is found.
 */
static int
get_attr(const char *path, const char *pattern, char **buf, int *val)
{
	int count = 1;

	if (!buf)
		return -1;

	if (!*buf) {
		count = get_stripe_attr_name(path, pattern, buf);
		if (count > 1) {
			/* pattern isn't good enough */
			fprintf(stderr, "ERROR: duplicate attributes found "
				"matching pattern: %s\n", pattern);
			free(*buf);
			*buf = NULL;
			return count;
		} else if (count < 1) {
			return count;
		}
	}

	if (get_stripe_attr_val(path, *buf, val) < 0)
		return -1;

	return count;
}

/*
 * validate_and_open_files()
 *
 * Open the provided source files and validate the extended attributes. Verify
 * that the geometric attributes are consistent across all of the files and
 * print a warning if any files are missing. We proceed without error in the
 * latter case to support partial recovery.
 */
static struct
file_stripe_info *validate_and_open_files(char *paths[], int count)
{
	int i, val, tmp;
	struct stat sbuf;
	char *stripe_count_attr = NULL;
	char *stripe_size_attr = NULL;
	char *stripe_index_attr = NULL;
	char *stripe_coalesce_attr = NULL;
	struct file_stripe_info *finfo = NULL;

	for (i = 0; i < count; i++) {
		if (!paths[i])
			goto err;

		/*
		 * Check the stripe count first so we can allocate the info
		 * struct with the appropriate number of fds.
		 */
		if (get_attr(paths[i], ATTRNAME_STRIPE_COUNT,
				&stripe_count_attr, &val) != 1) {
			fprintf(stderr, "ERROR: %s: attribute: '%s'\n",
				paths[i], ATTRNAME_STRIPE_COUNT);
			goto err;
		}
		if (!finfo) {
			finfo = alloc_file_stripe_info(val);
			if (!finfo)
				goto err;

			if (val != count)
				fprintf(stderr, "WARNING: %s: stripe-count "
					"(%d) != file count (%d). Result may "
					"be incomplete.\n", paths[i], val,
					count);

			finfo->stripe_count = val;
		} else if (val != finfo->stripe_count) {
			fprintf(stderr, "ERROR %s: invalid stripe count: %d "
				"(expected %d)\n", paths[i], val,
				finfo->stripe_count);
			goto err;
		}

		/*
		 * Get and validate the chunk size.
		 */
		if (get_attr(paths[i], ATTRNAME_STRIPE_SIZE, &stripe_size_attr,
				&val) != 1) {
			fprintf(stderr, "ERROR: %s: attribute: '%s'\n",
				paths[i], ATTRNAME_STRIPE_SIZE);
			goto err;
		}

		if (!finfo->stripe_size) {
			finfo->stripe_size = val;
		} else if (val != finfo->stripe_size) {
			fprintf(stderr, "ERROR: %s: invalid stripe size: %d "
				"(expected %d)\n", paths[i], val,
				finfo->stripe_size);
			goto err;
		}

		/*
		 * stripe-coalesce is a backward compatible attribute. If the
		 * attribute does not exist, assume a value of zero for the
		 * traditional stripe format.
		 */
		tmp = get_attr(paths[i], ATTRNAME_STRIPE_COALESCE,
				&stripe_coalesce_attr, &val);
		if (!tmp) {
			val = 0;
		} else if (tmp != 1) {
			fprintf(stderr, "ERROR: %s: attribute: '%s'\n",
				paths[i], ATTRNAME_STRIPE_COALESCE);
			goto err;
		}

		if (finfo->coalesce == INVALID_FD) {
			finfo->coalesce = val;
		} else if (val != finfo->coalesce) {
			fprintf(stderr, "ERROR: %s: invalid coalesce flag\n",
				paths[i]);
			goto err;
		}

		/*
		 * Get/validate the stripe index and open the file in the
		 * appropriate fd slot.
		 */
		if (get_attr(paths[i], ATTRNAME_STRIPE_INDEX,
				&stripe_index_attr, &val) != 1) {
			fprintf(stderr, "ERROR: %s: attribute: '%s'\n",
				paths[i], ATTRNAME_STRIPE_INDEX);
			goto err;
		}
		if (finfo->fd[val] != INVALID_FD) {
			fprintf(stderr, "ERROR: %s: duplicate stripe index: "
				"%d\n", paths[i], val);
			goto err;
		}

		finfo->fd[val] = open(paths[i], O_RDONLY);
		if (finfo->fd[val] < 0)
			goto err;

		/*
		 * Get the creation mode for the file.
		 */
		if (fstat(finfo->fd[val], &sbuf) < 0)
			goto err;
		if (finfo->mode == INVALID_MODE) {
			finfo->mode = sbuf.st_mode;
		} else if (sbuf.st_mode != finfo->mode) {
			fprintf(stderr, "ERROR: %s: invalid mode\n", paths[i]);
			goto err;
		}
	}

	free(stripe_count_attr);
	free(stripe_size_attr);
	free(stripe_index_attr);
	free(stripe_coalesce_attr);

	return finfo;
err:

	free(stripe_count_attr);
	free(stripe_size_attr);
	free(stripe_index_attr);
	free(stripe_coalesce_attr);

	if (finfo) {
		close_files(finfo);
		free(finfo);
	}

	return NULL;
}

static int
close_files(struct file_stripe_info *finfo)
{
	int i, ret;

	if (!finfo)
		return -1;

	for (i = 0; i < finfo->stripe_count; i++) {
		if (finfo->fd[i] == INVALID_FD)
			continue;

		ret = close(finfo->fd[i]);
		if (ret < 0)
			return ret;
	}

	return ret;
}

/*
 * Generate the original file using files striped in the coalesced format.
 * Data in the striped files is stored at a coalesced offset based on the
 * stripe number.
 *
 * Walk through the finfo fds (which are already ordered) and and iteratively
 * copy stripe_size bytes from the source files to the target file. If a source
 * file is missing, seek past the associated stripe_size bytes in the target
 * file.
 */
static int
generate_file_coalesce(int target, struct file_stripe_info *finfo)
{
	char *buf;
	int ret = 0;
	int r, w, i;

	buf = malloc(finfo->stripe_size);
	if (!buf)
		return -1;

	i = 0;
	while (1) {
		if (finfo->fd[i] == INVALID_FD) {
			if (lseek(target, finfo->stripe_size, SEEK_CUR) < 0)
				break;

			i = (i + 1) % finfo->stripe_count;
			continue;
		}

		r = read(finfo->fd[i], buf, finfo->stripe_size);
		if (r < 0) {
			ret = r;
			break;
		}
		if (!r)
			break;

		w = write(target, buf, r);
		if (w < 0) {
			ret = w;
			break;
		}

		i = (i + 1) % finfo->stripe_count;
	}

	free(buf);
	return ret;
}

/*
 * Generate the original file using files striped with the traditional stripe
 * format. Data in the striped files is stored at the equivalent offset from
 * the source file.
 */
static int
generate_file_traditional(int target, struct file_stripe_info *finfo)
{
	int i, j, max_ret, ret;
	char buf[finfo->stripe_count][4096];

	do {
		char newbuf[4096] = {0, };

		max_ret = 0;
		for (i = 0; i < finfo->stripe_count; i++) {
			memset(buf[i], 0, 4096);
			ret = read(finfo->fd[i], buf[i], 4096);
			if (ret > max_ret)
				max_ret = ret;
		}
		for (i = 0; i < max_ret; i++)
			for (j = 0; j < finfo->stripe_count; j++)
				newbuf[i] |= buf[j][i];
			write(target, newbuf, max_ret);
	} while (max_ret);

	return 0;
}

static int
generate_file(int target, struct file_stripe_info *finfo)
{
	if (finfo->coalesce)
		return generate_file_coalesce(target, finfo);

	return generate_file_traditional(target, finfo);
}

static void
usage(char *name)
{
	fprintf(stderr, "Usage: %s [-o <outputfile>] <inputfile1> "
		"<inputfile2> ...\n", name);
}

int
main(int argc, char *argv[])
{
	int file_count, opt;
	char *opath = NULL;
	int targetfd;
	struct file_stripe_info *finfo;

	while ((opt = getopt(argc, argv, "o:")) != -1) {
		switch (opt) {
		case 'o':
			opath = optarg;
			break;
		default:
			usage(argv[0]);
			return -1;
		}
	}

	file_count = argc - optind;

	if (!opath || !file_count) {
		usage(argv[0]);
		return -1;
	}

	finfo = validate_and_open_files(&argv[optind], file_count);
	if (!finfo)
		goto err;

	targetfd = open(opath, O_RDWR|O_CREAT, finfo->mode);
	if (targetfd < 0)
		goto err;

	if (generate_file(targetfd, finfo) < 0)
		goto err;

	if (fsync(targetfd) < 0)
		fprintf(stderr, "ERROR: %s\n", strerror(errno));
	if (close(targetfd) < 0)
		fprintf(stderr, "ERROR: %s\n", strerror(errno));

	close_files(finfo);
	free(finfo);

	return 0;

err:
	if (finfo) {
		close_files(finfo);
		free(finfo);
	}

	return -1;
}