/* This file is part of GNU Pies.
   Copyright (C) 2008-2019 Sergey Poznyakoff

   GNU Pies 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 3, or (at your option)
   any later version.

   GNU Pies 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 GNU Pies.  If not, see <http://www.gnu.org/licenses/>. */

#include <pies.h>

#define SET_LIMIT_AS      0x0001
#define SET_LIMIT_CPU     0x0002
#define SET_LIMIT_DATA    0x0004
#define SET_LIMIT_FSIZE   0x0008
#define SET_LIMIT_NPROC   0x0010
#define SET_LIMIT_CORE    0x0020
#define SET_LIMIT_MEMLOCK 0x0040
#define SET_LIMIT_NOFILE  0x0080
#define SET_LIMIT_RSS     0x0100
#define SET_LIMIT_STACK   0x0200
#define SET_LIMIT_LOGINS  0x0400
#define SET_LIMIT_PRIO    0x0800

struct limits_rec
{
  unsigned set;
  rlim_t   limit_as;
  rlim_t   limit_cpu;
  rlim_t   limit_data;
  rlim_t   limit_fsize;
  rlim_t   limit_nproc;
  rlim_t   limit_core;
  rlim_t   limit_memlock;
  rlim_t   limit_nofile;
  rlim_t   limit_rss;
  rlim_t   limit_stack;
  int      limit_logins;
  int      limit_prio;  
};

int
do_set_limit (int rlimit, rlim_t limit)
{
  struct rlimit rlim;

  debug (1, ("setting limit %d to %lu", rlimit, (unsigned long) limit));
  rlim.rlim_cur = limit;
  rlim.rlim_max = limit;

  if (setrlimit(rlimit, &rlim))
    {
      //FIXME: arg?
      logfuncall ("setrlimit", NULL, errno);
      return 1;
    }
  return 0;
}

static int
set_prio (int prio)
{
  debug (2, ("Setting priority to %d", prio));
  if (setpriority (PRIO_PROCESS, 0, prio))
    {
      logfuncall ("setpriority", NULL, errno);
      return 1;
    }
  return 0;
}

/* Counts the number of user logins and check against the limit */
static int
check_logins (const char *name, int limit)
{
  logmsg (LOG_NOTICE, _("L limit is not implemented"));
  return 0;
}

int
set_limits (const char *name, struct limits_rec *lrec)
{
  int rc = 0;

  if (!lrec)
    return 0;

  debug (2, ("setting limits for %s", name));

#if defined(RLIMIT_AS)
  if (lrec->set & SET_LIMIT_AS)
    rc |= do_set_limit(RLIMIT_AS, lrec->limit_as);
#endif
#if defined(RLIMIT_CPU)
  if (lrec->set & SET_LIMIT_CPU)
    rc |= do_set_limit(RLIMIT_CPU, lrec->limit_cpu);
#endif
#if defined(RLIMIT_DATA)
  if (lrec->set & SET_LIMIT_DATA)
    rc |= do_set_limit(RLIMIT_DATA, lrec->limit_data);
#endif
#if defined(RLIMIT_FSIZE)
  if (lrec->set & SET_LIMIT_FSIZE) 
    rc |= do_set_limit(RLIMIT_FSIZE, lrec->limit_fsize);
#endif
#if defined(RLIMIT_NPROC)
  if (lrec->set & SET_LIMIT_NPROC)
    rc |= do_set_limit(RLIMIT_NPROC, lrec->limit_nproc);
#endif
#if defined(RLIMIT_CORE)
  if (lrec->set & SET_LIMIT_CORE)
    rc |= do_set_limit(RLIMIT_CORE, lrec->limit_core);
#endif
#if defined(RLIMIT_MEMLOCK)
  if (lrec->set & SET_LIMIT_MEMLOCK)
    rc |= do_set_limit(RLIMIT_MEMLOCK, lrec->limit_memlock);
#endif
#if defined(RLIMIT_NOFILE)
  if (lrec->set & SET_LIMIT_NOFILE)
    rc |= do_set_limit(RLIMIT_NOFILE, lrec->limit_nofile);
#endif
#if defined(RLIMIT_RSS)
  if (lrec->set & SET_LIMIT_RSS)
    rc |= do_set_limit(RLIMIT_RSS, lrec->limit_rss);
#endif
#if defined(RLIMIT_STACK)
  if (lrec->set & SET_LIMIT_STACK)
    rc |= do_set_limit(RLIMIT_STACK, lrec->limit_stack);
#endif
  if (lrec->set & SET_LIMIT_LOGINS)
    rc |= check_logins(name, lrec->limit_logins);
  if (lrec->set & SET_LIMIT_PRIO)
    rc |= set_prio(lrec->limit_logins);
  return rc;
}

int
getlimit (char **ptr, rlim_t *rlim, int mul)
{
  if (**ptr == '-')
    {
      *rlim = RLIM_INFINITY;
      ++*ptr;
    }
  else
    {
      unsigned long val;

      val = strtoul (*ptr, ptr, 10);
      if (val == 0)
	return 1;
      *rlim = val * mul;
    }
  return 0;
}

