CFBox/src/applets/timeout.cpp at main · Awesome-Embedded-Learning-Studio/CFBox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#include <cstdio>
#include <cstring>
#include <csignal>
#include <string>
#include <vector>
#include <sys/wait.h>
#include <unistd.h>

#include <cfbox/args.hpp>
#include <cfbox/help.hpp>
#include <cfbox/error.hpp>

namespace {
constexpr cfbox::help::HelpEntry HELP = {
    .name    = "timeout",
    .version = CFBOX_VERSION_STRING,
    .one_line = "run a command with a time limit",
    .usage   = "timeout DURATION COMMAND [ARGS]...",
    .options = "  -s SIG   send SIG on timeout (default: TERM)\n"
               "  -k DUR   also send KILL after DUR",
    .extra   = "DURATION supports floating point (e.g. 2.5 for 2.5 seconds).",
};
} // namespace

auto timeout_main(int argc, char* argv[]) -> int {
    auto parsed = cfbox::args::parse(argc, argv, {
        cfbox::args::OptSpec{'s', true, "signal"},
        cfbox::args::OptSpec{'k', true, "kill-after"},
    });

    if (parsed.has_long("help"))    { cfbox::help::print_help(HELP); return 0; }
    if (parsed.has_long("version")) { cfbox::help::print_version(HELP); return 0; }

    const auto& pos = parsed.positional();
    if (pos.size() < 2) {
        CFBOX_ERR("timeout", "missing operand");
        return 1;
    }

    char* end = nullptr;
    double duration = std::strtod(std::string{pos[0]}.c_str(), &end);
    if (end == std::string{pos[0]}.c_str() || duration <= 0) {
        CFBOX_ERR("timeout", "invalid duration '%.*s'", static_cast<int>(pos[0].size()), pos[0].data());
        return 1;
    }

    int sig = SIGTERM;
    if (auto s = parsed.get_any('s', "signal")) {
        auto name = std::string{*s};
        if (name == "TERM" || name == "15") sig = SIGTERM;
        else if (name == "KILL" || name == "9") sig = SIGKILL;
        else if (name == "INT" || name == "2") sig = SIGINT;
        else if (name == "HUP" || name == "1") sig = SIGHUP;
        else {
            CFBOX_ERR("timeout", "unknown signal '%s'", name.c_str());
            return 1;
        }
    }

    pid_t pid = fork();
    if (pid < 0) {
        CFBOX_ERR("timeout", "fork failed: %s", std::strerror(errno));
        return 1;
    }

    if (pid == 0) {
        // Child
        std::vector<std::string> arg_storage;
        arg_storage.reserve(pos.size() - 1);
        for (std::size_t i = 1; i < pos.size(); ++i) {
            arg_storage.emplace_back(pos[i]);
        }
        std::vector<char*> cmd_args;
        cmd_args.reserve(arg_storage.size() + 1);
        for (auto& s : arg_storage) cmd_args.push_back(s.data());
        cmd_args.push_back(nullptr);
        execvp(cmd_args[0], cmd_args.data());
        CFBOX_ERR("timeout", "failed to execute '%s': %s", cmd_args[0], std::strerror(errno));
        _exit(125);
    }

    // Parent: wait with timeout
    auto usecs = static_cast<useconds_t>(duration * 1000000);
    usleep(usecs);

    int status;
    pid_t result = waitpid(pid, &status, WNOHANG);
    if (result == 0) {
        // Still running, send signal
        kill(pid, sig);
        if (auto k = parsed.get_any('k', "kill-after")) {
            double kill_dur = std::strtod(std::string{*k}.c_str(), nullptr);
            auto kill_usecs = static_cast<useconds_t>(kill_dur * 1000000);
            usleep(kill_usecs);
            result = waitpid(pid, &status, WNOHANG);
            if (result == 0) {
                kill(pid, SIGKILL);
                waitpid(pid, &status, 0);
            }
        } else {
            waitpid(pid, &status, 0);
        }
        return 124;
    }

    if (WIFEXITED(status)) return WEXITSTATUS(status);
    if (WIFSIGNALED(status)) return 128 + WTERMSIG(status);
    return 1;
}