/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2004 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */
/*
 * Copyright 2015, Joyent, Inc.
 */

#define __EXTENSIONS__  /* For strtok_r */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <ctype.h>
#include <string.h>
#include <signal.h>
#include <limits.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <sys/mkdev.h>
#include <libproc.h>
#include <priv.h>

#define TRUE    1
#define FALSE   0

static  int     interrupt;
static  char    *command;
static  int     Fflag;
static  int     kbytes = FALSE;
static  int     mbytes = FALSE;
static  char    set_current[RLIM_NLIMITS];
static  char    set_maximum[RLIM_NLIMITS];
static  struct rlimit64 rlimit[RLIM_NLIMITS];

static  void    intr(int);
static  int     parse_limits(int, char *);
static  void    show_limits(struct ps_prochandle *);
static  int     set_limits(struct ps_prochandle *);

static void
usage()
{
        (void) fprintf(stderr,
            "usage:\n"
            "    For each process, report all resource limits:\n"
            "\t%s [-km] pid ...\n"
            "\t-k\treport file sizes in kilobytes\n"
            "\t-m\treport file/memory sizes in megabytes\n"
            "    For each process, set specified resource limits:\n"
            "\t%s -{cdfnstv} soft,hard ... pid ...\n"
            "\t-c soft,hard\tset core file size limits\n"
            "\t-d soft,hard\tset data segment (heap) size limits\n"
            "\t-f soft,hard\tset file size limits\n"
            "\t-n soft,hard\tset file descriptor limits\n"
            "\t-s soft,hard\tset stack segment size limits\n"
            "\t-t soft,hard\tset CPU time limits\n"
            "\t-v soft,hard\tset virtual memory size limits\n"
            "\t(default units are as shown by the output of '%s pid')\n",
            command, command, command);
        exit(2);
}

int
main(int argc, char **argv)
{
        int retc = 0;
        int opt;
        int errflg = 0;
        int set = FALSE;
        struct ps_prochandle *Pr;

        if ((command = strrchr(argv[0], '/')) != NULL)
                command++;
        else
                command = argv[0];

        while ((opt = getopt(argc, argv, "Fkmc:d:f:n:s:t:v:")) != EOF) {
                switch (opt) {
                case 'F':               /* force grabbing (no O_EXCL) */
                        Fflag = PGRAB_FORCE;
                        break;
                case 'k':
                        kbytes = TRUE;
                        mbytes = FALSE;
                        break;
                case 'm':
                        kbytes = FALSE;
                        mbytes = TRUE;
                        break;
                case 'c':       /* core file size */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_CORE, optarg);
                        break;
                case 'd':       /* data segment size */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_DATA, optarg);
                        break;
                case 'f':       /* file size */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_FSIZE, optarg);
                        break;
                case 'n':       /* file descriptors */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_NOFILE, optarg);
                        break;
                case 's':       /* stack segment size */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_STACK, optarg);
                        break;
                case 't':       /* CPU time */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_CPU, optarg);
                        break;
                case 'v':       /* virtual memory size */
                        set = TRUE;
                        errflg += parse_limits(RLIMIT_VMEM, optarg);
                        break;
                default:
                        errflg = 1;
                        break;
                }
        }

        argc -= optind;
        argv += optind;

        if (errflg || argc <= 0)
                usage();

        /* catch signals from terminal */
        if (sigset(SIGHUP, SIG_IGN) == SIG_DFL)
                (void) sigset(SIGHUP, intr);
        if (sigset(SIGINT, SIG_IGN) == SIG_DFL)
                (void) sigset(SIGINT, intr);
        if (sigset(SIGQUIT, SIG_IGN) == SIG_DFL)
                (void) sigset(SIGQUIT, intr);
        (void) sigset(SIGPIPE, intr);
        (void) sigset(SIGTERM, intr);

        while (--argc >= 0 && !interrupt) {
                psinfo_t psinfo;
                char *arg;
                pid_t pid;
                int gret;

                (void) fflush(stdout);  /* process-at-a-time */

                /* get the specified pid and the psinfo struct */
                if ((pid = proc_arg_psinfo(arg = *argv++, PR_ARG_PIDS,
                    &psinfo, &gret)) == -1) {
                        (void) fprintf(stderr, "%s: cannot examine %s: %s\n",
                            command, arg, Pgrab_error(gret));
                        retc = 1;
                } else if ((Pr = Pgrab(pid, Fflag, &gret)) != NULL) {
                        if (Pcreate_agent(Pr) == 0) {
                                if (set) {
                                        if (set_limits(Pr) != 0)
                                                retc = 1;
                                } else {
                                        proc_unctrl_psinfo(&psinfo);
                                        (void) printf("%d:\t%.70s\n",
                                            (int)pid, psinfo.pr_psargs);
                                        show_limits(Pr);
                                }
                                Pdestroy_agent(Pr);
                        } else {
                                (void) fprintf(stderr,
                                    "%s: cannot control process %d\n",
                                    command, (int)pid);
                                retc = 1;
                        }
                        Prelease(Pr, 0);
                } else {
                        if ((gret == G_SYS || gret == G_SELF) && !set) {
                                proc_unctrl_psinfo(&psinfo);
                                (void) printf("%d:\t%.70s\n", (int)pid,
                                    psinfo.pr_psargs);
                                if (gret == G_SYS)
                                        (void) printf("  [system process]\n");
                                else
                                        show_limits(NULL);
                        } else {
                                (void) fprintf(stderr,
                                    "%s: %s: %d\n",
                                    command, Pgrab_error(gret), (int)pid);
                                retc = 1;
                        }
                }
        }

        if (interrupt)
                retc = 1;
        return (retc);
}