/* Parse limits string and fill appropriate fields in lrec.
 
   The string consists of _commands_, optionally separated by any amount
   of whitespace.  A command has the following form:
 
            [AaCcDdFfMmNnRrSsTtUuLlPp](-|[0-9]+)
 
   i.e. a letter followed by number or a dash. The latters stands for
   'unlimited'. Commands are interpreted as follows:
 
         Command   ulimit  setrlimit()     The limit it sets
                   option    arg
         -------------------------------------------------------------
            [Aa]   a       RLIMIT_AS        max address space (KB)
            [Cc]   c       RLIMIT_CORE      max core file size (KB)
            [Dd]   d       RLIMIT_DATA      max data size (KB)
            [Ff]   f       RLIMIT_FSIZE     Maximum filesize (KB)
            [Mm]   m       RLIMIT_MEMLOCK   max locked-in-memory address
                                            space (KB)
            [Nn]   n       RLIMIT_NOFILE    max number of open files
            [Rr]   r       RLIMIT_RSS       max resident set size (KB)
            [Ss]   s       RLIMIT_STACK     max stack size (KB)
            [Tt]   t       RLIMIT_CPU       max CPU time (MIN)
            [Uu]   u       RLIMIT_NPROC     max number of processes
            [Ll]   l       (none)           max number of logins (N/A)
            [Pp]   p       (none)           process priority -20..20
                                            (negative = high priority)
 */
int
parse_limits (limits_record_t *plrec, char *str, char **endp)
{
  int c;
  struct limits_rec *lrec = grecs_malloc (sizeof (*lrec));
  *plrec = lrec;
  lrec->set = 0;
  while ((c = *str++))
    {
      if (c == ' ' || c == '\t')
	continue;
      switch (c)
	{
	case 'a':
	case 'A':
	  /* RLIMIT_AS - max address space (KB) */
	  if (!getlimit (&str, &lrec->limit_as, 1024))
	    lrec->set |= SET_LIMIT_AS;
	  break;

	case 't':
	case 'T':
	  /* RLIMIT_CPU - max CPU time (MIN) */
	  if (!getlimit (&str, &lrec->limit_cpu, 60))
	    lrec->set |= SET_LIMIT_CPU;
	  break;

	case 'd':
	case 'D':
	  /* RLIMIT_DATA - max data size (KB) */
	  if (!getlimit (&str, &lrec->limit_data, 1024))
	    lrec->set |= SET_LIMIT_DATA;
	  break;

	case 'f':
	case 'F':
	  /* RLIMIT_FSIZE - Maximum filesize (KB) */
	  if (!getlimit (&str, &lrec->limit_fsize, 1024))
	    lrec->set |= SET_LIMIT_FSIZE;
	  break;

	case 'u':
	case 'U':
	  /* RLIMIT_NPROC - max number of processes */
	  if (!getlimit (&str, &lrec->limit_nproc, 1))
	    lrec->set |= SET_LIMIT_NPROC;
	  break;

	case 'c':
	case 'C':
	  /* RLIMIT_CORE - max core file size (KB) */
	  if (!getlimit (&str, &lrec->limit_core, 1024))
	    lrec->set |= SET_LIMIT_CORE;
	  break;
	  
	case 'm':
	case 'M':
	  /* RLIMIT_MEMLOCK - max locked-in-memory
	   * address space (KB)
	   */
	  if (!getlimit (&str, &lrec->limit_memlock, 1024)) 
	    lrec->set |= SET_LIMIT_MEMLOCK;
	  break;

	case 'n':
	case 'N':
	  /* RLIMIT_NOFILE - max number of open files */
	  if (!getlimit (&str, &lrec->limit_nofile, 1))
	    lrec->set |= SET_LIMIT_NOFILE;
	  break;

	case 'r':
	case 'R':
	  /* RLIMIT_RSS - max resident set size (KB) */
	  if (!getlimit (&str, &lrec->limit_rss, 1024))
	    lrec->set |= SET_LIMIT_RSS;
	  break;

	case 's':
	case 'S':
	  /* RLIMIT_STACK - max stack size (KB) */
	  if (!getlimit (&str, &lrec->limit_stack, 1024))
	    lrec->set |= SET_LIMIT_STACK;
	  break;

	case 'l':
	case 'L':
	  lrec->limit_logins = strtol (str, &str, 10);
	  if (lrec->limit_logins >= 0)
	    lrec->set |= SET_LIMIT_LOGINS;
	  break;

	case 'p':
	case 'P':
	  lrec->limit_prio = strtol (str, &str, 10);
	  if (lrec->limit_prio > 0)
	    lrec->set |= SET_LIMIT_PRIO;
	  break;
	  
	default:
	  *endp = str-1;
	  return 1;
	}
    }
  return 0;
}

void
free_limits (limits_record_t rec)
{
  if (rec)
    free (rec);
}

int
limits_cmp (limits_record_t a, limits_record_t b)
{
  if (!a)
    return !!b;
  else if (!b)
    return 1;
  return memcmp (a, b, sizeof (*a));
}