1 | // SPDX-License-Identifier: GPL-3.0-or-later |
2 | #include "global.hpp" |
3 | #include "unit/unit.hpp" |
4 | |
5 | #include <argparse/libargparse.h> |
6 | #include <filesystem> |
7 | #include <functional> |
8 | #include <glob.h> |
9 | #include <iostream> |
10 | #include <map> |
11 | #include <set> |
12 | #include <signal.h> |
13 | #include <spawn.h> |
14 | #include <stdio.h> |
15 | #include <string> |
16 | #include <sys/stat.h> |
17 | #include <sys/wait.h> |
18 | #include <toml++/toml.hpp> |
19 | #include <unistd.h> |
20 | #include <vector> |
21 | |
22 | #define RED(text) "\033[1;31m" text "\033[0m" |
23 | #define GREEN(text) "\033[1;32m" text "\033[0m" |
24 | |
25 | #define FAILED() RED("[FAILED]") |
26 | #define OK() GREEN("[ OK ]") |
27 | #define STARTING() "\033[0m " |
28 | |
29 | std::map<std::string, pid_t> service_pid; |
30 | bool debug = false; |
31 | inline GlobalConfig global_config; |
32 | inline std::map<std::string, std::shared_ptr<Unit>> units; |
33 | |
34 | static std::vector<std::string> startup_order_for_unit(const std::string &id) |
35 | { |
36 | std::set<std::string> visited; |
37 | std::vector<std::string> order; |
38 | |
39 | std::function<void(const std::string &)> visit = [&](const std::string id) |
40 | { |
41 | if (visited.contains(x: id)) |
42 | return; |
43 | |
44 | visited.insert(x: id); |
45 | const auto &unit = units[id]; |
46 | if (!unit) |
47 | { |
48 | std::cerr << "unit " << id << " does not exist" << std::endl; |
49 | return; |
50 | } |
51 | |
52 | for (const auto &dep_id : unit->depends_on) |
53 | visit(dep_id); |
54 | order.push_back(x: id); |
55 | }; |
56 | |
57 | visit(id); |
58 | return order; |
59 | } |
60 | |
61 | static bool start_unit_tree(const std::string &id) |
62 | { |
63 | const auto order = startup_order_for_unit(id); |
64 | for (const auto &unit_id : order) |
65 | { |
66 | const auto unit = units[unit_id]; |
67 | |
68 | if (debug) |
69 | std::cout << STARTING() << "Starting " << unit->description << " (" << unit->id << ")" << std::endl; |
70 | if (!unit->start()) |
71 | { |
72 | std::cerr << FAILED() << " Failed to start " << unit->description << ": " << unit->error_reason() << std::endl; |
73 | return false; |
74 | } |
75 | else |
76 | { |
77 | if (unit->type == "target" ) |
78 | std::cout << OK() << " Reached target " << unit->description << std::endl; |
79 | else |
80 | std::cout << OK() << " Started " << unit->description << std::endl; |
81 | } |
82 | } |
83 | |
84 | return true; |
85 | } |
86 | |
87 | static void sigsegv_handler(int sig) |
88 | { |
89 | if (sig == SIGSEGV) |
90 | { |
91 | std::cout << RED("Segmentation fault" ) << std::endl; |
92 | while (true) |
93 | sched_yield(); |
94 | } |
95 | } |
96 | |
97 | static void sigchild_handler(int sig) |
98 | { |
99 | MOS_UNUSED(sig); |
100 | } |
101 | |
102 | #define DYN_ERROR_CODE (__COUNTER__ + 1) |
103 | |
104 | static const argparse_arg_t longopts[] = { |
105 | { .full: "help" , .abbr: 'h', .argtype: ARGPARSE_NONE, .help: "show this help" }, |
106 | { .full: "config" , .abbr: 'C', .argtype: ARGPARSE_REQUIRED, .help: "configuration file, default: /initrd/config/init.conf" }, |
107 | { .full: "shell" , .abbr: 'S', .argtype: ARGPARSE_REQUIRED, .help: "shell to start, default: /initrd/programs/mossh" }, |
108 | {}, |
109 | }; |
110 | |
111 | int main(int argc, const char *argv[]) |
112 | { |
113 | struct sigaction sa; |
114 | sa.sa_handler = SIG_IGN; |
115 | sa.sa_flags = SA_RESTART; |
116 | sigaction(SIGINT, &sa, NULL); |
117 | sigaction(SIGQUIT, &sa, NULL); |
118 | sigaction(SIGTERM, &sa, NULL); |
119 | |
120 | sa.sa_handler = sigchild_handler; |
121 | sigaction(SIGCHLD, &sa, NULL); |
122 | |
123 | sa.sa_handler = sigsegv_handler; |
124 | sigaction(SIGSEGV, &sa, NULL); |
125 | |
126 | std::filesystem::path config_file = "/initrd/config/init-config.toml" ; |
127 | std::string shell = "/initrd/programs/mossh" ; |
128 | argparse_state_t state; |
129 | argparse_init(options: &state, argv); |
130 | while (true) |
131 | { |
132 | const int option = argparse_long(options: &state, longopts, NULL); |
133 | if (option == -1) |
134 | break; |
135 | |
136 | switch (option) |
137 | { |
138 | case 'C': config_file = state.optarg; break; |
139 | case 'S': shell = state.optarg; break; |
140 | case 'h': argparse_usage(options: &state, args: longopts, usage: "the init program" ); return 0; |
141 | default: break; |
142 | } |
143 | } |
144 | |
145 | if (getpid() != 1) |
146 | { |
147 | for (int i = 0; i < argc; i++) |
148 | printf(format: "argv[%d] = %s\n" , i, argv[i]); |
149 | puts(string: "init: not running as PID 1, exiting..." ); |
150 | return DYN_ERROR_CODE; |
151 | } |
152 | |
153 | if (debug) |
154 | std::cout << "init: using config file " << config_file << std::endl; |
155 | |
156 | if (!std::filesystem::exists(p: config_file)) |
157 | { |
158 | std::cerr << "init: config file " << config_file << " does not exist" << std::endl; |
159 | return DYN_ERROR_CODE; |
160 | } |
161 | |
162 | load_configurations(config_path: config_file); |
163 | |
164 | if (!start_unit_tree(id: global_config.default_target)) |
165 | { |
166 | std::cerr << RED("init: failed to start default target" ) << std::endl; |
167 | return DYN_ERROR_CODE; |
168 | } |
169 | |
170 | // start the shell |
171 | const char **shell_argv = (const char **) malloc(size: sizeof(char *)); |
172 | int shell_argc = 1; |
173 | shell_argv[0] = shell.c_str(); |
174 | |
175 | const char *arg; |
176 | argparse_init(options: &state, argv); // reset the options |
177 | while ((arg = argparse_arg(options: &state))) |
178 | { |
179 | shell_argc++; |
180 | shell_argv = (const char **) realloc(pointer: shell_argv, size: shell_argc * sizeof(char *)); |
181 | shell_argv[shell_argc - 1] = arg; |
182 | } |
183 | shell_argv = (const char **) realloc(pointer: shell_argv, size: (shell_argc + 1) * sizeof(char *)); |
184 | shell_argv[shell_argc] = NULL; |
185 | |
186 | start_shell:; |
187 | const pid_t shell_pid = fork(); |
188 | if (shell_pid == 0) |
189 | if (execv(shell.c_str(), (char **) shell_argv) <= 0) |
190 | return DYN_ERROR_CODE; |
191 | |
192 | while (true) |
193 | { |
194 | int status = 0; |
195 | const pid_t pid = waitpid(pid: -1, status: &status, flags: 0); |
196 | if (pid == shell_pid) |
197 | { |
198 | puts(string: "init: shell exited, restarting..." ); |
199 | goto start_shell; |
200 | } |
201 | |
202 | if (pid > 0) |
203 | { |
204 | if (WIFEXITED(status)) |
205 | std::cout << "init: process " << pid << " exited with status " << WEXITSTATUS(status) << std::endl; |
206 | else if (WIFSIGNALED(status)) |
207 | std::cout << "init: process " << pid << " killed by signal " << WTERMSIG(status) << std::endl; |
208 | } |
209 | |
210 | // check if any service has exited |
211 | for (const auto &[id, spid] : service_pid) |
212 | { |
213 | if (spid == -1) |
214 | continue; |
215 | |
216 | if (pid == spid) |
217 | { |
218 | std::cout << "init: service " << id << " exited" << std::endl; |
219 | service_pid[id] = -1; |
220 | } |
221 | } |
222 | } |
223 | |
224 | return 0; |
225 | } |
226 | |