static void
intr(int sig)
{
        interrupt = sig;
}

/* ------ begin specific code ------ */

/*
 * Compute a limit, given a string:
 *      unlimited       unlimited
 *      nnn k           nnn kilobytes
 *      nnn m           nnn megabytes (minutes for CPU time)
 *      nnn h           nnn hours (for CPU time only)
 *      mm : ss         minutes and seconds (for CPU time only)
 */
static int
limit_value(int which, char *arg, rlim64_t *limit)
{
        rlim64_t value;
        rlim64_t unit;
        char *lastc;

        if (strcmp(arg, "unlimited") == 0) {
                *limit = RLIM64_INFINITY;
                return (0);
        }

        if (which == RLIMIT_CPU && strchr(arg, ':') != NULL) {
                char *minutes = strtok_r(arg, " \t:", &lastc);
                char *seconds = strtok_r(NULL, " \t", &lastc);
                rlim64_t sec;

                if (seconds != NULL && strtok_r(NULL, " \t", &lastc) != NULL)
                        return (1);
                value = strtoull(minutes, &lastc, 10);
                if (*lastc != '\0' || value > RLIM64_INFINITY / 60)
                        return (1);
                if (seconds == NULL || *seconds == '\0')
                        sec = 0;
                else {
                        sec = strtoull(seconds, &lastc, 10);
                        if (*lastc != '\0' || sec > 60)
                                return (1);
                }
                value = value * 60 + sec;
                if (value > RLIM64_INFINITY)
                        value = RLIM64_INFINITY;
                *limit = value;
                return (0);
        }

        switch (*(lastc = arg + strlen(arg) - 1)) {
        case 'k':
                unit = 1024;
                *lastc = '\0';
                break;
        case 'm':
                if (which == RLIMIT_CPU)
                        unit = 60;
                else
                        unit = 1024 * 1024;
                *lastc = '\0';
                break;
        case 'h':
                if (which == RLIMIT_CPU)
                        unit = 60 * 60;
                else
                        return (1);
                *lastc = '\0';
                break;
        default:
                switch (which) {
                case RLIMIT_CPU:        unit = 1;       break;
                case RLIMIT_FSIZE:      unit = 512;     break;
                case RLIMIT_DATA:       unit = 1024;    break;
                case RLIMIT_STACK:      unit = 1024;    break;
                case RLIMIT_CORE:       unit = 512;     break;
                case RLIMIT_NOFILE:     unit = 1;       break;
                case RLIMIT_VMEM:       unit = 1024;    break;
                }
                break;
        }

        value = strtoull(arg, &lastc, 10);
        if (*lastc != '\0' || value > RLIM64_INFINITY / unit)
                return (1);

        value *= unit;
        if (value > RLIM64_INFINITY)
                value = RLIM64_INFINITY;
        *limit = value;
        return (0);
}

