/*
 *	UCW Library -- Strong random generator
 *
 *	(c) 2020--2022 Pavel Charvat <pchar@ucw.cz>
 *
 *	This software may be freely distributed and used according to the terms
 *	of the GNU Lesser General Public License.
 */

#undef LOCAL_DEBUG

#include <ucw/lib.h>
#include <ucw/random.h>
#include <ucw/threads.h>

#include <errno.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>

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

/*** getrandom() helpers ***/

#ifdef CONFIG_UCW_GETRANDOM
static bool strongrand_getrandom_detect(void)
{
  static int static_detected;
  int detected;
  ucwlib_lock();
  detected = static_detected;
  if (detected)
    {
      ucwlib_unlock();
      return detected > 0;
    }
  byte buf[1];
  int err;
  while (1)
    {
      int n = getrandom(buf, 1, GRND_NONBLOCK);
      if (n >= 0)
	{
	  ASSERT(n == 1);
	  detected = 1;
	  break;
	}
      else if (errno != EINTR)
	{
	  err = errno;
	  detected = -1;
	  break;
	}
    }
  static_detected = detected;
  ucwlib_unlock();
  if (detected > 0)
    {
      DBG("RANDOM: Kernel supports getrandom()");
      return true;
    }
  else if (err == ENOSYS)
    {
      DBG("RANDOM: Kernel does not support getrandom()");
      return false;
    }
  else
    {
      // We print an error also for EAGAIN -- for urandom it should be possible only during early boot.
      msg(L_ERROR, "RANDOM: Failed to call getrandom(): %s -> assuming no support", strerror(err));
      return false;
    }
}

bool strongrand_getrandom_mem_try(void *buf, size_t size)
{
  if (!strongrand_getrandom_detect())
    return false;
  while (size)
    {
      ssize_t n = getrandom(buf, size, 0);
      if (n < 0)
	{
	  if (errno == EINTR)
	    continue;
	  die("RANDOM: Failed to call getrandom(): %m");
	}
      buf = (byte *)buf + n;
      size -= n;
    }
  return true;
}
#else
bool strongrand_getrandom_mem_try(void *buf UNUSED, size_t size UNUSED)
{
  return false;
}
#endif

/*** Generic interface ***/

void strongrand_close(struct strongrand *c)
{
  if (c && c->close)
    c->close(c);
}

void strongrand_reset(struct strongrand *c)
{
  if (c->reset)
    c->reset(c);
}

void strongrand_mem(struct strongrand *c, void *ptr, size_t size)
{
  ASSERT(c->read);
  while (size)
    {
      int n = MIN(size, 0x7fffffff);
      c->read(c, ptr, n);
      size -= n;
      ptr = (byte *)ptr + n;
    }
}

/*** Kernel random source ***/

struct strongrand_std {
  struct strongrand sr;
  int fd;		// Non-negative descriptor or special STRONGRAND_FD_x value
  uint flags;		// STRONGRAND_STD_x
  uint buf_size;	// Size of @buf or 0
  uint buf_avail;	// Remaining buffered bytes, <= @buf_size
  pid_t last_pid;	// Last known process id; maintained with STRONGRAND_STD_x_RESET
  byte buf[];		// Buffered randomness
};

// Use getrandom()
#define STRONGRAND_FD_KERNEL_GETRANDOM -1

// Not yet open /dev/[u]random
#define STRONGRAND_FD_KERNEL_DEVICE -2

static void strongrand_std_close(struct strongrand *sr)
{
  struct strongrand_std *srs = (struct strongrand_std *)sr;
  if (srs->fd >= 0)
    close(srs->fd);
  DBG("RANDOM[%s]: Closed", sr->name);
  xfree(srs);
}

static void strongrand_std_reset(struct strongrand *sr)
{
  struct strongrand_std *srs = (struct strongrand_std *)sr;
  srs->buf_avail = 0;
}

static void strongrand_std_open_device(struct strongrand_std *srs)
{
  ASSERT(srs->fd == STRONGRAND_FD_KERNEL_DEVICE);
  ASSERT(srs->flags & (STRONGRAND_STD_URANDOM | STRONGRAND_STD_RANDOM));
  const char *path = (srs->flags & STRONGRAND_STD_URANDOM) ? "/dev/urandom" : "/dev/random";
  int fd = open(path, O_RDONLY);
  if (fd < 0)
    die("RANDOM[%s]: Cannot open %s: %m", srs->sr.name, path);
  srs->fd = fd;
  DBG("RANDOM[%s]: Opened device, path=%s, fd=%d", srs->sr.name, path, fd);
}

