#include <pthread.h>

#include <umem.h>

#include <sha2.h>

static const zio_cksum_t zero_cksum = { 0 };

typedef struct dedup_arg {

	int	inputfd;

	int	outputfd;

	libzfs_handle_t  *dedup_hdl;

} dedup_arg_t;

typedef struct dataref {

	uint64_t ref_guid;

	uint64_t ref_object;

	uint64_t ref_offset;

} dataref_t;

typedef struct dedup_entry {

	struct dedup_entry	*dde_next;

	zio_cksum_t dde_chksum;

	dataref_t dde_ref;

} dedup_entry_t;

#define	MAX_DDT_PHYSMEM_PERCENT		20

#define	SMALLEST_POSSIBLE_MAX_DDT_MB		128

typedef struct dedup_table {

	dedup_entry_t	**dedup_hash_array;

	umem_cache_t	*ddecache;

	uint64_t	max_ddt_size;  /* max dedup table size in bytes */

	uint64_t	cur_ddt_size;  /* current dedup table size in bytes */

	uint64_t	ddt_count;

	int		numhashbits;

	boolean_t	ddt_full;

} dedup_table_t;

static int

high_order_bit(uint64_t n)

{

	int count;

	for (count = 0; n != 0; count++)

		n >>= 1;

	return (count);

}

static size_t

ssread(void *buf, size_t len, FILE *stream)

{

	size_t outlen;

	if ((outlen = fread(buf, len, 1, stream)) == 0)

		return (0);

	return (outlen);

}

static void

ddt_hash_append(libzfs_handle_t *hdl, dedup_table_t *ddt, dedup_entry_t **ddepp,

    zio_cksum_t *cs, dataref_t *dr)

{

	dedup_entry_t	*dde;

	if (ddt->cur_ddt_size >= ddt->max_ddt_size) {

		if (ddt->ddt_full == B_FALSE) {

			zfs_error_aux(hdl, dgettext(TEXT_DOMAIN,

			    "Dedup table full.  Deduplication will continue "

			    "with existing table entries"));

			ddt->ddt_full = B_TRUE;

		}

		return;

	}

	if ((dde = umem_cache_alloc(ddt->ddecache, UMEM_DEFAULT))

	    != NULL) {

		assert(*ddepp == NULL);

		dde->dde_next = NULL;

		dde->dde_chksum = *cs;

		dde->dde_ref = *dr;

		*ddepp = dde;

		ddt->cur_ddt_size += sizeof (dedup_entry_t);

		ddt->ddt_count++;

	}

}

/*

 * Using the specified dedup table, do a lookup for an entry with

 * the checksum cs.  If found, return the block's reference info

 * in *dr. Otherwise, insert a new entry in the dedup table, using

 * the reference information specified by *dr.

 *

 * return value:  true - entry was found

 *		  false - entry was not found

 */

static boolean_t

ddt_update(libzfs_handle_t *hdl, dedup_table_t *ddt, zio_cksum_t *cs,

    dataref_t *dr)

{

	uint32_t hashcode;

	dedup_entry_t **ddepp;

	hashcode = BF64_GET(cs->zc_word[0], 0, ddt->numhashbits);

	for (ddepp = &(ddt->dedup_hash_array[hashcode]); *ddepp != NULL;

	    ddepp = &((*ddepp)->dde_next)) {

		if (ZIO_CHECKSUM_EQUAL(((*ddepp)->dde_chksum), *cs)) {

			*dr = (*ddepp)->dde_ref;

			return (B_TRUE);

		}

	}

	ddt_hash_append(hdl, ddt, ddepp, cs, dr);

	return (B_FALSE);

}

static int

cksum_and_write(const void *buf, uint64_t len, zio_cksum_t *zc, int outfd)

{

	fletcher_4_incremental_native(buf, len, zc);

	return (write(outfd, buf, len));

}

