/* $Id: rlimsys.c,v 5.3 2012/07/24 23:16:35 ksb Exp $
 * resource limited system	Kevin S Braunsdorf		(ksb)
 */
#include "machine.h"

#if RESOURCE

#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <signal.h>
#include <stdio.h>
#include <ctype.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>

#include "mk.h"
#include "rlimsys.h"

#if defined(AIX)
#define	vfork	fork
#endif

extern char *progname;
extern int errno;

/* resource records
 */
typedef struct RRnode {
	int rtype;		/* resource type from resource.h	*/
	int fsys;		/* done by kernel or us?		*/
	char *pcrname;		/* user name				*/
	char *pcalias;		/* alias name				*/
	int fset;		/* was given by user?			*/
	struct rlimit ourlim;	/* rlim_cur, rlim_max			*/
} RES_REC;

#define MY_CLOCK	0
/* resources we know about in alpha order
 */
RES_REC aRRLimits[] = {
	{ MY_CLOCK,	0, "clock", (char *)0},
#if defined(RLIMIT_CPU) /* max cpu time in seconds */
	{ RLIMIT_CPU,	1, "cpu", (char *)0},
#endif
#if defined(RLIMIT_FSIZE) /* max file size, ouch! */
	{ RLIMIT_FSIZE,	1, "fsize", (char *)0},
#endif
#if defined(RLIMIT_DATA) /* max data space */
	{ RLIMIT_DATA,	1, "data", (char *)0},
#endif
#if defined(RLIMIT_STACK) /* max stack space */
	{ RLIMIT_STACK,	1, "stack", (char *)0},
#endif
#if defined(RLIMIT_CORE) /* max core dump size */
	{ RLIMIT_CORE,	1, "core", (char *)0},
#endif
#if defined(RLIMIT_RSS) /* resident set size */
	{ RLIMIT_RSS,	1, "rss", (char *)0},
#endif
#if defined(RLIMIT_MEMLOCK) /* wired in (locked in) memory pages */
	{ RLIMIT_MEMLOCK, 1, "memlock", "wired"},
#endif
#if defined(RLIMIT_NPROC) /* number of processes */
	{ RLIMIT_NPROC,	1, "nproc", "processes"},
#endif
#if defined(RLIMIT_NOFILE) /* number of open files */
	{ RLIMIT_NOFILE, 1, "nofile", "files"},
#endif
#if defined(RLIMIT_SBSIZE) /* socket buffer space */
	{ RLIMIT_SBSIZE, 1, "sbsize", "sobuffers"},
#endif
#if defined(RLIMIT_AS)	/* total virtual memeory */
	{ RLIMIT_AS,	1, "as", "vmem"},
#endif
};


/* init the resource limit stuff, make mongo syscalls			(ksb)
 */
char *
rinit(int fList)
{
	register int i;
	register RES_REC *pRR;
	register size_t wList;
	register char *pcRet;

	wList = 0;
	for (i = 0; i < sizeof(aRRLimits)/sizeof(aRRLimits[0]); ++i) {
		pRR = & aRRLimits[i];
		pRR->fset = 0;
		if (pRR->fsys) {
			getrlimit(pRR->rtype, & pRR->ourlim);
		} else {
			pRR->ourlim.rlim_cur = RLIM_INFINITY;
			pRR->ourlim.rlim_max = RLIM_INFINITY;
		}
		if ((char *)0 != pRR->pcalias)
			wList += strlen(pRR->pcalias)+1;
		wList += strlen(pRR->pcrname)+1;
	}
	if (!fList)
		return (char *)0;

	wList = (wList|7)+9; /* leading "," and \000 plus rounding */
	if ((char *)0 == (pcRet = malloc(wList))) {
		return "out of memory";
	}
	*pcRet = '\000';
	for (i = 0; i < sizeof(aRRLimits)/sizeof(aRRLimits[0]); ++i) {
		pRR = & aRRLimits[i];
		strcat(pcRet, ",");
		if ((char *)0 != pRR->pcalias) {
			strcat(pcRet, pRR->pcalias);
			strcat(pcRet, "|");
		}
		strcat(pcRet, pRR->pcrname);
	}
	return pcRet+1;	/* caller cannot free this */
}

/* parse a resource limit statement of the form				(ksb)
 *	<resource>=number_cur/number_max
 * (we eat leading blanks, as all good cvt routines should)
 */
