/*
 * random.c - module to access kernel random sources.
 *
 * This file is part of zsh, the Z shell.
 *
 * Copyright (c) 2022 Clinton Bunch
 * All rights reserved.
 *
 * Permission is hereby granted, without written agreement and without
 * license or royalty fees, to use, copy, modify, and distribute this
 * software and to distribute modified versions of this software for any
 * purpose, provided that the above copyright notice and the following
 * two paragraphs appear in all copies of this software.
 *
 * In no event shall Clinton Bunch or the Zsh Development Group be liable
 * to any party for direct, indirect, special, incidental, or consequential
 * damages arising out of the use of this software and its documentation,
 * even if Clinton Bunch and the Zsh Development Group have been advised of
 * the possibility of such damage.
 *
 * Clinton Bunch and the Zsh Development Group specifically disclaim any
 * warranties, including, but not limited to, the implied warranties of
 * merchantability and fitness for a particular purpose.  The software
 * provided hereunder is on an "as is" basis, and Clinton Bunch and the
 * Zsh Development Group have no obligation to provide maintenance,
 * support, updates, enhancements, or modifications.
 *
 */

#include "random.mdh"
#include "random.pro"

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>

#include <stdbool.h>

#ifdef HAVE_SYS_RANDOM_H
#include <sys/random.h>
#endif

/* Simplify select URANDOM specific code */
#if !defined(HAVE_ARC4RANDOM) && !defined(HAVE_GETRANDOM)
#define USE_URANDOM
#endif

/* buffer to pre-load integers for SRANDOM to lessen the context switches */
uint32_t rand_buff[8];
static int buf_cnt = -1;

#ifdef USE_URANDOM
/* File descriptor for /dev/urandom */
int randfd = -1;
#endif /* USE_URANDOM */

/*
 *  Function takes an input buffer and converts the characters to hex and
 *  places them in an outbut buffer twice the size. Returns -1 if it can't.
 *  Returns the size of the output otherwise.
 */

/**/
static int
ztr_to_hex(char *out, size_t outlen, uint8_t *input, size_t inlen)
{
    int i = 0, j = 0;

    if (outlen < (inlen * 2 + 1))
	return -1;

    for(i = 0, j = 0; i < inlen; i++, j+=2) {
	if (j < outlen) {
	    sprintf((char *)(out + j),"%02X",input[i]);
	} else {
	    return -1;
	}
    }
    out[j]='\0';
    return j;
}

/**/
static ssize_t
getrandom_buffer(void *buf, size_t len)
{
    ssize_t ret;
    size_t  val    = 0;
    uint8_t *bufptr = buf;

    do {
#ifdef HAVE_ARC4RANDOM
	/* Apparently, the creaters of arc4random didn't believe in errors */
        arc4random_buf(buf,len);
        ret = len;
#elif defined(HAVE_GETRANDOM)
        ret=getrandom(bufptr,(len - val),0);
#else
        ret=read(randfd,bufptr,(len - val));
#endif
        if (ret < 0) {
	    if (errno != EINTR || errflag || retflag || breaks || contflag) {
		zwarn("Unable to get random data: %e", errno);
	        return -1;
	    }
        }
	bufptr += ret;
	val    += ret;
    } while (ret < len);
    return ret;
}

/*
 * Implements the getrandom builtin to return a string of random bytes of
 * user-specified length.  It either prints them to stdout or returns them
 * in a parameter passed on the command line.
 *
 */

