/*

 * CDDL HEADER START

 *

 * The contents of this file are subject to the terms of the

 * Common Development and Distribution License (the "License").

 * You may not use this file except in compliance with the License.

 *

 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE

 * or http://www.opensolaris.org/os/licensing.

 * See the License for the specific language governing permissions

 * and limitations under the License.

 *

 * When distributing Covered Code, include this CDDL HEADER in each

 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.

 * If applicable, add the following below this CDDL HEADER, with the

 * fields enclosed by brackets "[]" replaced with your own identifying

 * information: Portions Copyright [yyyy] [name of copyright owner]

 *

 * CDDL HEADER END

 */

/*

 * Copyright 2010 Sun Microsystems, Inc.  All rights reserved.

 * Use is subject to license terms.

 */

#include <assert.h>

#include <ctype.h>

#include <errno.h>

#include <libintl.h>

#include <stdio.h>

#include <stdlib.h>

#include <strings.h>

#include <unistd.h>

#include <stddef.h>

#include <fcntl.h>

#include <sys/mount.h>

#include <pthread.h>

#include <umem.h>

#include <libzfs.h>

#include "zfs_namecheck.h"

#include "zfs_prop.h"

#include "zfs_fletcher.h"

#include "libzfs_impl.h"

#include <sha2.h>

#include <sys/zio_checksum.h>

#include <sys/ddt.h>

/* in libzfs_dataset.c */

extern void zfs_setprop_error(libzfs_handle_t *, zfs_prop_t, int, char *);

static int zfs_receive_impl(libzfs_handle_t *, const char *, recvflags_t,

    int, const char *, nvlist_t *, avl_tree_t *, char **);

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;

	uint64_t dde_prop;

	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, uint64_t prop, 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_prop = prop;

		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,

    uint64_t prop, 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) &&

		    (*ddepp)->dde_prop == prop) {

			*dr = (*ddepp)->dde_ref;

			return (B_TRUE);

		}

	}

	ddt_hash_append(hdl, ddt, ddepp, cs, prop, 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 = {0};

	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_BACKUP_FEATURE_DEDUPPROPS);

			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);

			/*

			 * Use the existing checksum if it's dedup-capable,

			 * else calculate a SHA256 checksum for it.

			 */

			if (ZIO_CHECKSUM_EQUAL(drrw->drr_key.ddk_cksum,

			    zero_cksum) ||

			    !DRR_IS_DEDUP_CAPABLE(drrw->drr_checksumflags)) {

				SHA256_CTX	ctx;

				zio_cksum_t	tmpsha256;

				SHA256Init(&ctx);

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

				SHA256Final(&tmpsha256, &ctx);

				drrw->drr_key.ddk_cksum.zc_word[0] =

				    BE_64(tmpsha256.zc_word[0]);

				drrw->drr_key.ddk_cksum.zc_word[1] =

				    BE_64(tmpsha256.zc_word[1]);

				drrw->drr_key.ddk_cksum.zc_word[2] =

				    BE_64(tmpsha256.zc_word[2]);

				drrw->drr_key.ddk_cksum.zc_word[3] =

				    BE_64(tmpsha256.zc_word[3]);

				drrw->drr_checksumtype = ZIO_CHECKSUM_SHA256;

				drrw->drr_checksumflags = DRR_CHECKSUM_DEDUP;

			}

			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_key.ddk_cksum, drrw->drr_key.ddk_prop,

			    &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_checksumtype =

				    drrw->drr_checksumtype;

				wbr_drrr->drr_checksumflags =

				    drrw->drr_checksumtype;

				wbr_drrr->drr_key.ddk_cksum =

				    drrw->drr_key.ddk_cksum;

				wbr_drrr->drr_key.ddk_prop =

				    drrw->drr_key.ddk_prop;

				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);

}

/*

 * Routines for dealing with the AVL tree of fs-nvlists

 */

typedef struct fsavl_node {

	avl_node_t fn_node;

	nvlist_t *fn_nvfs;

	char *fn_snapname;

	uint64_t fn_guid;

} fsavl_node_t;

static int

fsavl_compare(const void *arg1, const void *arg2)

{

	const fsavl_node_t *fn1 = arg1;

	const fsavl_node_t *fn2 = arg2;

	if (fn1->fn_guid > fn2->fn_guid)

		return (+1);

	else if (fn1->fn_guid < fn2->fn_guid)

		return (-1);

	else

		return (0);

}

/*

 * Given the GUID of a snapshot, find its containing filesystem and

 * (optionally) name.

 */

static nvlist_t *

fsavl_find(avl_tree_t *avl, uint64_t snapguid, char **snapname)

{

	fsavl_node_t fn_find;

	fsavl_node_t *fn;

	fn_find.fn_guid = snapguid;

	fn = avl_find(avl, &fn_find, NULL);

	if (fn) {

		if (snapname)

			*snapname = fn->fn_snapname;

		return (fn->fn_nvfs);

	}

	return (NULL);

}

static void

fsavl_destroy(avl_tree_t *avl)

{

	fsavl_node_t *fn;

	void *cookie;

	if (avl == NULL)

		return;

	cookie = NULL;

	while ((fn = avl_destroy_nodes(avl, &cookie)) != NULL)

		free(fn);

	avl_destroy(avl);

	free(avl);

}

/*

 * Given an nvlist, produce an avl tree of snapshots, ordered by guid

 */

static avl_tree_t *

fsavl_create(nvlist_t *fss)

{

	avl_tree_t *fsavl;

	nvpair_t *fselem = NULL;

	if ((fsavl = malloc(sizeof (avl_tree_t))) == NULL)