S2OJ/judge_client/1/uoj_judger/run/run_interaction.cpp

264 lines
5.9 KiB
C++
Raw Normal View History

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <exception>
#include <system_error>
#include <thread>
#include <unistd.h>
#include <asm/unistd.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/resource.h>
#include <sys/user.h>
#include <fcntl.h>
#include <cstring>
#include <string>
#include <vector>
#include <set>
#include <argp.h>
#include "uoj_env.h"
using namespace std;
struct RunResult {
int result;
int ust;
int usm;
int exit_code;
RunResult(int _result, int _ust = -1, int _usm = -1, int _exit_code = -1)
: result(_result), ust(_ust), usm(_usm), exit_code(_exit_code) {
if (result != RS_AC) {
ust = -1, usm = -1;
}
}
};
struct PipeConfig {
int from, to;
int from_fd, to_fd;
string saving_file_name; // empty for none
PipeConfig() {
}
PipeConfig(string str) {
if (sscanf(str.c_str(), "%d:%d-%d:%d", &from, &from_fd, &to, &to_fd) != 4) {
throw invalid_argument("bad init str for PipeConfig");
}
from -= 1;
to -= 1;
}
};
struct RunInteractionConfig {
vector<string> cmds;
vector<PipeConfig> pipes;
};
struct RunCmdData {
string cmd;
pid_t pid;
vector<int> ipipes, opipes;
};
struct PipeData {
PipeConfig config;
int ipipefd[2], opipefd[2];
thread io_thread;
exception_ptr eptr;
};
class RunInteraction {
private:
vector<RunCmdData> cmds;
vector<PipeData> pipes;
void prepare_fd() { // me
for (int i = 0; i < (int)pipes.size(); i++) {
close(pipes[i].ipipefd[1]);
close(pipes[i].opipefd[0]);
}
}
void prepare_fd_for_cmd(int id) {
freopen("/dev/null", "r", stdin);
freopen("/dev/null", "w", stdout);
freopen("/dev/null", "w", stderr);
for (int i = 0; i < (int)pipes.size(); i++) {
if (pipes[i].config.from == id) {
dup2(pipes[i].ipipefd[1], 128 + pipes[i].ipipefd[1]);
}
if (pipes[i].config.to == id) {
dup2(pipes[i].opipefd[0], 128 + pipes[i].opipefd[0]);
}
close(pipes[i].ipipefd[0]);
close(pipes[i].ipipefd[1]);
close(pipes[i].opipefd[0]);
close(pipes[i].opipefd[1]);
}
for (int i = 0; i < (int)pipes.size(); i++) {
if (pipes[i].config.from == id) {
dup2(128 + pipes[i].ipipefd[1], pipes[i].config.from_fd);
close(128 + pipes[i].ipipefd[1]);
}
if (pipes[i].config.to == id) {
dup2(128 + pipes[i].opipefd[0], pipes[i].config.to_fd);
close(128 + pipes[i].opipefd[0]);
}
}
}
void wait_pipe_io(int pipe_id) {
FILE *sf = NULL;
if (!pipes[pipe_id].config.saving_file_name.empty())
sf = fopen(pipes[pipe_id].config.saving_file_name.c_str(), "w");
int ifd = pipes[pipe_id].ipipefd[0];
int ofd = pipes[pipe_id].opipefd[1];
FILE *inf = fdopen(ifd, "r");
FILE *ouf = fdopen(ofd, "w");
try {
pipes[pipe_id].eptr = nullptr;
const int L = 4096;
char buf[L];
while (true) {
int c = fgetc(inf);
if (c == EOF) {
if (errno) {
throw system_error(errno, system_category());
}
break;
}
if (fputc(c, ouf) == EOF) {
throw system_error(errno, system_category());
}
fflush(ouf);
if (fputc(c, sf) == EOF) {
throw system_error(errno, system_category());
}
}
} catch (exception &e) {
cerr << e.what() << endl;
pipes[pipe_id].eptr = current_exception();
}
fclose(sf);
fclose(inf);
fclose(ouf);
}
public:
RunInteraction(const RunInteractionConfig &config) {
cmds.resize(config.cmds.size());
for (int i = 0; i < (int)config.cmds.size(); i++) {
cmds[i].cmd = config.cmds[i];
}
pipes.resize(config.pipes.size());
for (int i = 0; i < (int)config.pipes.size(); i++) {
pipes[i].config = config.pipes[i];
cmds[pipes[i].config.from].opipes.push_back(i);
cmds[pipes[i].config.to].ipipes.push_back(i);
}
for (int i = 0; i < (int)pipes.size(); i++) {
if (pipe(pipes[i].ipipefd) == -1 || pipe(pipes[i].opipefd) == -1) {
throw system_error(errno, system_category());
}
}
for (int i = 0; i < (int)cmds.size(); i++) {
cmds[i].pid = fork();
if (cmds[i].pid == 0) {
prepare_fd_for_cmd(i);
system(cmds[i].cmd.c_str());
exit(0);
} else if (cmds[i].pid == -1) {
throw system_error(errno, system_category());
}
}
prepare_fd();
for (int i = 0; i < (int)pipes.size(); i++) {
pipes[i].io_thread = thread(&RunInteraction::wait_pipe_io, this, i);
}
}
void join() {
int status;
while (wait(&status) > 0);
for (int i = 0; i < (int)pipes.size(); i++) {
pipes[i].io_thread.join();
}
}
};
error_t run_interaction_argp_parse_opt (int key, char *arg, struct argp_state *state) {
RunInteractionConfig *config = (RunInteractionConfig*)state->input;
try {
switch (key) {
case 'p':
config->pipes.push_back(PipeConfig(arg));
break;
case 's':
if (config->pipes.empty()) {
argp_usage(state);
}
config->pipes.back().saving_file_name = arg;
break;
case ARGP_KEY_ARG:
config->cmds.push_back(arg);
break;
case ARGP_KEY_END:
if (state->arg_num == 0) {
argp_usage(state);
}
break;
default:
return ARGP_ERR_UNKNOWN;
}
} catch (exception &e) {
argp_usage(state);
}
return 0;
}
RunInteractionConfig parse_args(int argc, char **argv) {
argp_option run_interaction_argp_options[] = {
{"add-pipe" , 'p', "PIPE" , 0, "Add a pipe <from>:<fd>-<to>:<fd> (fd < 128)" , 1},
{"save-pipe" , 's', "FILE" , 0, "Set last pipe saving file" , 2},
{0}
};
char run_interaction_argp_args_doc[] = "cmd1 cmd2 ...";
char run_interaction_argp_doc[] = "run_interaction: a tool to run multiple commands with interaction";
argp run_interaction_argp = {
run_interaction_argp_options,
run_interaction_argp_parse_opt,
run_interaction_argp_args_doc,
run_interaction_argp_doc
};
RunInteractionConfig config;
argp_parse(&run_interaction_argp, argc, argv, ARGP_NO_ARGS | ARGP_IN_ORDER, 0, &config);
return config;
}
int main(int argc, char **argv) {
signal(SIGPIPE, SIG_IGN);
RunInteraction ri(parse_args(argc, argv));
ri.join();
return 0;
}