/**/
static int
bin_getrandom(char *name, char **argv, Options ops, int func)
{

    zulong len      = 8;
    size_t byte_len = 0, outlen; /* buffer lengths */
    int    pad      = 0;

    bool integer_out = false;         /* boolean */
    int i, j;                       /*loop counters */

    /*maximum string length of a 32-bit integer + null */
    char int2str[11];

    uint8_t   *buf;
    uint32_t  *int_buf = NULL, int_val = 0; /* buffer for integer return */

    char *scalar = NULL, *arrname = NULL; /* parameter names */

    char *outbuff;  /* hex string or metafy'd output */
    char **array = NULL;   /* array value for -a */

    /* Vailues for -U *nd -L */
    zulong   upper = UINT32_MAX, lower = 0, diff = UINT32_MAX;
    uint32_t min   = 0, rej = 0;  /* Reject below min and count */
    bool     bound = false; /* boolean to set if upper or lower bound set */
    /* End of variable declarations */


    /* store out put in a scalar variable */
    if (OPT_ISSET(ops, 's')) {
	scalar = OPT_ARG(ops, 's');
	if (!isident(scalar)) {
	    zwarnnam(name, "argument to -s not an identifier: %s", scalar);
	    return 1;
	}
    }

    /* create array of descimal numbers */
    if (OPT_ISSET(ops,'a')) {
	arrname = OPT_ARG(ops, 'a');
	if (!isident(arrname)) {
	    zwarnnam(name,"argument to -a not an identifier: %s", arrname);
	    return 1;
	}
    }

    /* specify number of random bytes or integers to output */
    if (OPT_ISSET(ops, 'l')) {

	if (zstrtoul_underscore(OPT_ARG(ops, 'l'), &len)) {
	    if (len > 64 || len <= 0) {
	        zwarnnam(name,
			"length must be between 1 and 64 you specified: %d",
		        len);
	        return 1;
	    }
	} else {
	    zwarnnam(name, "Couldn't convert -l %s to a number",
		    OPT_ARG(ops, 'l'));
	    return 1;
	}
    }

    if (OPT_ISSET(ops, 'U')) {
	if (!zstrtoul_underscore(OPT_ARG(ops, 'U'), &upper)) {
	    zwarnnam(name, "Couldn't convert -U %s to a number",
		    OPT_ARG(ops, 'U'));
	    return 1;
        } else {
	    bound = true;
	}
    }

    if (OPT_ISSET(ops, 'L')) {
	if (!zstrtoul_underscore(OPT_ARG(ops, 'L'), &lower)) {
	    zwarnnam(name, "Couldn't convert -L %s to a number",
		    OPT_ARG(ops, 'L'));
	    return 1;
	} else {
	    bound = true;
	}
    }

    /* integer rather than byte output to array */
    if (OPT_ISSET(ops, 'i')) {
	if (OPT_ISSET(ops,'s') && len > 1) {
	    zwarnnam(name,"-i only makes sense with -s when 1ength is 1");
	    return 1;
	}
	integer_out = true;
    }

    if (bound) {
        if (scalar && len > 1) {
            zwarnnam(name,"Bounds only make sense with -s when length is 1");
	    return 1;
        }

	if (lower > upper) {
	    zwarnnam(name,"-U must be larger than -L.");
	    return 1;
	}

        diff = upper - lower;

	if(!diff) {
	    zwarnnam(name,"Upper and Lower bounds cannot be the same.");
	    return 1;
	}
    }


    if (!byte_len)
	byte_len = len;


    if (bound) {
	pad = len / 16 + 1;
        byte_len = (len + pad) * sizeof(uint32_t);
    }
    if (integer_out)
	byte_len = len * sizeof(uint32_t);

    if (byte_len > 256)
	byte_len=256;

    buf = zalloc(byte_len);

    if (getrandom_buffer(buf, byte_len) < 0) {
	zwarnnam(name,"Couldn't get random data");
	return 1;
    }

    /* Raw output or hex string */
    if (!integer_out && !bound && !arrname) {
        if (OPT_ISSET(ops, 'r')) {
	    outbuff = (char *) buf;
	    outlen  = len;
        } else {
	    outlen=len*2;
	    outbuff = (char*) zalloc(outlen+1);
	    ztr_to_hex(outbuff, outlen+1, buf, len);
        }

        if (scalar) {
	    setsparam(scalar, metafy(outbuff, outlen, META_DUP));
        } else {
            fwrite(outbuff,1,outlen,stdout);
	    if (!OPT_ISSET(ops, 'r'))
		printf("\n");
	}
        if (outbuff !=  (void *) buf) {
	    zfree(outbuff,outlen+1);
	}
	return 0;
    } else {
	if (OPT_ISSET(ops, 'r')) {
	    zwarnnam(name, "-r can't be specified with bounds or -i or -a");
	    return 1;
	}
    }

    if(arrname) {
	array = (char **) zalloc((len+1)*sizeof(char *));
	array[len] = NULL;
    }

    if(integer_out || bound)
	int_buf=(uint32_t *)buf;

    min = -diff % diff;

    j = 0; rej = 0;
    for(i = 0; i < len; i++) {
	if(integer_out) {
	    int_val = int_buf[i];
	} else if (!bound) {
	    int_val=buf[i];
	} else {
	    while (int_buf[j] < min ) { /*Reject */
		j++; rej++;
		if (j * sizeof(uint32_t) >= byte_len ){
		    if (getrandom_buffer(buf, byte_len) < 0) {
			zwarnnam(name,"Can't get enough random data");
			return 1;
		    }
		    j = 0;
		}
	    }
	    int_val = int_buf[j++] % diff + lower;
	}
	sprintf(int2str, "%u", int_val);
	if (arrname) {
	    array[i]=ztrdup(int2str);
	} else if (scalar && len == 1) {
	    setiparam(scalar, int_val);
	} else {
	    printf("%s ",int2str);
	}
    }
    if (arrname) {
	setaparam(arrname,array);
    } else if (!scalar) {
	printf("\n");
    }


    zfree(buf, byte_len);
    return 0;
}


