/*
  webid-oidc, implementation of the Solid specification
  Copyright (C) 2020, 2021  Vivien Kraus

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as
  published by the Free Software Foundation, either version 3 of the
  License, or (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif /* not HAVE_CONFIG_H */

#include <string.h>
#include <nettle/yarrow.h>
#include <threads.h>
#include <assert.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <webid-oidc/random.h>
#include <gettext.h>
#define _(s) dgettext (PACKAGE, s)

struct webid_oidc_random_context
{
  struct yarrow256_ctx ctx;
  enum webid_oidc_random_context_spawn_flags flags;
  enum webid_oidc_random_context_state state;
  unsigned int state_progress;
  uint8_t seed[YARROW256_SEED_FILE_SIZE];
  mtx_t mutex;
};

struct webid_oidc_random_context *
webid_oidc_random_context_alloc (void)
{
  struct webid_oidc_random_context *ret =
    malloc (sizeof (struct webid_oidc_random_context));
  if (ret && mtx_init (&(ret->mutex), mtx_plain) != thrd_success)
    {
      free (ret);
      ret = NULL;
    }
  if (ret)
    {
      yarrow256_init (&(ret->ctx), 0, NULL);
      ret->flags = 0;
      ret->state = WEBID_OIDC_RANDOM_CONTEXT_WAIT_LOCK_FILE;
    }
  return ret;
}

void
webid_oidc_random_context_free (struct webid_oidc_random_context *ctx)
{
  if (ctx)
    {
      mtx_destroy (&(ctx->mutex));
    }
  free (ctx);
}

enum webid_oidc_random_context_state
webid_oidc_random_context_get_state (const struct webid_oidc_random_context
				     *ctx)
{
  return ctx->state;
}

int
webid_oidc_random_context_get_wait_bytes (const struct
					  webid_oidc_random_context *ctx,
					  size_t max, size_t start,
					  uint8_t * bytes, size_t *n_bytes)
{
  if (ctx->state == WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_WRITE)
    {
      if (start >= YARROW256_SEED_FILE_SIZE)
	{
	  *n_bytes = 0;
	}
      else
	{
	  if (max + start > YARROW256_SEED_FILE_SIZE)
	    {
	      max = YARROW256_SEED_FILE_SIZE - start;
	    }
	  memcpy (bytes, ctx->seed + start, max);
	  *n_bytes = YARROW256_SEED_FILE_SIZE - start;
	}
      return 1;
    }
  return 0;
}

int
webid_oidc_random_context_ok (struct webid_oidc_random_context *ctx)
{
  switch (ctx->state)
    {
    case WEBID_OIDC_RANDOM_CONTEXT_WAIT_LOCK_FILE:
      ctx->state = WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_READ;
      ctx->state_progress = 0;
      break;
    case WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_WRITE:
      ctx->state = WEBID_OIDC_RANDOM_CONTEXT_WAIT_UNLOCK_FILE;
      break;
    case WEBID_OIDC_RANDOM_CONTEXT_WAIT_UNLOCK_FILE:
      ctx->state = WEBID_OIDC_RANDOM_CONTEXT_READY;
      break;
    default:
      return 0;
    }
  return 1;
}

int
webid_oidc_random_context_bytes (struct webid_oidc_random_context *ctx,
				 size_t n_bytes, const uint8_t * bytes)
{
  switch (ctx->state)
    {
    case WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_READ:
      if (n_bytes >= (YARROW256_SEED_FILE_SIZE - ctx->state_progress))
	{
	  n_bytes = (YARROW256_SEED_FILE_SIZE - ctx->state_progress);
	}
      memcpy (ctx->seed + ctx->state_progress, bytes, n_bytes);
      ctx->state_progress += n_bytes;
      if (ctx->state_progress == YARROW256_SEED_FILE_SIZE)
	{
	  yarrow256_seed (&(ctx->ctx), YARROW256_SEED_FILE_SIZE, ctx->seed);
	  yarrow256_random (&(ctx->ctx), YARROW256_SEED_FILE_SIZE, ctx->seed);
	  ctx->state = WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_WRITE;
	}
      break;
    default:
      return 0;
    }
  return 1;
}

int
webid_oidc_random_context_spawn (struct webid_oidc_random_context *parent,
				 struct webid_oidc_random_context *child,
				 enum webid_oidc_random_context_spawn_flags
				 flags)
{
  if (!webid_oidc_random_context_get
      (parent, YARROW256_SEED_FILE_SIZE, child->seed))
    {
      return 0;
    }
  child->flags = flags;
  child->state = WEBID_OIDC_RANDOM_CONTEXT_READY;
  yarrow256_seed (&(child->ctx), YARROW256_SEED_FILE_SIZE, child->seed);
  return 1;
}

