feat: refactor launcher.c

Co-authored-by: cxy004 <cxy004@qq.com>
This commit is contained in:
Paul Pan 2022-10-11 16:35:55 +08:00
parent b1ae874394
commit 05372f10e8

View File

@ -13,18 +13,21 @@
#include <time.h> #include <time.h>
#include <unistd.h> #include <unistd.h>
struct Config { enum ConfigIndex {
char *memory_limit; memory_limit = 0,
char *nproc_limit; nproc_limit,
char *time_limit; time_limit,
char *sandbox_path; sandbox_path,
char *sandbox_template; sandbox_template,
char *sandbox_action; sandbox_action,
char *file_input; file_input,
char *file_output; file_output,
char *file_info; file_info,
char *program; program,
} config; CONFIG_INDEX_MAX
};
char *config[CONFIG_INDEX_MAX];
void print_help(char *self) { void print_help(char *self) {
LOG_WARN("Usage:"); LOG_WARN("Usage:");
@ -44,57 +47,35 @@ void print_help(char *self) {
} }
void parse(int argc, char *argv[]) { void parse(int argc, char *argv[]) {
static struct option long_options[] = { static struct option options[] = {
{"memory_limit", required_argument, NULL, 0}, [memory_limit] = {"memory_limit", required_argument, NULL, 0},
{"nproc_limit", required_argument, NULL, 0}, [nproc_limit] = {"nproc_limit", required_argument, NULL, 0},
{"time_limit", required_argument, NULL, 0}, [time_limit] = {"time_limit", required_argument, NULL, 0},
{"sandbox_path", required_argument, NULL, 0}, [sandbox_path] = {"sandbox_path", required_argument, NULL, 0},
{"sandbox_template", required_argument, NULL, 0}, [sandbox_template] = {"sandbox_template", required_argument, NULL, 0},
{"sandbox_action", required_argument, NULL, 0}, [sandbox_action] = {"sandbox_action", required_argument, NULL, 0},
{"file_input", required_argument, NULL, 0}, [file_input] = {"file_input", required_argument, NULL, 0},
{"file_output", required_argument, NULL, 0}, [file_output] = {"file_output", required_argument, NULL, 0},
{"file_info", required_argument, NULL, 0}, [file_info] = {"file_info", required_argument, NULL, 0},
{"program", required_argument, NULL, 0}, [program] = {"program", required_argument, NULL, 0},
{"help", no_argument, NULL, 0}, [CONFIG_INDEX_MAX] = {"help", no_argument, NULL, 0},
{NULL, 0, NULL, 0} [CONFIG_INDEX_MAX + 1] = {NULL, 0, NULL, 0}
}; };
int c, idx = 0; int c, idx = 0;
while ((c = getopt_long_only(argc, argv, "", long_options, &idx)) != -1) { while ((c = getopt_long_only(argc, argv, "", options, &idx)) != -1) {
if (c != 0) break; if (c != 0) break;
const char *key = long_options[idx].name; if (idx < CONFIG_INDEX_MAX)
char *val = optarg; config[idx] = optarg;
else if (idx == CONFIG_INDEX_MAX) {
if (strcmp(key, "memory_limit") == 0) {
config.memory_limit = val;
} else if (strcmp(key, "nproc_limit") == 0) {
config.nproc_limit = val;
} else if (strcmp(key, "time_limit") == 0) {
config.time_limit = val;
} else if (strcmp(key, "sandbox_path") == 0) {
config.sandbox_path = val;
} else if (strcmp(key, "sandbox_template") == 0) {
config.sandbox_template = val;
} else if (strcmp(key, "sandbox_action") == 0) {
config.sandbox_action = val;
} else if (strcmp(key, "file_input") == 0) {
config.file_input = val;
} else if (strcmp(key, "file_output") == 0) {
config.file_output = val;
} else if (strcmp(key, "file_info") == 0) {
config.file_info = val;
} else if (strcmp(key, "program") == 0) {
config.program = val;
} else if (strcmp(key, "help") == 0) {
print_help(argv[0]); print_help(argv[0]);
exit(0); exit(0);
} }
} }
char **cfg = (char **)&config; for (int i = 0; i < CONFIG_INDEX_MAX; i++) {
for (int i = 0; i < 10; i++) { if (!config[i]) {
if (!cfg[i]) {
print_help(argv[0]); print_help(argv[0]);
LOG_ERR("Missing arguments"); LOG_ERR("Missing arguments");
exit(ERR_ARGUMENTS); exit(ERR_ARGUMENTS);
@ -103,44 +84,44 @@ void parse(int argc, char *argv[]) {
} }
void launch_child() { void launch_child() {
char *args[] = {config.program, NULL}; char *args[] = {config[program], NULL};
char *env[7]; char *env[7];
{ /* build env */ {
env[0] = malloc(sizeof("LD_PRELOAD=") + strlen(config.sandbox_path) + 1); env[0] = malloc(sizeof("LD_PRELOAD=") + strlen(config[sandbox_path]) + 1);
sprintf(env[0], "LD_PRELOAD=%s", config.sandbox_path); sprintf(env[0], "LD_PRELOAD=%s", config[sandbox_path]);
env[1] = malloc(sizeof(LIMIT_MEMORY "=") + strlen(config.memory_limit) + 1); env[1] = malloc(sizeof(LIMIT_MEMORY "=") + strlen(config[memory_limit]) + 1);
sprintf(env[1], LIMIT_MEMORY "=%s", config.memory_limit); sprintf(env[1], LIMIT_MEMORY "=%s", config[memory_limit]);
env[2] = malloc(sizeof(LIMIT_NPROC "=") + strlen(config.nproc_limit) + 1); env[2] = malloc(sizeof(LIMIT_NPROC "=") + strlen(config[nproc_limit]) + 1);
sprintf(env[2], LIMIT_NPROC "=%s", config.nproc_limit); sprintf(env[2], LIMIT_NPROC "=%s", config[nproc_limit]);
env[3] = malloc(sizeof(LIMIT_TIME "=") + strlen(config.time_limit) + 1); env[3] = malloc(sizeof(LIMIT_TIME "=") + strlen(config[time_limit]) + 1);
sprintf(env[3], LIMIT_TIME "=%s", config.time_limit); sprintf(env[3], LIMIT_TIME "=%s", config[time_limit]);
env[4] = malloc(sizeof(SANDBOX_TEMPLATE "=") + strlen(config.sandbox_template) + 1); env[4] = malloc(sizeof(SANDBOX_TEMPLATE "=") + strlen(config[sandbox_template]) + 1);
sprintf(env[4], SANDBOX_TEMPLATE "=%s", config.sandbox_template); sprintf(env[4], SANDBOX_TEMPLATE "=%s", config[sandbox_template]);
env[5] = malloc(sizeof(SANDBOX_ACTION "=") + strlen(config.sandbox_action) + 1); env[5] = malloc(sizeof(SANDBOX_ACTION "=") + strlen(config[sandbox_action]) + 1);
sprintf(env[5], SANDBOX_ACTION "=%s", config.sandbox_action); sprintf(env[5], SANDBOX_ACTION "=%s", config[sandbox_action]);
env[6] = NULL; env[6] = NULL;
} }
{ /* build stdin */ {
int fd = open(config.file_input, O_RDONLY); int fd = open(config[file_input], O_RDONLY);
dup2(fd, STDIN_FILENO); dup2(fd, STDIN_FILENO);
close(fd); close(fd);
} }
{ /* build stdout */ {
int fd = open(config.file_output, O_WRONLY | O_CREAT, 0644); int fd = open(config[file_output], O_WRONLY | O_CREAT, 0644);
dup2(fd, STDOUT_FILENO); dup2(fd, STDOUT_FILENO);
close(fd); close(fd);
} }
execve(config.program, args, env); execve(config[program], args, env);
} }
void dump_info(FILE *dest, struct rusage *usage, int status, long long time_usage) { void dump_info(FILE *dest, struct rusage *usage, int status, long long time_usage) {
@ -174,7 +155,7 @@ int main(int argc, char *argv[]) {
exit(ERR_FORK); exit(ERR_FORK);
} else if (killer == 0) { // Killer } else if (killer == 0) { // Killer
long limit = (strtol(config.time_limit, NULL, 10) + 1000) / 1000 + 2; long limit = (strtol(config[time_limit], NULL, 10) + 1000) / 1000 + 2;
LOG_INFO("Killer started, time limit: %lds", limit); LOG_INFO("Killer started, time limit: %lds", limit);
sleep(limit); // two more seconds sleep(limit); // two more seconds
LOG_WARN("Killer killed child"); LOG_WARN("Killer killed child");
@ -201,7 +182,7 @@ int main(int argc, char *argv[]) {
exit(ERR_WAIT); exit(ERR_WAIT);
} }
FILE *info = fopen(config.file_info, "w"); FILE *info = fopen(config[file_info], "w");
dump_info(info, &usage, status, time_usage); dump_info(info, &usage, status, time_usage);
exit(0); exit(0);
} }