static int
parse_limits(int which, char *arg)
{
        char *lastc;
        char *soft = strtok_r(arg, " \t,", &lastc);
        char *hard = strtok_r(NULL, " \t", &lastc);
        struct rlimit64 *rp = &rlimit[which];

        if (hard != NULL && strtok_r(NULL, " \t", &lastc) != NULL)
                return (1);

        if (soft == NULL || *soft == '\0') {
                rp->rlim_cur = 0;
                set_current[which] = FALSE;
        } else {
                if (limit_value(which, soft, &rp->rlim_cur) != 0)
                        return (1);
                set_current[which] = TRUE;
        }

        if (hard == NULL || *hard == '\0') {
                rp->rlim_max = 0;
                set_maximum[which] = FALSE;
        } else {
                if (limit_value(which, hard, &rp->rlim_max) != 0)
                        return (1);
                set_maximum[which] = TRUE;
        }
        if (set_current[which] && set_maximum[which] &&
            rp->rlim_cur > rp->rlim_max)
                return (1);

        return (0);
}

static void
limit_adjust(struct rlimit64 *rp, int units)
{
        if (rp->rlim_cur != RLIM64_INFINITY)
                rp->rlim_cur /= units;
        if (rp->rlim_max != RLIM64_INFINITY)
                rp->rlim_max /= units;
}

static char *
limit_values(struct rlimit64 *rp)
{
        static char buffer[64];
        char buf1[32];
        char buf2[32];
        char *s1;
        char *s2;

        if (rp->rlim_cur == RLIM64_INFINITY)
                s1 = "unlimited";
        else {
                (void) sprintf(s1 = buf1, "%lld", rp->rlim_cur);
                if (strlen(s1) < 8)
                        (void) strcat(s1, "\t");
        }

        if (rp->rlim_max == RLIM64_INFINITY)
                s2 = "unlimited";
        else {
                (void) sprintf(s2 = buf2, "%lld", rp->rlim_max);
        }

        (void) sprintf(buffer, "%s\t%s", s1, s2);

        return (buffer);
}

static void
show_limits(struct ps_prochandle *Pr)
{
        struct rlimit64 rlim;
        int resource;
        char buf[32];
        char *s;

        (void) printf("   resource\t\t current\t maximum\n");

        for (resource = 0; resource < RLIM_NLIMITS; resource++) {
                if (pr_getrlimit64(Pr, resource, &rlim) != 0)
                        continue;

                switch (resource) {
                case RLIMIT_CPU:
                        s = "  time(seconds)\t\t";
                        break;
                case RLIMIT_FSIZE:
                        if (kbytes) {
                                s = "  file(kbytes)\t\t";
                                limit_adjust(&rlim, 1024);
                        } else if (mbytes) {
                                s = "  file(mbytes)\t\t";
                                limit_adjust(&rlim, 1024 * 1024);
                        } else {
                                s = "  file(blocks)\t\t";
                                limit_adjust(&rlim, 512);
                        }
                        break;
                case RLIMIT_DATA:
                        if (mbytes) {
                                s = "  data(mbytes)\t\t";
                                limit_adjust(&rlim, 1024 * 1024);
                        } else {
                                s = "  data(kbytes)\t\t";
                                limit_adjust(&rlim, 1024);
                        }
                        break;
                case RLIMIT_STACK:
                        if (mbytes) {
                                s = "  stack(mbytes)\t\t";
                                limit_adjust(&rlim, 1024 * 1024);
                        } else {
                                s = "  stack(kbytes)\t\t";
                                limit_adjust(&rlim, 1024);
                        }
                        break;
                case RLIMIT_CORE:
                        if (kbytes) {
                                s = "  coredump(kbytes)\t";
                                limit_adjust(&rlim, 1024);
                        } else if (mbytes) {
                                s = "  coredump(mbytes)\t";
                                limit_adjust(&rlim, 1024 * 1024);
                        } else {
                                s = "  coredump(blocks)\t";
                                limit_adjust(&rlim, 512);
                        }
                        break;
                case RLIMIT_NOFILE:
                        s = "  nofiles(descriptors)\t";
                        break;
                case RLIMIT_VMEM:
                        if (mbytes) {
                                s = "  vmemory(mbytes)\t";
                                limit_adjust(&rlim, 1024 * 1024);
                        } else {
                                s = "  vmemory(kbytes)\t";
                                limit_adjust(&rlim, 1024);
                        }
                        break;
                default:
                        (void) sprintf(buf, "  rlimit #%d\t", resource);
                        s = buf;
                        break;
                }

                (void) printf("%s%s\n", s, limit_values(&rlim));
        }
}

