/*
 * set_limits.c
 * This program demonstrates how to set a limit.
 *
 * to compile:  cc -O -pipe set_limits.c  -o set_limits
 *
 */
#include <stdio.h>
#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <unistd.h>

/* 
 * Our option macros.
 */
#define CPU        ('c')
#define FSIZE      ('f')
#define DATA       ('d')
#define STACK      ('s')
#define CORE       ('C')
#define RSS        ('r')
#define NPROC      ('n')
#define NOFILE     ('o')
#define MEMLOCK    ('m')
#define VMEM       ('v')
#define HELP       ('h')
#define NONE       (RLIM_NLIMITS + 1)
#define OPTIONS    ("c:f:s:C:d:r:n:o:m:v:h")

#define USAGE      ("Usage:  -[opt] <value>\n\n\t-c CPU\n\t-f File size\n\t-d Data segment\n\t-s Stack size \
                    \n\t-C Core size\n\t-r Resident set size\n\t-n Number of children \
                    \n\t-o Open files\n\t-m Memory lock\n\t-v Virtual memory\n\t-h Help\n\n")

#define BYTE(x)     ( x * 1024 )

int main(int argc, char **argv)
{
	int i, ch, which;
	char *name;
	struct rlimit n_lim, o_lim;

	while( (ch = getopt(argc, argv, OPTIONS)) != -1 )
	{
		/*
		 * We are going to take a shortcut and set both soft and hard limits
		 * to be the same. As you know, they dont have to be the same,
		 * the soft limit can and normally should be set below the
		 * hard limit.
		 */
		switch( ch )
		{
			case CPU:
				name = "RLIMIT_CPU";
				which = RLIMIT_CPU;
				break;  
			case FSIZE:
				name = "RLIMIT_FSIZE";
				which = RLIMIT_FSIZE;
				break;

			case DATA:
				name = "RLIMIT_DATA";
				which = RLIMIT_DATA;
				break;

			case STACK:
				name = "RLIMIT_STACK";
				which = RLIMIT_STACK;
				break;

			case CORE:
				name = "RLIMIT_CORE";
				which = RLIMIT_CORE;
				break;

			case RSS:
				name = "RLIMIT_RSS";
				which = RLIMIT_RSS;
				break;

			case NPROC:
				name = "RLIMIT_NPROC";
				which = RLIMIT_NPROC;
				break;

			case NOFILE:
				name = "RLIMIT_NOFILE";
				which = RLIMIT_NOFILE;
				break;

			case MEMLOCK:
				name = "RLIMIT_MEMLOCK";
				which = RLIMIT_MEMLOCK;
				break;

			case VMEM:
				name = "RLIMIT_VMEM";
				which = RLIMIT_VMEM;
				break;

			case HELP:
			default:
				which = NONE;
				printf(USAGE);
				break;

		} /* switch */

		if( which != NONE )
		{
			n_lim.rlim_cur = atoi(optarg);
			n_lim.rlim_max = n_lim.rlim_cur;

			/* Get the old limit first */
			getrlimit(which, &o_lim);
			/* Now set it */
			if( setrlimit(which, &n_lim) != 0 )
			{
				printf("Error cannot set value !\n");
			}
			else
			{
				printf("%s :\n", name);
				printf("\told \tHard: %d \t Soft: %d \n", o_lim.rlim_max, o_lim.rlim_cur );
				printf("\tnew \tHard: %d \t Soft: %d \n", n_lim.rlim_max, n_lim.rlim_cur );               
			}

		} /* if( which) */

	} /* while */

	/* FIN */
	return(0);
}