/* Program running GMP-ECM factoring processes in parallel, minimizing
   simultaneous step2.

Copyright 2005, 2006, 2007 Torbjorn Granlund.

This program is free software; you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation; either version 2.1 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 General Public License for more details.

You should have received a copy of the GNU General Public License along with
this program; see the file COPYING.  If not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA. */

#include <sys/types.h>		/* for fork, wait4 */
#include <unistd.h>		/* for fork */
#include <fcntl.h>		/* for open */
#include <sys/wait.h>		/* for wait4 */
#include <sys/time.h>		/* for gettimeofday, wait4, setpriority */
#include <sys/resource.h>	/* for wait4, setpriority */
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <getopt.h>
#include <glob.h>
#include <signal.h>


/*
  To perform only step 1:  $ ./ecm -save toto B1 1 < composite
  Then to perform step 2:  $ ./ecm -resume toto B1 [B2]
*/


enum foo {NONE = 0, STEP1_RUNNING, STEP2_RUNNING, STEP1_DONE, STEP2_DONE};

struct jobinfo
{
  pid_t pid;
  enum foo state;
  char tmp_file[50];
  char cofac_file[50];
};

#define RUNNING(x) (x.state == STEP1_RUNNING || x.state == STEP2_RUNNING)

char *progname;

unsigned long
myrandom ()
{
  unsigned long ran;
  int fd;

  fd = open ("/dev/urandom", O_RDONLY);
  if (fd >= 0)
    {
      unsigned char buf[sizeof (long)];
      size_t nread;
      nread = read (fd, buf, sizeof (long));
      if (nread != sizeof (long))
	goto stupid;
      ran = buf[0] + (buf[1] << 8) + (buf[2] << 16) + (buf[3] << 24);
      if (sizeof (long) > 4)
	{
	  unsigned long ran2;
	  ran2 = buf[4] + (buf[5] << 8);
	  ran += (ran2 << 31) << 1;
	}
      close (fd);
      return ran;
    }
  else
    {
      static int flag = 0;
    stupid:
      if (flag == 0)
	{
	  struct timeval tp;
	  gettimeofday (&tp, NULL);
	  srand48 ((tp.tv_sec << 16) + tp.tv_usec + getpid ());
	  flag = 1;
	}
      ran = mrand48 ();
    }
  return ran;
}

char *
pathfind (const char *command)
{
  char *path, *p, *buf;
  int len, clen;

  clen = strlen (command);

  path = getenv ("PATH");
  if (path == NULL)
    abort ();

  buf = malloc (strlen (path) + clen + 2);

  for (;;)
    {
      p = strchr (path, ':');
      if (p == NULL)
	len = strlen (path);
      else
	len = p - path;

      memcpy (buf, path, len);
      if (buf[len - 1] != '/')
	{
	  buf[len] = '/';
	  memcpy (buf + len + 1, command, clen + 1);
	}
      else
	{
	  memcpy (buf + len, command, clen + 1);
	}
      if (access (buf, X_OK) == 0)
	return buf;

      path += len + 1;
    }

  free (buf);
  return NULL;
}

#define BUFSIZE 65536

enum algo
{
  ALGO_PM1,
  ALGO_PP1,
  ALGO_ECM
};

int B1_opt, B2_opt, algo, helpme, verbose, nproc_opt, mail_opt, ns2_opt;

static struct option longopts[] =
  {
    { "B1",     required_argument,  &B1_opt,     1 },
    { "B2",     required_argument,  &B2_opt,     1 },
    { "pm1",    no_argument,        &algo,       ALGO_PM1 },
    { "pp1",    no_argument,        &algo,       ALGO_PP1 },
    { "ecm",    no_argument,        &algo,       ALGO_ECM },
    { "help",   no_argument,        &helpme,     1 },
    { "mail",   required_argument,  &mail_opt,   0 },
    { "verbose",no_argument,        &verbose,    1 },
    { "n",      required_argument,  &nproc_opt,  1 },
    { "nproc",  required_argument,  &nproc_opt,  1 },
    { "n2",     required_argument,  &ns2_opt,    1 },
    { "nstep2", required_argument,  &ns2_opt,    1 },
    { NULL,     0,                  NULL,        0 }
  };