/*

 * This function is started in a separate thread when the dedup option

 * has been requested.  The main send thread determines the list of

 * snapshots to be included in the send stream and makes the ioctl calls

 * for each one.  But instead of having the ioctl send the output to the

 * the output fd specified by the caller of zfs_send()), the

 * ioctl is told to direct the output to a pipe, which is read by the

 * alternate thread running THIS function.  This function does the

 * dedup'ing by:

 *  1. building a dedup table (the DDT)

 *  2. doing checksums on each data block and inserting a record in the DDT

 *  3. looking for matching checksums, and

 *  4.  sending a DRR_WRITE_BYREF record instead of a write record whenever

 *      a duplicate block is found.

 * The output of this function then goes to the output fd requested

 * by the caller of zfs_send().

 */

static void *

cksummer(void *arg)

{

	dedup_arg_t *dda = arg;

	char *buf = malloc(1<<20);

	dmu_replay_record_t thedrr;

	dmu_replay_record_t *drr = &thedrr;

	struct drr_begin *drrb = &thedrr.drr_u.drr_begin;

	struct drr_end *drre = &thedrr.drr_u.drr_end;

	struct drr_object *drro = &thedrr.drr_u.drr_object;

	struct drr_write *drrw = &thedrr.drr_u.drr_write;

	FILE *ofp;

	int outfd;

	dmu_replay_record_t wbr_drr;

	struct drr_write_byref *wbr_drrr = &wbr_drr.drr_u.drr_write_byref;

	dedup_table_t ddt;

	zio_cksum_t stream_cksum;

	uint64_t physmem = sysconf(_SC_PHYS_PAGES) * sysconf(_SC_PAGESIZE);

	uint64_t numbuckets;

	ddt.max_ddt_size =

	    MAX((physmem * MAX_DDT_PHYSMEM_PERCENT)/100,

	    SMALLEST_POSSIBLE_MAX_DDT_MB<<20);

	numbuckets = ddt.max_ddt_size/(sizeof (dedup_entry_t));

	/*

	 * numbuckets must be a power of 2.  Increase number to

	 * a power of 2 if necessary.

	 */

	if (!ISP2(numbuckets))

		numbuckets = 1 << high_order_bit(numbuckets);

	ddt.dedup_hash_array = calloc(numbuckets, sizeof (dedup_entry_t *));

	ddt.ddecache = umem_cache_create("dde", sizeof (dedup_entry_t), 0,

	    NULL, NULL, NULL, NULL, NULL, 0);

	ddt.cur_ddt_size = numbuckets * sizeof (dedup_entry_t *);

	ddt.numhashbits = high_order_bit(numbuckets) - 1;

	ddt.ddt_full = B_FALSE;

	/* Initialize the write-by-reference block. */

	wbr_drr.drr_type = DRR_WRITE_BYREF;

	wbr_drr.drr_payloadlen = 0;

	outfd = dda->outputfd;

	ofp = fdopen(dda->inputfd, "r");

	while (ssread(drr, sizeof (dmu_replay_record_t), ofp) != 0) {

		switch (drr->drr_type) {

		case DRR_BEGIN:

		{

			int	fflags;

			ZIO_SET_CHECKSUM(&stream_cksum, 0, 0, 0, 0);

			/* set the DEDUP feature flag for this stream */

			fflags = DMU_GET_FEATUREFLAGS(drrb->drr_versioninfo);

			fflags |= DMU_BACKUP_FEATURE_DEDUP;

			DMU_SET_FEATUREFLAGS(drrb->drr_versioninfo, fflags);

			if (cksum_and_write(drr, sizeof (dmu_replay_record_t),

			    &stream_cksum, outfd) == -1)

				goto out;

			if (DMU_GET_STREAM_HDRTYPE(drrb->drr_versioninfo) ==

			    DMU_COMPOUNDSTREAM && drr->drr_payloadlen != 0) {

				int sz = drr->drr_payloadlen;

				if (sz > 1<<20) {

					free(buf);

					buf = malloc(sz);

				}

				(void) ssread(buf, sz, ofp);

				if (ferror(stdin))

					perror("fread");

				if (cksum_and_write(buf, sz, &stream_cksum,

				    outfd) == -1)

					goto out;

			}

			break;

		}

		case DRR_END:

		{

			/* use the recalculated checksum */

			ZIO_SET_CHECKSUM(&drre->drr_checksum,

			    stream_cksum.zc_word[0], stream_cksum.zc_word[1],

			    stream_cksum.zc_word[2], stream_cksum.zc_word[3]);

			if ((write(outfd, drr,

			    sizeof (dmu_replay_record_t))) == -1)

				goto out;

			break;

		}

		case DRR_OBJECT:

		{

			if (cksum_and_write(drr, sizeof (dmu_replay_record_t),

			    &stream_cksum, outfd) == -1)

				goto out;

			if (drro->drr_bonuslen > 0) {

				(void) ssread(buf,

				    P2ROUNDUP((uint64_t)drro->drr_bonuslen, 8),

				    ofp);

				if (cksum_and_write(buf,

				    P2ROUNDUP((uint64_t)drro->drr_bonuslen, 8),

				    &stream_cksum, outfd) == -1)

					goto out;

			}

			break;

		}

		case DRR_FREEOBJECTS:

		{

			if (cksum_and_write(drr, sizeof (dmu_replay_record_t),

			    &stream_cksum, outfd) == -1)

				goto out;

			break;

		}

		case DRR_WRITE:

		{

			dataref_t	dataref;

			(void) ssread(buf, drrw->drr_length, ofp);

			/*

			 * If the block doesn't already have a dedup

			 * checksum, calculate one.

			 */

			if (ZIO_CHECKSUM_EQUAL(drrw->drr_blkcksum,

			    zero_cksum)) {

				SHA256_CTX	ctx;

				zio_cksum_t	tmpsha256;

				SHA256Init(&ctx);

				SHA256Update(&ctx, buf, drrw->drr_length);

				SHA256Final(&tmpsha256, &ctx);

				drrw->drr_blkcksum.zc_word[0] =

				    BE_64(tmpsha256.zc_word[0]);

				drrw->drr_blkcksum.zc_word[1] =

				    BE_64(tmpsha256.zc_word[1]);

				drrw->drr_blkcksum.zc_word[2] =

				    BE_64(tmpsha256.zc_word[2]);

				drrw->drr_blkcksum.zc_word[3] =

				    BE_64(tmpsha256.zc_word[3]);

			}

			dataref.ref_guid = drrw->drr_toguid;

			dataref.ref_object = drrw->drr_object;

			dataref.ref_offset = drrw->drr_offset;

			if (ddt_update(dda->dedup_hdl, &ddt,

			    &drrw->drr_blkcksum, &dataref)) {

				/* block already present in stream */

				wbr_drrr->drr_object = drrw->drr_object;

				wbr_drrr->drr_offset = drrw->drr_offset;

				wbr_drrr->drr_length = drrw->drr_length;

				wbr_drrr->drr_toguid = drrw->drr_toguid;

				wbr_drrr->drr_refguid = dataref.ref_guid;

				wbr_drrr->drr_refobject =

				    dataref.ref_object;

				wbr_drrr->drr_refoffset =

				    dataref.ref_offset;

				wbr_drrr->drr_blkcksum = drrw->drr_blkcksum;

				if (cksum_and_write(&wbr_drr,

				    sizeof (dmu_replay_record_t), &stream_cksum,

				    outfd) == -1)

					goto out;

			} else {

				/* block not previously seen */

				if (cksum_and_write(drr,

				    sizeof (dmu_replay_record_t), &stream_cksum,

				    outfd) == -1)

					goto out;

				if (cksum_and_write(buf,

				    drrw->drr_length,

				    &stream_cksum, outfd) == -1)

					goto out;

			}

			break;

		}

		case DRR_FREE:

		{

			if (cksum_and_write(drr, sizeof (dmu_replay_record_t),

			    &stream_cksum, outfd) == -1)

				goto out;

			break;

		}

		default:

			(void) printf("INVALID record type 0x%x\n",

			    drr->drr_type);

			/* should never happen, so assert */

			assert(B_FALSE);

		}

	}

out:

	umem_cache_destroy(ddt.ddecache);

	free(ddt.dedup_hash_array);

	free(buf);

	(void) fclose(ofp);

	return (NULL);

}