int
webid_oidc_random_context_get (struct webid_oidc_random_context *ctx,
			       size_t request_size, uint8_t * data)
{
  if (ctx->state == WEBID_OIDC_RANDOM_CONTEXT_READY)
    {
      if (!(ctx->flags & WEBID_OIDC_RANDOM_CONTEXT_THREAD_UNSAFE))
	{
	  if (mtx_lock (&(ctx->mutex)) != thrd_success)
	    {
	      assert (0);
	    }
	}
      yarrow256_random (&(ctx->ctx), request_size, data);
      if (!(ctx->flags & WEBID_OIDC_RANDOM_CONTEXT_THREAD_UNSAFE))
	{
	  if (mtx_unlock (&(ctx->mutex)) != thrd_success)
	    {
	      assert (0);
	    }
	}
    }
  return (ctx->state == WEBID_OIDC_RANDOM_CONTEXT_READY);
}

static struct webid_oidc_random_context *global_ctx = NULL;

void
webid_oidc_random_init (void)
{
  char *cache_dir, *pkg_cache_dir, *filename;
  char *home = getenv ("HOME");
  char *xdg_cache_home = getenv ("XDG_CACHE_HOME");
  char *application = getenv ("WEBID_OIDC_APPLICATION_NAME");
  static const char *default_application = PACKAGE;
  FILE *seed_file;
  FILE *system_rng_file;
  struct flock lock;
  if (global_ctx)
    {
      webid_oidc_random_context_free (global_ctx);
      global_ctx = NULL;
    }
  global_ctx = webid_oidc_random_context_alloc ();
  if (!global_ctx)
    {
      fprintf (stderr, _("Could not set the global random generator up.\n"));
      abort ();
    }
  if (!application)
    {
      application = (char *) default_application;
    }
  if (xdg_cache_home)
    {
      cache_dir = malloc (strlen (xdg_cache_home) + 1);
      if (!cache_dir)
	{
	  fprintf (stderr,
		   _
		   ("Could not set the global random generator up: out of memory.\n"));
	  abort ();
	}
      strcpy (cache_dir, xdg_cache_home);
    }
  else if (home)
    {
      cache_dir = malloc (strlen (home) + strlen ("/.cache") + 1);
      if (!cache_dir)
	{
	  fprintf (stderr,
		   _
		   ("Could not set the global random generator up: out of memory.\n"));
	  abort ();
	}
      strcpy (cache_dir, home);
      strcat (cache_dir, "/.cache");
    }
  else
    {
      cache_dir = malloc (strlen ("/var/cache") + 1);
      if (!cache_dir)
	{
	  fprintf (stderr,
		   _
		   ("Could not set the global random generator up: out of memory.\n"));
	  abort ();
	}
      strcpy (cache_dir, "/var/cache");
    }
  pkg_cache_dir = malloc (strlen (cache_dir)
			  + strlen ("/") + strlen (application) + 1);
  if (!pkg_cache_dir)
    {
      fprintf (stderr,
	       _
	       ("Could not set the global random generator up: out of memory.\n"));
      abort ();
    }
  strcpy (pkg_cache_dir, cache_dir);
  strcat (pkg_cache_dir, "/");
  strcat (pkg_cache_dir, application);
  filename = malloc (strlen (pkg_cache_dir) + strlen ("/seed") + 1);
  if (!filename)
    {
      fprintf (stderr,
	       _
	       ("Could not set the global random generator up: out of memory.\n"));
      abort ();
    }
  strcpy (filename, pkg_cache_dir);
  strcat (filename, "/seed");
  seed_file = fopen (filename, "r+");
  if (!seed_file)
    {
      fprintf (stderr,
	       _
	       ("Warning: could not open the seed file, maybe the parent directory is missing...\n"));
      if (mkdir (cache_dir, 0777) != 0)
	{
	  fprintf (stderr,
		   _("Warning: could not create the cache directory '%s'.\n"),
		   cache_dir);
	  perror (_("when creating the cache directory"));
	}
      if (mkdir (pkg_cache_dir, 0777) != 0)
	{
	  fprintf (stderr,
		   _
		   ("Warning: could not create the package cache directory '%s'.\n"),
		   pkg_cache_dir);
	  perror (_("when creating the package cache directory"));
	}
      seed_file = fopen (filename, "w+");
    }
  if (!seed_file)
    {
      fprintf (stderr, _("Could not open the seed file '%s'.\n"), filename);
      perror (_("when opening the seed file"));
      abort ();
    }
  system_rng_file = fopen ("/dev/random", "r");
  while (global_ctx->state != WEBID_OIDC_RANDOM_CONTEXT_READY)
    {
      int error;
      int c;
      uint8_t byte = 0;
      size_t i, n;
      switch (global_ctx->state)
	{
	case WEBID_OIDC_RANDOM_CONTEXT_WAIT_LOCK_FILE:
	  lock.l_type = F_WRLCK;
	  lock.l_whence = SEEK_SET;
	  lock.l_start = 0;
	  lock.l_len = 0;
	  error = fcntl (fileno (seed_file), F_SETLKW, &lock);
	  if (error != 0)
	    {
	      fprintf (stderr, _("Could not lock the seed file '%s'.\n"),
		       filename);
	      perror (_("when locking the seed file"));
	      abort ();
	    }
	  if (!webid_oidc_random_context_ok (global_ctx))
	    {
	      assert (0);
	    }
	  break;
	case WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_READ:
	  c = fgetc (seed_file);
	  if (c == EOF && system_rng_file == NULL)
	    {
	      fprintf (stderr,
		       _
		       ("Warning: the seed file '%s' is too short. This weakens the random number generator. Please write more random data in it.\n"),
		       filename);
	    }
	  else if (c == EOF)
	    {
	      c = fgetc (system_rng_file);
	    }
	  byte = c;
	  if (!webid_oidc_random_context_bytes (global_ctx, 1, &byte))
	    {
	      assert (0);
	    }
	  break;
	case WEBID_OIDC_RANDOM_CONTEXT_WAIT_FILE_WRITE:
	  if (fseek (seed_file, 0, SEEK_SET) != 0)
	    {
	      fprintf (stderr,
		       _("Could not update the seed file '%s'.\n"), filename);
	      perror (_("when rewinding the seed file"));
	      abort ();
	    }
	  if (!webid_oidc_random_context_get_wait_bytes
	      (global_ctx, 0, 0, NULL, &n))
	    {
	      assert (0);
	    }
	  for (i = 0; i < n; i++)
	    {
	      size_t remaining;
	      if (!webid_oidc_random_context_get_wait_bytes
		  (global_ctx, 1, i, &byte, &remaining))
		{
		  assert (0);
		}
	      assert (i + remaining == n);
	      fputc (byte, seed_file);
	    }
	  if (!webid_oidc_random_context_ok (global_ctx))
	    {
	      assert (0);
	    }
	  break;
	case WEBID_OIDC_RANDOM_CONTEXT_WAIT_UNLOCK_FILE:
	  lock.l_type = F_UNLCK;
	  lock.l_whence = SEEK_SET;
	  lock.l_start = 0;
	  lock.l_len = 0;
	  error = fcntl (fileno (seed_file), F_SETLKW, &lock);
	  if (error != 0)
	    {
	      fprintf (stderr, _("Could not unlock the seed file '%s'.\n"),
		       filename);
	      perror (_("when unlocking the seed file"));
	      abort ();
	    }
	  if (!webid_oidc_random_context_ok (global_ctx))
	    {
	      assert (0);
	    }
	  break;
	default:
	  assert (0);
	}
    }
  fclose (seed_file);
  fclose (system_rng_file);
}

void
webid_oidc_random (size_t request_size, uint8_t * data)
{
  static _Thread_local struct webid_oidc_random_context *thrd_ctx = NULL;
  if (thrd_ctx == NULL)
    {
      assert (global_ctx);
      thrd_ctx = webid_oidc_random_context_alloc ();
      if (thrd_ctx == NULL)
	{
	  fprintf (stderr,
		   _
		   ("Could not set the thread-local random generator up.\n"));
	  abort ();
	}
      if (!webid_oidc_random_context_spawn
	  (global_ctx, thrd_ctx, WEBID_OIDC_RANDOM_CONTEXT_THREAD_UNSAFE))
	{
	  fprintf (stderr,
		   _
		   ("The random module has not been initialized. Please call webid_oidc_random_init first.\n"));
	  abort ();
	}
    }
  if (!webid_oidc_random_context_get (thrd_ctx, request_size, data))
    {
      assert (0);
    }
}