int sortfunc_exp (const void *, const void *);
int sortfunc_cf (const void *, const void *);
void cleanup_at_sig (int);

struct jobinfo *gjiv;
int gnprocs;

int
main (int argc, char *argv[], char *envp[])
{
  int nprocs, max_procs_in_step2;
  int i;
  double B1, B2;
  struct jobinfo *jiv;
  pid_t pid;
  int wstat;
  char *tmpdir;
  int (*result_chan)[2];		/* for reading output from passed */
  int n_running_procs;
  int fd;
  char *ecmfactor;
  char buf[BUFSIZE + 1];
  char sigma[21];
  int next_cofac_i;
  size_t nread;
  struct rusage rus;
  unsigned used_ms;
  char *mailaddr;
  int pass;

  nprocs = 1;			/* default */
  max_procs_in_step2 = 0;	/* bad value used as flag */

  B1 = 0;

  progname = argv[0];

  while (getopt_long_only (argc, argv, "", longopts, NULL) != -1)
    {
      if (B1_opt)
	{
	  B1 = strtod (optarg, 0);
	  B2 = B1 * 400;
	  B1_opt = 0;
	}
      else if (B2_opt)
	{
	  B2 = strtod (optarg, 0);
	  B2_opt = 0;
	}
      else if (mail_opt)
	{
	  mailaddr = optarg;
	  mail_opt = 0;
	}
      else if (nproc_opt)
	{
	  nprocs = strtoul (optarg, 0, 0);
	  nproc_opt = 0;
	}
      else if (ns2_opt)
	{
	  max_procs_in_step2 = strtoul (optarg, 0, 0);
	  ns2_opt = 0;
	}
      else if (helpme)
	{
	  printf ("usage: %s [-help] [-B1=VAL] [-B2=VAL] [-pm1 | -pp1 | -ecm] "filespec"\n", argv[0]);
	  printf ("   filespec is one shell-type files expression, such as \"c*\"\n");
	  printf ("      which specifies the cofactor files\n");
	  exit (0);
	}
    }

  argc -= optind;
  argv += optind;

  if (max_procs_in_step2 == 0)
    max_procs_in_step2 = nprocs / 2;

  if (argc == 0)
    {
      printf ("No files to factor\n");
      exit (0);
    }

  if (B1 == 0)
    {
      fprintf (stderr, "%s: missing B1 value\n", progname);
      exit (1);
    }

  next_cofac_i = 0;

  ecmfactor = pathfind ("ecmfactor");

  result_chan = malloc (nprocs * sizeof (int [2]));

  jiv = malloc (nprocs * sizeof (struct jobinfo));
  for (i = 0; i < nprocs; i++)
    {
      jiv[i].pid = 0;
      jiv[i].state = NONE;
      pipe (result_chan[i]);
    }

  tmpdir = getenv ("TMPDIR");
  if (tmpdir == NULL)
    tmpdir = "/tmp";

  n_running_procs = 0;

  pass = 0;

  gjiv = jiv;			/* reachable from signal handler */
  gnprocs = nprocs;		/* reachable from signal handler */

  signal (SIGHUP, cleanup_at_sig);
  signal (SIGILL, cleanup_at_sig);
  signal (SIGSEGV, cleanup_at_sig);
  signal (SIGINT, cleanup_at_sig);
  signal (SIGQUIT, cleanup_at_sig);
  signal (SIGTERM, cleanup_at_sig);

  setpriority (PRIO_PROCESS, 0, 20);

  for (;;)
    {
      int n_step2_procs = 0;
      char B1arg[15], B2arg[20];
      glob_t cofac_files;

      sprintf (B1arg, "%.0f", B1);
      sprintf (B2arg, "%.0f", B2);

      for (i = 0; i < nprocs; i++)
	{
	  n_step2_procs += (jiv[i].state == STEP2_RUNNING);
	}
      for (i = 0; i < nprocs; i++)
	{
	  if (! RUNNING (jiv[i]))
	    {
	      if (jiv[i].state == STEP1_DONE)
		{
		  if (n_step2_procs >= max_procs_in_step2)
		    continue;
		  printf ("starting step 2 job %02d for %-15s B2=%s\n", i, jiv[i].cofac_file, B2arg);
		  fflush (stdout);
		  pid = fork ();
		  if (pid == 0)
		    {
		      /* Child */
		      char *outv[6], **op = outv;
		      dup2 (result_chan[i][1], 1);
		      close (result_chan[i][0]);
		      close (result_chan[i][1]);
		      *op++ = "ecmfactor";
		      if (algo == ALGO_PM1)
			*op++ = "-pm1";
		      else if (algo == ALGO_PP1)
			*op++ = "-pp1";
		      *op++ = "-resume";
		      *op++ = jiv[i].tmp_file;
		      *op++ = B1arg;
		      *op++ = B2arg;
		      *op = NULL;
		      execve (ecmfactor, outv, envp);
		      fprintf (stderr, "cannot execute %s\n", ecmfactor);
		      abort ();
		    }
		  n_running_procs++;
		  jiv[i].pid = pid;
		  jiv[i].state = STEP2_RUNNING;
		  continue;
		}

	      for (;;)
		{
		  char filename[50];
		  if (next_cofac_i == 0)
		    {
		      /* We ran out of globbed file names, alternatively, the
			 program just started.  Glob. */
		      glob(argv[0], GLOB_NOSORT, NULL, &cofac_files);
		      pass++;
		      printf ("PASS %d: %d cofactor files\n", pass, cofac_files.gl_pathc);
		      fflush (stdout);
		      if (cofac_files.gl_pathc == 0)
			exit (0);
		      qsort (cofac_files.gl_pathv, cofac_files.gl_pathc,
			     sizeof (char *), sortfunc_cf);
		      sprintf (sigma, "%lu", myrandom ());
		    }
		  strcpy (filename, cofac_files.gl_pathv[next_cofac_i]);
		  next_cofac_i++;
		  if (next_cofac_i == cofac_files.gl_pathc)
		    {
		      globfree (&cofac_files);
		      next_cofac_i = 0;
		    }
		  fd = open (filename, O_RDONLY);
		  if (fd != -1)
		    {
		      strcpy (jiv[i].cofac_file, filename);
		      break;
		    }
		  usleep (50000);
		}

	      sprintf (jiv[i].tmp_file, "%s/ecm-save-%u", tmpdir, i);
	      unlink (jiv[i].tmp_file);
	      printf ("starting step 1 job %02d for %-15s B1=%s, s=%s\n", i, jiv[i].cofac_file, B1arg, sigma);
	      fflush (stdout);
	      pid = fork ();
	      if (pid == 0)
		{
		  /* Child */
		  char *outv[8], **op = outv;
		  dup2 (result_chan[i][1], 1);
		  close (result_chan[i][0]);
		  close (result_chan[i][1]);
		  dup2 (fd, 0);
		  close (fd);
		  *op++ = "ecmfactor";
		  if (algo == ALGO_PM1)
		    *op++ = "-pm1";
		  else if (algo == ALGO_PP1)
		    *op++ = "-pp1";
		  *op++ = "-save";
		  *op++ = jiv[i].tmp_file;
		  *op++ = "-sigma";
		  *op++ = sigma;
		  *op++ = B1arg;
		  *op++ = "1";
		  *op = NULL;
		  execve (ecmfactor, outv, envp);
		  fprintf (stderr, "cannot execute %s\n", ecmfactor);
		  abort ();
		}
	      close (fd);
	      n_running_procs++;
	      jiv[i].pid = pid;
	      jiv[i].state = STEP1_RUNNING;
	      continue;
	    }
	}

      pid = wait4 (0, &wstat, 0,  &rus);
      if (pid == -1)
	{
	  fprintf (stderr, "wait returned error %d\n", errno);
	  abort ();
	}
      n_running_procs--;

      used_ms = rus.ru_utime.tv_sec * 1000 + rus.ru_utime.tv_usec / 1000;

      if (WIFSIGNALED (wstat))
	{
	  nread = read (result_chan[i][0], buf, BUFSIZE);
	  fprintf (stderr, "*** child got signal %d\n", WTERMSIG (wstat));
	  fwrite (buf, 1, nread, stderr);
	  exit (1);		/* FIXME: kill all, cleanup, and exit */
	}
      for (i = 0; i < nprocs; i++)
	{
	  if (jiv[i].pid == pid)
	    goto yee;
	}
      abort ();
    yee:
      if (jiv[i].state == STEP1_RUNNING)
	{
	  nread = read (result_chan[i][0], buf, BUFSIZE);
	  jiv[i].state = STEP1_DONE;
	  printf ("finished step 1 job %02d for %-15s (used %u ms)\n", i, jiv[i].cofac_file, used_ms);
	  fflush (stdout);
	}
      else if (jiv[i].state == STEP2_RUNNING)
	{
	  unlink (jiv[i].tmp_file);
	  nread = read (result_chan[i][0], buf, BUFSIZE);
	  jiv[i].state = STEP2_DONE;
	  printf ("finished step 2 job %02d for %-15s (used %u ms)\n", i, jiv[i].cofac_file, used_ms);
	  fflush (stdout);
	}
      else
	abort ();

      char *s;
      buf[nread] = 0;
      s = strstr (buf, "Found");

      if (s != NULL)
	{
	  char tofilename[100];

	  printf (">>> Found new factor for %s <<<\n", jiv[i].cofac_file);
	  fflush (stdout);

	  sprintf (tofilename, "../cracked/%s", jiv[i].cofac_file);
	  fd = open (tofilename, O_WRONLY | O_CREAT, 0600);
	  write (fd, buf, nread);
	  close (fd);
	  unlink (jiv[i].cofac_file);

	  if (mail_opt)
	    {
	      FILE *fs;
	      char *popen_str;

	      popen_str = malloc (strlen (mailaddr + 100));
	      sprintf (popen_str, "mail -s \"NEW FACTOR\" %s", mailaddr);
	      fs = popen (popen_str, "w");
	      free (popen_str);
	      fwrite (buf, nread, 1, fs);
	      pclose (fs);
	    }
	}
    }
}

int
sortfunc_exp (const void *aa, const void *ba)
{
  const char *a = *(char **) aa;
  const char *b = *(char **) ba;

  a = strchr (a, ',') + 1;
  b = strchr (b, ',') + 1;
  return strcmp (a, b);
}

int
sortfunc_cf (const void *aa, const void *ba)
{
  const char *a = *(char **) aa;
  const char *b = *(char **) ba;
  int ai, bi;
  const char *aplus, *bplus;

  aplus = strchr (a, '+');
  bplus = strchr (b, '+');

  if (aplus == NULL || bplus == NULL)
    return strcmp (a, b);

  ai = aplus - a;
  bi = bplus - b;

  if (ai != bi)
    return ai - bi;
  return strcmp (a, b);
}

void
cleanup_at_sig (int sig)
{
  int i;
  for (i = 0; i < gnprocs; i++)
    {
      kill (gjiv[i].pid, SIGKILL);
      unlink (gjiv[i].tmp_file);
    }
  signal (sig, SIG_DFL);
  kill (getpid (), sig);
}