cmd/plimit/plimit.c

/*
 * 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);
}