Executor.cpp

executor.cpp

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cassert>
#include <errno.h>

#include <pwd.h>
#include <unistd.h>
#include <sys/resource.h>
#include <sys/signal.h>
#include <sys/time.h>
#include <sys/wait.h>

char *stdin_file, *stdout_file, *stderr_file, *exec_path, *exec_file, *user;
bool is_trust;
long time_limit;


pid_t exec_pid, alarm_pid;
uid_t user_uid;
long rtime0;

void report_error(const char *message) {
	fprintf(stderr, "%s", message);
	exit(1);
}

void report_warning(const char *message) {
	fprintf(stderr, "Warning: %s\n", message);
}

void read_parameter(int argc, char *argv[]) {
	if (argc < 2) report_error("Parameter Error.\n");
	for (int i = 0; i < argc; ++ i) {
		if (argv[i][0] != '-' || argv[i][1] != '-') {
			exec_path = argv[i];
			char *p = argv[i];
			while (p = strchr(p, '/')) {
				exec_file = ++ p;
			}
		} else {
			char *parameter = new char[strlen(argv[i])];
			char *value = strchr(argv[i], '=');
			if (value) ++ value;
			strcpy(parameter, strtok(argv[i], "="));
			if (strcmp(parameter, "--stdin-file") == 0) {
				stdin_file = value;
			} else if (strcmp(parameter, "--stdout-file") == 0) {
				stdout_file = value;
			} else if (strcmp(parameter, "--stderr-file") == 0) {
				stderr_file = value;
			} else if (strcmp(parameter, "--time-limit") == 0) {
				sscanf(value, "%ld", &time_limit);
			} else if (strcmp(parameter, "--user") == 0) {
				user = value;
			} else if (strcmp(parameter, "--trust") == 0) {
				is_trust = *value == 't';
			} else {
				report_warning("Parameter not exist.");
			}
		}
	}
}

void executor() {
	if (stdin_file) {
		freopen(stdin_file, "r", stdin);
	}
	if (stdout_file)
		freopen(stdout_file, "w", stdout);
//	if (stderr_file)
//		freopen(stderr_file, "w", stderr);
	if (user) {
		setuid(user_uid);
	}

	if (time_limit) {
		rlimit rl;
		rl.rlim_cur = (long)ceil(time_limit / 1000.0);
		rl.rlim_max = rl.rlim_cur + 1;
		setrlimit(RLIMIT_CPU, &rl);
	}

	// large stack
	{
		rlimit rl;
		rl.rlim_cur = 256 << 20;
		rl.rlim_max = rl.rlim_cur;
		setrlimit(RLIMIT_STACK, &rl);
	}

	// limit number process
	{
		rlimit rl;
		rl.rlim_cur = 4;
		rl.rlim_max = rl.rlim_cur;
		setrlimit(RLIMIT_NPROC, &rl);
	}

	// limit memory
	{
		rlimit rl;
		rl.rlim_cur = 1;
		rl.rlim_cur = rl.rlim_cur << 32;
		rl.rlim_max = rl.rlim_cur;
		setrlimit(RLIMIT_AS, &rl);
	}

	// limit output
	{
		rlimit rl;
		rl.rlim_cur = 64 << 20;
		rl.rlim_max = rl.rlim_cur;
		setrlimit(RLIMIT_FSIZE, &rl);
	}

	// limit open file number
	{
		rlimit rl;
		rl.rlim_cur = 0;
		rl.rlim_max = 0;
		setrlimit(RLIMIT_NOFILE, &rl);
	}

	int ret = execlp(exec_path, exec_file, NULL);
	if (ret) report_error("Exec error.\n");
}

void monitor() {
	alarm_pid = vfork();
	if (!alarm_pid) {
		sleep((long)ceil(time_limit / 1000.0) + 1);
		exit(0);
	} else if (alarm_pid < 0) {
		kill(exec_pid, SIGKILL);
		report_error("Fork error.");
	}
	int status;
	pid_t ret_pid;
	rusage rinfo;
	timeval tinfo;
	while (ret_pid = wait3(&status, 0, &rinfo)) {
		gettimeofday(&tinfo, NULL);
		long utime = rinfo.ru_utime.tv_sec * 1000 + rinfo.ru_utime.tv_usec / 1000;
		long stime = rinfo.ru_stime.tv_sec * 1000 + rinfo.ru_stime.tv_usec / 1000;
		long rtime = (tinfo.tv_sec * 1000 + tinfo.tv_usec / 1000) - rtime0;
		long memory = rinfo.ru_maxrss;
		if (ret_pid == exec_pid) {
			if (WIFEXITED(status)) {
				printf("OK\n");
				printf("%d\n", WEXITSTATUS(status));
				printf("%ld %ld %ld\n", utime, stime, rtime);
				printf("%ld\n", memory);
			} else if (WIFSIGNALED(status)) {
				printf("??\n");
				printf("%d %s\n", WTERMSIG(status), sys_siglist[WTERMSIG(status)]);
				printf("%ld %ld %ld\n", utime, stime, rtime);
				printf("%ld\n", memory);
			} else if (WIFSTOPPED(status)) {
				kill(exec_pid, SIGKILL);
			} else {
				kill(exec_pid, SIGKILL);
			}
			break;
		} else if (ret_pid == alarm_pid){
			kill(exec_pid, SIGKILL);
		} else {
			kill(exec_pid, SIGKILL);
			report_error("Wait error.\n");
		}
	}
	kill(alarm_pid, SIGKILL);
}

int main(int argc, char *argv[]) {
	read_parameter(argc, argv);

	{
		timeval tinfo;
		gettimeofday(&tinfo, NULL);
		rtime0 = tinfo.tv_sec * 1000 + tinfo.tv_usec / 1000;
	}

	if (user) {
		passwd *pw = getpwnam(user);
		if (!user) report_error("User set but not exist.\n");
		user_uid = pw->pw_uid;
		if (getuid()) report_error("Only root can setuid.\n");
	} else if (!getuid()) {
		report_error("Cannot run as root.\n");
	}

	exec_pid = fork();
	if (!exec_pid) {
		executor();
	} else if (exec_pid > 0) {
		monitor();
	} else {
		report_error("Fork error.\n");
	}

	return 0;
}