char *
rparse(char *pchState)
{
	register int i;
	register char *pchTemp;
	register RES_REC *pRR;
	register int fSetMax = 0, fSetCur = 0;

	while (isspace(*pchState)) {
		++pchState;
	}
	if ((char *)0 == (pchTemp = strchr(pchState, '='))) {
		return pchState;
	}

	for (i = 0; i < sizeof(aRRLimits)/sizeof(aRRLimits[0]); ++i) {
		pRR = & aRRLimits[i];
		if (0 == strncmp(pRR->pcrname, pchState, strlen(pRR->pcrname)))
			break;
		if ((char *)0 == pRR->pcalias)
			continue;
		if (0 == strncmp(pRR->pcalias, pchState, strlen(pRR->pcalias)))
			break;
	}
	if (sizeof(aRRLimits)/sizeof(aRRLimits[0]) == i) {
		return pchState;
	}

	do {
		++pchTemp;
	} while (isspace(*pchTemp));
	if (isdigit(*pchTemp)) {
		pRR->fset = 1;
		pRR->ourlim.rlim_cur = atol(pchTemp);
		fSetCur = 1;
		do {
			++pchTemp;
		} while (isdigit(*pchTemp));
	}
	if ('/' == *pchTemp)
		++pchTemp;
	while (isspace(*pchTemp))
		++pchTemp;

	if (isdigit(*pchTemp)) {
		pRR->fset = 1;
		pRR->ourlim.rlim_max = atol(pchTemp);
		fSetMax = 1;
		do {
			++pchTemp;
		} while (isdigit(*pchTemp));
	}
	while (isspace(*pchTemp))
		++pchTemp;
	if (fSetMax != fSetCur) {	/* only one set, set both same	*/
		if (fSetMax) {
			pRR->ourlim.rlim_cur = pRR->ourlim.rlim_max;
		} else {
			pRR->ourlim.rlim_max = pRR->ourlim.rlim_cur;
		}
	}
	return pchTemp;
}

int r_fTrace = 0;
static int iChildPid, fWarned;

/* trap a time out for a "real time" alarm
 */
void
do_alarm()
{
	register unsigned next;

	if (kill(iChildPid, fWarned++ ? SIGALRM : SIGKILL) < 0) {
		fprintf(stderr, "%s: kill: %d: %s\n", progname, iChildPid, strerror(errno));
		exit(1);
	}
	next = (unsigned)(aRRLimits[MY_CLOCK].ourlim.rlim_max -
		aRRLimits[MY_CLOCK].ourlim.rlim_cur);
	if (0 == next) {
		next = 2;
	}
	alarm(next);
}

/* emulate "system" with a CPU limit
 */
int
rlimsys(char *pchCmd)
{
	auto int waitbuf;
	register int wret;
	register void *atrap, *itrap, *qtrap;
	register int i;
	register RES_REC *pRR;
	auto int iRetCode;

	iChildPid = vfork();
	switch (iChildPid) {
	case 0:
		for (i = 0; i < sizeof(aRRLimits)/sizeof(aRRLimits[0]); ++i) {
			pRR = & aRRLimits[i];
			if (pRR->fset && pRR->fsys) {
				setrlimit(pRR->rtype, & pRR->ourlim);
			}
		}
		execl("/bin/sh", "sh", r_fTrace ? "-cx" : "-c", pchCmd, (char *)0);
		_exit(127);
	case -1:
		fprintf(stderr, "%s: fork: %s\n", progname, strerror(errno));
		return 1;
	default:
		break;
	}

	itrap = signal(SIGINT, SIG_IGN);
	qtrap = signal(SIGQUIT, SIG_IGN);

	if (aRRLimits[MY_CLOCK].fset) {
		atrap = signal(SIGALRM, do_alarm);
		alarm((unsigned)aRRLimits[MY_CLOCK].ourlim.rlim_cur);
	}
	for (;;) {
		if (iChildPid == (wret = wait3(& waitbuf, WUNTRACED, (struct rusage *)0))) {
			if (WIFSTOPPED(waitbuf)) {
				if (kill(iChildPid, SIGCONT) < 0) {
					break;
				}
				continue;
			}
			if (WIFSIGNALED(waitbuf)) {
				iRetCode = WTERMSIG(waitbuf);
			} else {
				iRetCode = WEXITSTATUS(waitbuf);
			}
			break;
		}
		if (wret == -1) {
			fprintf(stderr, "%s: wait: %d: %s\n", progname, iChildPid, strerror(errno));
			iRetCode = 1;
			break;
		}
	}

	(void) signal(SIGINT, itrap);
	(void) signal(SIGQUIT, qtrap);
	if (aRRLimits[MY_CLOCK].fset) {
		alarm((unsigned)0);
		(void) signal(SIGALRM, atrap);
	}
	return iRetCode;
}

#else
/* Fake init stub to force a NULL no resource list			(ksb)
 */
char *	/*ARGSUSED*/
rinit(int wDummy)
{
	return (char *)0;
}
#endif /* resource usage */