/*
 * Provides for the SRANDOM parameter and returns  a unisgned 32-bit random
 * integer.
 */

/**/
static zlong
get_srandom(UNUSED(Param pm)) {

    if(buf_cnt < 0) {
	getrandom_buffer((void*) rand_buff,sizeof(rand_buff));
	buf_cnt=7;
    }
    return rand_buff[buf_cnt--];
}

/*
 * Implements  the mathfunc zrandom and returns a random floating-point
 * number between 0 and 1.  Does this by getting a random 32-bt integer
 * and dividing it by UINT32_MAX.
 *
 */

/**/
static mnumber
math_zrandom(UNUSED(char *name), UNUSED(int argc), UNUSED(mnumber *argv),
	     UNUSED(int id))
{
    mnumber ret;
    double r;

    r = random_real();
    if (r < 0) {
	zwarnnam(name,"Random Data Failed");
    }
    ret.type = MN_FLOAT;
    ret.u.d = r;

    return ret;
}

static struct builtin bintab[] = {
    BUILTIN("getrandom", 0, bin_getrandom, 0, 8, 0, "ria:s:l:%U:%L:%", NULL),
};

static const struct gsu_integer srandom_gsu =
{ get_srandom, nullintsetfn, stdunsetfn };

static struct paramdef patab[] = {
    SPECIALPMDEF("SRANDOM", PM_INTEGER|PM_READONLY, &srandom_gsu,NULL,NULL),
};

static struct mathfunc mftab[] = {
    NUMMATHFUNC("zrandom", math_zrandom, 0, 0, 0),
};

static struct features module_features = {
    bintab, sizeof(bintab)/sizeof(*bintab),
    NULL, 0,
    mftab, sizeof(mftab)/sizeof(*mftab),
    patab, sizeof(patab)/sizeof(*patab),
    0
};

/**/
int
setup_(UNUSED(Module m))
{
#ifdef USE_URANDOM
    /* Check for the existence of /dev/urnadom */

    struct stat st;

    if (lstat("/dev/urandom",&st) < 0) {
	zwarn("No kernel random pool found.");
	return 1;
    }

    if (!(S_ISCHR(st.st_mode)) ) {
	zwarn("No kernel random pool found.");
	return 1;
    }
#endif /* USE_URANDOM */
    return 0;
}

/**/
int
features_(Module m, char ***features)
{
    *features = featuresarray(m, &module_features);
    return 0;
}

/**/
int
enables_(Module m, int **enables)
{
    return handlefeatures(m, &module_features, enables);
}

/**/
int
boot_(Module m)
{
#ifdef USE_URANDOM
    /*
     * Open /dev/urandom.  Here because of a weird issue on HP-UX 11.31
     * When opening in setup_ open returned 0.  In boot_, it returns
     * an unused file descriptor. Decided against ifdef HPUX as it works
     * here just as well for other platforms.
     *
     */

    int tmpfd=-1;

    if ((tmpfd = open("/dev/urandom", O_RDONLY)) < 0) {
	zwarn("Could not access kernel random pool");
	return 1;
    }
    randfd = movefd(tmpfd);
    addmodulefd(randfd,FDT_MODULE);
    if (randfd < 0) {
	zwarn("Could not access kernel random pool");
	return 1;
    }
#endif /* USE_URANDOM */
    return 0;
}

