]> ruderich.org/simon Gitweb - coloredstderr/coloredstderr.git/blob - src/trackfds.h
tests/example_exec.c: Sort environment variables.
[coloredstderr/coloredstderr.git] / src / trackfds.h
1 /*
2  * Utility functions to track file descriptors.
3  *
4  * Copyright (C) 2013  Simon Ruderich
5  *
6  * This program is free software: you can redistribute it and/or modify
7  * it under the terms of the GNU General Public License as published by
8  * the Free Software Foundation, either version 3 of the License, or
9  * (at your option) any later version.
10  *
11  * This program is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14  * GNU General Public License for more details.
15  *
16  * You should have received a copy of the GNU General Public License
17  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
18  */
19
20 #ifndef TRACKFDS_H
21 #define TRACKFDS_H 1
22
23 /* Array of tracked file descriptors. Used for fast lookups for the normally
24  * used file descriptors (0 <= fd < TRACKFDS_STATIC_COUNT). */
25 static int tracked_fds[TRACKFDS_STATIC_COUNT];
26
27 /* List of tracked file descriptors >= TRACKFDS_STATIC_COUNT. */
28 static int *tracked_fds_list;
29 /* Current number of items in the list. */
30 static size_t tracked_fds_list_count;
31 /* Allocated items, used to reduce realloc()s. */
32 static size_t tracked_fds_list_space;
33
34
35 #ifdef DEBUG
36 static void tracked_fds_debug(void) {
37     size_t i;
38
39     for (i = 0; i < TRACKFDS_STATIC_COUNT; i++) {
40         if (tracked_fds[i]) {
41             debug("    tracked_fds[%d]: %d\n", i, tracked_fds[i]);
42         }
43     }
44     debug("    tracked_fds_list: %d/%d\t[%d]\n", tracked_fds_list_count,
45                                                  tracked_fds_list_space,
46                                                  getpid());
47     for (i = 0; i < tracked_fds_list_count; i++) {
48         debug("    tracked_fds_list[%d]: %d\n", i, tracked_fds_list[i]);
49     }
50 }
51 #endif
52
53 static int init_tracked_fds_list(size_t count) {
54     assert(count > 0);
55
56     /* Reduce reallocs. */
57     count += TRACKFDS_REALLOC_STEP;
58
59     tracked_fds_list = malloc(count * sizeof(*tracked_fds_list));
60     if (!tracked_fds_list) {
61 #ifdef WARNING
62         warning("malloc(tracked_fds_list, %d) failed [%d]\n",
63                 count * sizeof(*tracked_fds_list), getpid());
64 #endif
65         return 0;
66     }
67
68     tracked_fds_list_space = count;
69     return 1;
70 }
71
72 /* Load tracked file descriptors from the environment. The environment is used
73  * to pass the information to child processes.
74  *
75  * ENV_NAME_FDS and ENV_NAME_PRIVATE_FDS have the following format: Each
76  * descriptor as string followed by a comma; there's a trailing comma.
77  * Example: "2,4,". */
78 static void init_from_environment(void) {
79 #ifdef DEBUG
80     debug("init_from_environment()\t\t[%d]\n", getpid());
81 #endif
82     char const *env;
83
84     int saved_errno = errno;
85
86     initialized = 1;
87     tracked_fds_list_count = 0;
88
89     /* If ENV_NAME_FORCE_WRITE is set and not empty, allow writes to a non-tty
90      * device. Use with care! Mainly used for the test suite. */
91     env = getenv(ENV_NAME_FORCE_WRITE);
92     if (env && env[0] != '\0') {
93         force_write_to_non_tty = 1;
94     }
95
96     /* Prefer user defined list of file descriptors, fall back to file
97      * descriptors passed through the environment from the parent process. */
98     env = getenv(ENV_NAME_FDS);
99     if (env) {
100         used_fds_set_by_user = 1;
101     } else {
102         env = getenv(ENV_NAME_PRIVATE_FDS);
103     }
104     if (!env) {
105         errno = saved_errno;
106         return;
107     }
108 #ifdef DEBUG
109     debug("  getenv(\"%s\"): \"%s\"\n", ENV_NAME_FDS, env);
110     debug("  getenv(\"%s\"): \"%s\"\n", ENV_NAME_PRIVATE_FDS, env);
111 #endif
112
113     /* Environment must be treated read-only. */
114     char env_copy[strlen(env) + 1];
115     strcpy(env_copy, env);
116
117     char *x;
118
119     size_t count = 0;
120     for (x = env_copy; *x; x++) {
121         if (*x == ',') {
122             count++;
123         }
124     }
125
126     size_t i = 0;
127
128     /* Parse file descriptor numbers from environment string and store them as
129      * integers in tracked_fds and tracked_fds_list. */
130     char *last;
131     for (x = env_copy, last = env_copy; *x; x++) {
132         if (*x != ',') {
133             continue;
134         }
135         /* ',' at the beginning or double ',' - ignore. */
136         if (x == last) {
137             goto next;
138         }
139
140         if (i == count) {
141             break;
142         }
143
144         /* Replace ',' to null-terminate number for atoi(). */
145         *x = 0;
146
147         int fd = atoi(last);
148         if (fd < 0) {
149             goto next;
150
151         } else if (fd < TRACKFDS_STATIC_COUNT) {
152             tracked_fds[fd] = 1;
153         } else {
154             if (!tracked_fds_list) {
155                 /* Pessimistic count estimate, but allocating a few more
156                  * elements doesn't hurt. */
157                 if (!init_tracked_fds_list(count)) {
158                     /* Couldn't allocate memory, skip this entry. */
159                     goto next;
160                 }
161             }
162             tracked_fds_list[i++] = fd;
163 #ifdef DEBUG
164             debug("  large fd: %d\n", fd);
165 #endif
166         }
167
168 next:
169         last = x + 1;
170     }
171
172     tracked_fds_list_count = i;
173
174 #ifdef DEBUG
175     tracked_fds_debug();
176 #endif
177
178     errno = saved_errno;
179 }
180
181 static char *update_environment_buffer_entry(char *x, int fd) {
182     assert(fd >= 0);
183
184     int length = snprintf(x, 10 + 1, "%d", fd);
185     if (length >= 10 + 1 || length <= 0 /* shouldn't happen */) {
186         /* Integer too big to fit the buffer, skip it. */
187 #ifdef WARNING
188         warning("update_environment_buffer_entry(): truncated fd: %d [%d]\n",
189                 fd, getpid());
190 #endif
191         return x;
192     }
193
194     /* Write comma after number. */
195     x += length;
196     *x++ = ',';
197     /* Make sure the string is always null-terminated. */
198     *x = 0;
199
200     return x;
201 }
202 static void update_environment_buffer(char *x) {
203     assert(initialized);
204
205     size_t i;
206     for (i = 0; i < TRACKFDS_STATIC_COUNT; i++) {
207         if (tracked_fds[i]) {
208             x = update_environment_buffer_entry(x, (int)i);
209         }
210     }
211     for (i = 0; i < tracked_fds_list_count; i++) {
212         x = update_environment_buffer_entry(x, tracked_fds_list[i]);
213     }
214 }
215 inline static size_t update_environment_buffer_size(void) {
216     assert(initialized);
217
218     /* Use the maximum count (TRACKFDS_STATIC_COUNT) of used descriptors
219      * because it's simple and small enough not to be a problem.
220      *
221      * An integer (32-bit) has at most 10 digits, + 1 for the comma after each
222      * number. Bigger file descriptors (which shouldn't occur in reality) are
223      * skipped. */
224     return (TRACKFDS_STATIC_COUNT + tracked_fds_list_count)
225                * (10 + 1) + 1 /* to fit '\0' */;
226 }
227 static void update_environment(void) {
228 #ifdef DEBUG
229     debug("update_environment()\t\t[%d]\n", getpid());
230 #endif
231
232     /* If we haven't parsed the environment we also haven't modified it - so
233      * nothing to do. */
234     if (!initialized) {
235         return;
236     }
237
238     int saved_errno = errno;
239
240     char env[update_environment_buffer_size()];
241     env[0] = 0;
242
243     update_environment_buffer(env);
244
245 #if 0
246     debug("    setenv(\"%s\", \"%s\", 1)\n", ENV_NAME_PRIVATE_FDS, env);
247 #endif
248     setenv(ENV_NAME_PRIVATE_FDS, env, 1 /* overwrite */);
249
250     /* Child processes must use ENV_NAME_PRIVATE_FDS to get the updated list
251      * of tracked file descriptors, not the static list provided by the user
252      * in ENV_NAME_FDS.
253      *
254      * But only remove it if the static list in ENV_NAME_FDS was loaded by
255      * init_from_environment() and merged into ENV_NAME_PRIVATE_FDS. */
256     if (used_fds_set_by_user) {
257         unsetenv(ENV_NAME_FDS);
258     }
259
260     errno = saved_errno;
261 }
262
263
264
265 static void tracked_fds_add(int fd) {
266     assert(fd >= 0);
267
268     if (fd < TRACKFDS_STATIC_COUNT) {
269         tracked_fds[fd] = 1;
270 #if 0
271         debug("tracked_fds_add(): %-3d\t\t[%d]\n", fd, getpid());
272         tracked_fds_debug();
273 #endif
274         return;
275     }
276
277     if (tracked_fds_list_count >= tracked_fds_list_space) {
278         int saved_errno = errno;
279
280         size_t new_space = tracked_fds_list_space + TRACKFDS_REALLOC_STEP;
281         int *tmp = realloc(tracked_fds_list,
282                            sizeof(*tracked_fds_list) * new_space);
283         if (!tmp) {
284             /* We can do nothing, just ignore the error. We made sure not to
285              * destroy our state, so the new descriptor is ignored without any
286              * other consequences. */
287 #ifdef WARNING
288             warning("realloc(tracked_fds_list, %zu) failed! [%d]\n",
289                     sizeof(*tracked_fds_list) * new_space, getpid());
290 #endif
291             errno = saved_errno;
292             return;
293         }
294         errno = saved_errno;
295
296         tracked_fds_list = tmp;
297         tracked_fds_list_space = new_space;
298     }
299
300     tracked_fds_list[tracked_fds_list_count++] = fd;
301
302 #ifdef DEBUG
303     debug("tracked_fds_add(): %-3d\t\t[%d]\n", fd, getpid());
304     tracked_fds_debug();
305 #endif
306 }
307 static int tracked_fds_remove(int fd) {
308     assert(fd >= 0);
309
310     if (fd < TRACKFDS_STATIC_COUNT) {
311         int old_value = tracked_fds[fd];
312         tracked_fds[fd] = 0;
313
314 #if 0
315         debug("tracked_fds_remove(): %-3d\t[%d]\n", fd, getpid());
316         tracked_fds_debug();
317 #endif
318         return old_value; /* Found vs. not found. */
319     }
320
321     size_t i;
322     for (i = 0; i < tracked_fds_list_count; i++) {
323         if (fd != tracked_fds_list[i]) {
324             continue;
325         }
326
327         memmove(tracked_fds_list + i, tracked_fds_list + i + 1,
328                 sizeof(*tracked_fds_list) * (tracked_fds_list_count - i - 1));
329         tracked_fds_list_count--;
330
331 #ifdef DEBUG
332         debug("tracked_fds_remove(): %-3d\t[%d]\n", fd, getpid());
333         tracked_fds_debug();
334 #endif
335
336         /* Found. */
337         return 1;
338     }
339
340     /* Not found. */
341     return 0;
342 }
343
344 static int tracked_fds_find_slow(int fd) noinline;
345 /*
346  * tracked_fds_find() is called for each hook call and should be as fast as
347  * possible. As most file descriptors are < TRACKFDS_STATIC_COUNT, force the
348  * compiler to inline that part which is almost exclusively used.
349  *
350  * Inlining tracked_fds_add()/tracked_fds_remove() isn't worth the effort as
351  * they are not called often enough.
352  */
353 inline static int tracked_fds_find(int fd) always_inline;
354 inline static int tracked_fds_find(int fd) {
355     /* Invalid file descriptor. No assert() as we're called from the hooked
356      * macro. */
357     if (unlikely(fd < 0)) {
358         return 0;
359     }
360
361     if (fd < TRACKFDS_STATIC_COUNT) {
362         return tracked_fds[fd];
363     }
364
365     return tracked_fds_find_slow(fd);
366 }
367 static int tracked_fds_find_slow(int fd) {
368     assert(initialized);
369     assert(fd >= 0);
370
371     if (tracked_fds_list_count == 0) {
372         return 0;
373     }
374
375     size_t i;
376     for (i = 0; i < tracked_fds_list_count; i++) {
377         if (fd == tracked_fds_list[i]) {
378             return 1;
379         }
380     }
381     return 0;
382 }
383
384 #endif