pstree: fix width logic. +30 bytes
[oweals/busybox.git] / procps / pstree.c
1 /*
2  * pstree.c - display process tree
3  *
4  * Copyright (C) 1993-2002 Werner Almesberger
5  * Copyright (C) 2002-2009 Craig Small
6  * Copyright (C) 2010 Lauri Kasanen
7  *
8  * Based on pstree (PSmisc) 22.13.
9  *
10  * Licensed under GPLv2, see file LICENSE in this source tree.
11  */
12
13 //config:config PSTREE
14 //config:       bool "pstree"
15 //config:       default y
16 //config:       help
17 //config:         Display a tree of processes.
18
19 //applet:IF_PSTREE(APPLET(pstree, _BB_DIR_USR_BIN, _BB_SUID_DROP))
20
21 //kbuild:lib-$(CONFIG_PSTREE) += pstree.o
22
23 //usage:#define pstree_trivial_usage
24 //usage:        "[-p] [PID|USER]"
25 //usage:#define pstree_full_usage "\n\n"
26 //usage:       "Display process tree, optionally start from USER or PID\n"
27 //usage:     "\nOptions:"
28 //usage:     "\n        -p      Show pids"
29
30 #include "libbb.h"
31
32 #define PROC_BASE "/proc"
33
34 #define OPT_PID  (1 << 0)
35
36 struct child;
37
38 typedef struct proc {
39         char comm[COMM_LEN + 1];
40 //      char flags; - unused, delete?
41         pid_t pid;
42         uid_t uid;
43         struct child *children;
44         struct proc *parent;
45         struct proc *next;
46 } PROC;
47
48 /* For flags above */
49 //#define PFLAG_THREAD  0x01
50
51 typedef struct child {
52         PROC *child;
53         struct child *next;
54 } CHILD;
55
56 #define empty_2  "  "
57 #define branch_2 "|-"
58 #define vert_2   "| "
59 #define last_2   "`-"
60 #define single_3 "---"
61 #define first_3  "-+-"
62
63 struct globals {
64         /* 0-based. IOW: the number of chars we printer on current line */
65         unsigned cur_x;
66         unsigned output_width;
67
68         /* The buffers will be dynamically increased in size as needed */
69         unsigned capacity;
70         unsigned *width;
71         uint8_t *more;
72
73         PROC *list;
74
75         smallint dumped; /* used by dump_by_user */
76 };
77 #define G (*ptr_to_globals)
78 #define INIT_G() do { \
79         SET_PTR_TO_GLOBALS(xzalloc(sizeof(G))); \
80 } while (0)
81
82
83 /*
84  * Allocates additional buffer space for width and more as needed.
85  * The first call will allocate the first buffer.
86  *
87  * bufindex  the index that will be used after the call to this function.
88  */
89 static void ensure_buffer_capacity(int bufindex)
90 {
91         if (bufindex >= G.capacity) {
92                 G.capacity += 0x100;
93                 G.width = xrealloc(G.width, G.capacity * sizeof(G.width[0]));
94                 G.more = xrealloc(G.more, G.capacity * sizeof(G.more[0]));
95         }
96 }
97
98 #if ENABLE_FEATURE_CLEAN_UP
99 static void maybe_free_buffers(void)
100 {
101         free(G.width);
102         free(G.more);
103 }
104 #else
105 # define maybe_free_buffers() ((void)0)
106 #endif
107
108 /* NB: this function is never called with "bad" chars
109  * (control chars or chars >= 0x7f)
110  */
111 static void out_char(char c)
112 {
113         G.cur_x++;
114         if (G.cur_x == G.output_width)
115                 c = '+';
116         if (G.cur_x <= G.output_width)
117                 putchar(c);
118 }
119
120 /* NB: this function is never called with "bad" chars
121  * (control chars or chars >= 0x7f)
122  */
123 static void out_string(const char *str)
124 {
125         while (*str)
126                 out_char(*str++);
127 }
128
129 static void out_newline(void)
130 {
131         putchar('\n');
132         G.cur_x = 0;
133 }
134
135 static PROC *find_proc(pid_t pid)
136 {
137         PROC *walk;
138
139         for (walk = G.list; walk; walk = walk->next)
140                 if (walk->pid == pid)
141                         break;
142
143         return walk;
144 }
145
146 static PROC *new_proc(const char *comm, pid_t pid, uid_t uid)
147 {
148         PROC *new = xzalloc(sizeof(*new));
149
150         strcpy(new->comm, comm);
151         new->pid = pid;
152         new->uid = uid;
153         new->next = G.list;
154
155         G.list = new;
156         return G.list;
157 }
158
159 static void add_child(PROC *parent, PROC *child)
160 {
161         CHILD *new, **walk;
162         int cmp;
163
164         new = xmalloc(sizeof(*new));
165
166         new->child = child;
167         for (walk = &parent->children; *walk; walk = &(*walk)->next) {
168                 cmp = strcmp((*walk)->child->comm, child->comm);
169                 if (cmp > 0)
170                         break;
171                 if (cmp == 0 && (*walk)->child->uid > child->uid)
172                         break;
173         }
174         new->next = *walk;
175         *walk = new;
176 }
177
178 static void add_proc(const char *comm, pid_t pid, pid_t ppid,
179                         uid_t uid /*, char isthread*/)
180 {
181         PROC *this, *parent;
182
183         this = find_proc(pid);
184         if (!this)
185                 this = new_proc(comm, pid, uid);
186         else {
187                 strcpy(this->comm, comm);
188                 this->uid = uid;
189         }
190
191         if (pid == ppid)
192                 ppid = 0;
193 //      if (isthread)
194 //              this->flags |= PFLAG_THREAD;
195
196         parent = find_proc(ppid);
197         if (!parent)
198                 parent = new_proc("?", ppid, 0);
199
200         add_child(parent, this);
201         this->parent = parent;
202 }
203
204 static int tree_equal(const PROC *a, const PROC *b)
205 {
206         const CHILD *walk_a, *walk_b;
207
208         if (strcmp(a->comm, b->comm) != 0)
209                 return 0;
210         if ((option_mask32 /*& OPT_PID*/) && a->pid != b->pid)
211                 return 0;
212
213         for (walk_a = a->children, walk_b = b->children;
214           walk_a && walk_b;
215           walk_a = walk_a->next, walk_b = walk_b->next
216         ) {
217                 if (!tree_equal(walk_a->child, walk_b->child))
218                         return 0;
219         }
220
221         return !(walk_a || walk_b);
222 }
223
224 static int out_args(const char *mystr)
225 {
226         const char *here;
227         int strcount = 0;
228         char tmpstr[5];
229
230         for (here = mystr; *here; here++) {
231                 if (*here == '\\') {
232                         out_string("\\\\");
233                         strcount += 2;
234                 } else if (*here >= ' ' && *here < 0x7f) {
235                         out_char(*here);
236                         strcount++;
237                 } else {
238                         sprintf(tmpstr, "\\%03o", (unsigned char) *here);
239                         out_string(tmpstr);
240                         strcount += 4;
241                 }
242         }
243
244         return strcount;
245 }
246
247 static void
248 dump_tree(PROC *current, int level, int rep, int leaf, int last, int closing)
249 {
250         CHILD *walk, *next, **scan;
251         int lvl, i, add, offset, count, comm_len, first;
252         char tmp[sizeof(int)*3 + 4];
253
254         if (!current)
255                 return;
256
257         if (!leaf) {
258                 for (lvl = 0; lvl < level; lvl++) {
259                         i = G.width[lvl] + 1;
260                         while (--i >= 0)
261                                 out_char(' ');
262
263                         if (lvl == level - 1) {
264                                 if (last) {
265                                         out_string(last_2);
266                                 } else {
267                                         out_string(branch_2);
268                                 }
269                         } else {
270                                 if (G.more[lvl + 1]) {
271                                         out_string(vert_2);
272                                 } else {
273                                         out_string(empty_2);
274                                 }
275                         }
276                 }
277         }
278
279         add = 0;
280         if (rep > 1) {
281                 add += sprintf(tmp, "%d*[", rep);
282                 out_string(tmp);
283         }
284         comm_len = out_args(current->comm);
285         if (option_mask32 /*& OPT_PID*/) {
286                 comm_len += sprintf(tmp, "(%d)", (int)current->pid);
287                 out_string(tmp);
288         }
289         offset = G.cur_x;
290
291         if (!current->children) {
292                 while (closing--)
293                         out_char(']');
294                 out_newline();
295         }
296         ensure_buffer_capacity(level);
297         G.more[level] = !last;
298
299         G.width[level] = comm_len + G.cur_x - offset + add;
300         if (G.cur_x >= G.output_width) {
301                 //out_string(first_3); - why? it won't print anything
302                 //out_char('+');
303                 out_newline();
304                 return;
305         }
306
307         first = 1;
308         for (walk = current->children; walk; walk = next) {
309                 count = 0;
310                 next = walk->next;
311                 scan = &walk->next;
312                 while (*scan) {
313                         if (!tree_equal(walk->child, (*scan)->child))
314                                 scan = &(*scan)->next;
315                         else {
316                                 if (next == *scan)
317                                         next = (*scan)->next;
318                                 count++;
319                                 *scan = (*scan)->next;
320                         }
321                 }
322                 if (first) {
323                         out_string(next ? first_3 : single_3);
324                         first = 0;
325                 }
326
327                 dump_tree(walk->child, level + 1, count + 1,
328                                 walk == current->children, !next,
329                                 closing + (count ? 1 : 0));
330         }
331 }
332
333 static void dump_by_user(PROC *current, uid_t uid)
334 {
335         const CHILD *walk;
336
337         if (!current)
338                 return;
339
340         if (current->uid == uid) {
341                 if (G.dumped)
342                         putchar('\n');
343                 dump_tree(current, 0, 1, 1, 1, 0);
344                 G.dumped = 1;
345                 return;
346         }
347         for (walk = current->children; walk; walk = walk->next)
348                 dump_by_user(walk->child, uid);
349 }
350
351 static void handle_thread(const char *comm, pid_t pid, pid_t ppid, uid_t uid)
352 {
353         char threadname[COMM_LEN + 2];
354         sprintf(threadname, "{%.*s}", COMM_LEN - 2, comm);
355         add_proc(threadname, pid, ppid, uid/*, 1*/);
356 }
357
358 static void mread_proc(void)
359 {
360         procps_status_t *p = NULL;
361         pid_t parent = 0;
362         int flags = PSSCAN_COMM | PSSCAN_PID | PSSCAN_PPID | PSSCAN_UIDGID | PSSCAN_TASKS;
363
364         while ((p = procps_scan(p, flags)) != NULL) {
365 #if ENABLE_FEATURE_SHOW_THREADS
366                 if (p->pid != p->main_thread_pid)
367                         handle_thread(p->comm, p->pid, parent, p->uid);
368                 else
369 #endif
370                 {
371                         add_proc(p->comm, p->pid, p->ppid, p->uid/*, 0*/);
372                         parent = p->pid;
373                 }
374         }
375 }
376
377 int pstree_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
378 int pstree_main(int argc UNUSED_PARAM, char **argv)
379 {
380         pid_t pid = 1;
381         long uid = 0;
382
383         INIT_G();
384
385         get_terminal_width_height(1, &G.output_width, NULL);
386
387         getopt32(argv, "p");
388         argv += optind;
389
390         if (argv[0]) {
391                 if (argv[1])
392                         bb_show_usage();
393                 if (argv[0][0] >= '0' && argv[0][0] <= '9') {
394                         pid = xatoi(argv[0]);
395                 } else {
396                         uid = xuname2uid(argv[0]);
397                 }
398         }
399
400         mread_proc();
401
402         if (!uid)
403                 dump_tree(find_proc(pid), 0, 1, 1, 1, 0);
404         else {
405                 dump_by_user(find_proc(1), uid);
406                 if (!G.dumped) {
407                         bb_error_msg_and_die("no processes found");
408                 }
409         }
410
411         maybe_free_buffers();
412         return 0;
413 }