/**/
int
cleanup_(Module m)
{
    return setfeatureenables(m, &module_features, NULL);
}

/**/
int
finish_(UNUSED(Module m))
{
#ifdef USE_URANDOM
    if (randfd >= 0)
	zclose(randfd);
#endif /* USE_URANDOM */
    return 0;
}


/* Count the number of leading zeros, hopefully in gcc/clang by HW
 * instruction */
#if defined(__GNUC__) || defined(__clang__)
#define clz64(x) __builtin_clzll(x)
#else
#define clz64(x) _zclz64(x)

/**/
int
_zclz64(uint64_t x) {
    int n = 0;

    if (x == 0)
	return 64;

    if (!(x & 0xFFFFFFFF00000000)) {
	n+=32;
	x<<=32;
    }
    if (!(x & 0xFFFF000000000000)) {
	n+=16;
	x<<=16;
    }
    if (!(x & 0xFF00000000000000)) {
	n+=8;
	x<<=8;
    }
    if (!(x & 0xF000000000000000)) {
	n+=4;
	x<<=4;
    }
    if (!(x & 0xC000000000000000)) {
	n+=2;
	x<<=1;
    }
    if (!(x & 0x8000000000000000)) {
	n+=1;
    }
    return n;
}
#endif /* __GNU_C__ or __clang__ */

/**/
uint64_t
random64(void) {
    uint64_t r;

    if(getrandom_buffer(&r,sizeof(r)) < 0) {
	zwarn("zsh/random: Can't get sufficient random data");
	return 1; /* 0 will cause loop */
    }

    return r;
}

/* Following code is under the below copyright, despite changes for error
 * handling and removing GCCisms */

/*-
 * Copyright (c) 2014 Taylor R. Campbell
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

/*
 * Uniform random floats: How to generate a double-precision
 * floating-point numbers in [0, 1] uniformly at random given a uniform
 * random source of bits.
 *
 * See <http://mumble.net/~campbell/2014/04/28/uniform-random-float>
 * for explanation.
 *
 * Updated 2015-02-22 to replace ldexp(x, <constant>) by x * ldexp(1,
 * <constant>), since glibc and NetBSD libm both have slow software
 * bit-twiddling implementations of ldexp, but GCC can constant-fold
 * the latter.
 */

#include <math.h>
#include <stdint.h>

/*
 * random_real: Generate a stream of bits uniformly at random and
 * interpret it as the fractional part of the binary expansion of a
 * number in [0, 1], 0.00001010011111010100...; then round it.
 */

/**/
double
random_real(void)
{
	int exponent = 0;
	uint64_t significand = 0;
	uint64_t r = 0;
	unsigned shift;

	/*
	 * Read zeros into the exponent until we hit a one; the rest
	 * will go into the significand.
	 */
	while (significand == 0) {
		exponent -= 64;

		/* Get random64 and check for error */
		errno = 0;
		significand = random64();
		if (errno)
		    return -1;
		/*
		 * If the exponent falls below -1074 = emin + 1 - p,
		 * the exponent of the smallest subnormal, we are
		 * guaranteed the result will be rounded to zero.  This
		 * case is so unlikely it will happen in realistic
		 * terms only if random64 is broken.
		 */
		if (exponent < -1074)
			return 0;
	}

	/*
	 * There is a 1 somewhere in significand, not necessarily in
	 * the most significant position.  If there are leading zeros,
	 * shift them into the exponent and refill the less-significant
	 * bits of the significand.  Can't predict one way or another
	 * whether there are leading zeros: there's a fifty-fifty
	 * chance, if random64 is uniformly distributed.
	 */
	shift = clz64(significand);
	if (shift != 0) {

		errno = 0;
		r = random64();
		if (errno)
		    return -1;

		exponent -= shift;
		significand <<= shift;
		significand |= (r >> (64 - shift));
	}

	/*
	 * Set the sticky bit, since there is almost surely another 1
	 * in the bit stream.  Otherwise, we might round what looks
	 * like a tie to even when, almost surely, were we to look
	 * further in the bit stream, there would be a 1 breaking the
	 * tie.
	 */
	significand |= 1;

	/*
	 * Finally, convert to double (rounding) and scale by
	 * 2^exponent.
	 */
	return ldexp((double)significand, exponent);
}