static int
set_one_limit(struct ps_prochandle *Pr, int which, rlim64_t cur, rlim64_t max)
{
        struct rlimit64 rlim;
        int be_su = 0;
        prpriv_t *old_prpriv = NULL, *new_prpriv = NULL;
        priv_set_t *eset, *pset;
        int ret = 0;

        if (pr_getrlimit64(Pr, which, &rlim) != 0) {
                (void) fprintf(stderr,
                    "%s: unable to get process limit for pid %d: %s\n",
                    command, Pstatus(Pr)->pr_pid, strerror(errno));
                return (1);
        }

        if (!set_current[which])
                cur = rlim.rlim_cur;
        if (!set_maximum[which])
                max = rlim.rlim_max;

        if (max < cur)
                max = cur;

        if (max > rlim.rlim_max && Pr != NULL)
                be_su = 1;
        rlim.rlim_cur = cur;
        rlim.rlim_max = max;

        if (be_su) {
                new_prpriv = proc_get_priv(Pstatus(Pr)->pr_pid);
                if (new_prpriv == NULL) {
                        (void) fprintf(stderr,
                            "%s: unable to get process privileges for pid"
                            " %d: %s\n", command, Pstatus(Pr)->pr_pid,
                            strerror(errno));
                        return (1);
                }

                /*
                 * We only have to change the process privileges if it doesn't
                 * already have PRIV_SYS_RESOURCE.  In addition, we want to make
                 * sure that we don't leave a process with elevated privileges,
                 * so we make sure the process dies if we exit unexpectedly.
                 */
                eset = (priv_set_t *)
                    &new_prpriv->pr_sets[new_prpriv->pr_setsize *
                    priv_getsetbyname(PRIV_EFFECTIVE)];
                pset = (priv_set_t *)
                    &new_prpriv->pr_sets[new_prpriv->pr_setsize *
                    priv_getsetbyname(PRIV_PERMITTED)];
                if (!priv_ismember(eset, PRIV_SYS_RESOURCE)) {
                        /* Keep track of original privileges */
                        old_prpriv = proc_get_priv(Pstatus(Pr)->pr_pid);
                        if (old_prpriv == NULL) {
                                proc_free_priv(new_prpriv);
                                (void) fprintf(stderr,
                                    "%s: unable to get process privileges "
                                    "for pid %d: %s\n", command,
                                    Pstatus(Pr)->pr_pid, strerror(errno));
                                return (1);
                        }

                        (void) priv_addset(eset, PRIV_SYS_RESOURCE);
                        (void) priv_addset(pset, PRIV_SYS_RESOURCE);

                        if (Psetflags(Pr, PR_KLC) != 0 ||
                            Psetpriv(Pr, new_prpriv) != 0) {
                                (void) fprintf(stderr,
                                    "%s: unable to set process privileges for"
                                    " pid %d: %s\n", command,
                                    Pstatus(Pr)->pr_pid, strerror(errno));
                                (void) Punsetflags(Pr, PR_KLC);
                                proc_free_priv(new_prpriv);
                                proc_free_priv(old_prpriv);
                                return (1);
                        }
                }
        }

        if (pr_setrlimit64(Pr, which, &rlim) != 0) {
                (void) fprintf(stderr,
                    "%s: cannot set resource limit for pid %d: %s\n",
                    command, Pstatus(Pr)->pr_pid, strerror(errno));
                ret = 1;
        }

        if (old_prpriv != NULL) {
                if (Psetpriv(Pr, old_prpriv) != 0) {
                        /*
                         * If this fails, we can't leave a process hanging
                         * around with elevated privileges, so we'll have to
                         * release the process from libproc, knowing that it
                         * will be killed (since we set PR_KLC).
                         */
                        Pdestroy_agent(Pr);
                        (void) fprintf(stderr,
                            "%s: cannot relinquish privileges for pid %d."
                            " The process was killed.",
                            command, Pstatus(Pr)->pr_pid);
                        ret = 1;
                }
                if (Punsetflags(Pr, PR_KLC) != 0) {
                        (void) fprintf(stderr,
                            "%s: cannot relinquish privileges for pid %d."
                            " The process was killed.",
                            command, Pstatus(Pr)->pr_pid);
                        ret = 1;
                }

                proc_free_priv(old_prpriv);
        }

        if (new_prpriv != NULL)
                proc_free_priv(new_prpriv);

        return (ret);
}

static int
set_limits(struct ps_prochandle *Pr)
{
        int which;
        int retc = 0;

        for (which = 0; which < RLIM_NLIMITS; which++) {
                if (set_current[which] || set_maximum[which]) {
                        if (set_one_limit(Pr, which, rlimit[which].rlim_cur,
                            rlimit[which].rlim_max) != 0)
                                retc = 1;
                }
        }

        return (retc);
}