static void strongrand_std_read_unbuffered(struct strongrand *sr, byte *ptr, int size)
{
  struct strongrand_std *srs = (struct strongrand_std *)sr;
#ifdef CONFIG_UCW_GETRANDOM
  if (srs->fd == STRONGRAND_FD_KERNEL_GETRANDOM)
    {
      while (size)
	{
	  int n = getrandom(ptr, size, (srs->flags & STRONGRAND_STD_RANDOM) ? GRND_RANDOM : 0);
	  if (n < 0)
	    {
	      if (errno == EINTR || errno == EAGAIN)
		continue;
	      die("RANDOM[%s]: Failed to read %u bytes with getrandom(): %m", srs->sr.name, size);
	    }
	  DBG("RANDOM[%s]: Read %u bytes with getrandom()", srs->sr.name, n);
	  ASSERT(n <= size);
	  ptr += n;
	  size -= n;
	}
      return;
    }
#endif
  if (srs->fd < 0)
    {
      ASSERT(srs->flags & STRONGRAND_STD_DELAYED);
      strongrand_std_open_device(srs);
    }
  while (size)
    {
      int n = read(srs->fd, ptr, size);
      if (n < 0)
	{
	  if (errno == EINTR || errno == EAGAIN)
	    continue;
	  die("RANDOM[%s]: Failed to read %u bytes: %m", srs->sr.name, size);
	}
      ASSERT(n <= size);
      DBG("RANDOM[%s]: Read %u bytes", srs->sr.name, n);
      ptr += n;
      size -= n;
    }
}

static void strongrand_std_read_buffered(struct strongrand *sr, byte *ptr, int size)
{
  struct strongrand_std *srs = (struct strongrand_std *)sr;
  if (srs->flags & (STRONGRAND_STD_AUTO_RESET | STRONGRAND_STD_ASSERT_RESET))
    {
      pid_t pid = getpid();
      if (pid != srs->last_pid)
	{
	  DBG("RANDOM[%s]: Detected changed pid from %u to %u", sr->name, (uint)srs->last_pid, (uint)pid);
	  srs->last_pid = pid;
	  if (srs->flags & STRONGRAND_STD_ASSERT_RESET)
	    ASSERT(!srs->buf_avail); // Note: Does not catch all cases of wrong usage (already read multiples of buf_size)
	  srs->buf_avail = 0;
	}
    }
  while (size)
    {
      if (!srs->buf_avail)
	{
	  DBG("RANDOM[%s]: Refilling buffer", srs->sr.name);
	  strongrand_std_read_unbuffered(sr, srs->buf, srs->buf_size);
	  srs->buf_avail = srs->buf_size;
	}
      uint n = MIN(srs->buf_avail, (uint)size);
      memcpy(ptr, srs->buf + (srs->buf_size - srs->buf_avail), n);
      srs->buf_avail -= n;
      size -= n;
      ptr += n;
    }
}

static struct strongrand *strongrand_std_open_internal(int fd, const char *name, uint flags, uint buf_size)
{
  struct strongrand_std *srs = xmalloc(sizeof(*srs) + buf_size);
  bzero(srs, sizeof(*srs));
  srs->sr.name = name;
  srs->sr.close = strongrand_std_close;
  srs->sr.reset = buf_size ? strongrand_std_reset : NULL;
  srs->sr.read = buf_size ? strongrand_std_read_buffered : strongrand_std_read_unbuffered;
  srs->fd = fd;
  srs->flags = flags;
  srs->buf_size = buf_size;
  if (flags & (STRONGRAND_STD_AUTO_RESET | STRONGRAND_STD_ASSERT_RESET))
    srs->last_pid = getpid();
  return &srs->sr;
}

struct strongrand *strongrand_std_open(uint flags, uint buf_size)
{
  const char *name;
  if (flags & STRONGRAND_STD_URANDOM)
    {
      ASSERT(!(flags & STRONGRAND_STD_RANDOM));
      name = "urandom";
    }
  else
    {
      ASSERT(flags & STRONGRAND_STD_RANDOM);
      name = "random";
    }
#ifdef CONFIG_UCW_GETRANDOM
  if (strongrand_getrandom_detect())
    return strongrand_std_open_internal(STRONGRAND_FD_KERNEL_GETRANDOM, name, flags, buf_size);
#endif
  struct strongrand *sr = strongrand_std_open_internal(STRONGRAND_FD_KERNEL_DEVICE, name, flags, buf_size);
  if (!(flags & STRONGRAND_STD_DELAYED))
    strongrand_std_open_device((struct strongrand_std *)sr);
  return sr;
}