1#include <stdio.h>
2#include <stdlib.h>
3#include <stdint.h>
4#include <errno.h>
5#include <string.h>
6#include <unistd.h>
7#include <libproc.h>
8#include <sys/types.h>
9#include <sys/sysctl.h>
10#include <Kernel/kern/ledger.h>
11#include <mach/mach_types.h>
12
13extern int ledger(int cmd, caddr_t arg1, caddr_t arg2, caddr_t arg3);
14
15int pid = -1;
16char *group_print = NULL;
17char *resource_print = NULL;
18int diff_mode = 0;
19
20struct proc_list {
21	int pid;
22	int seen;
23	char command[32];
24	struct ledger *ledger;
25	struct proc_list *next;
26};
27
28struct proc_list *procs = NULL;
29struct ledger_template_info *template = NULL;
30int entry_cnt = 0;
31
32struct ledger {
33	int64_t id;
34	int seen;
35	int64_t entries;
36	struct ledger_entry_info *info;
37	struct ledger_entry_info *old_info;
38	struct ledger *next;
39};
40
41struct ledger *ledgers = NULL;
42
43static void
44get_template_info()
45{
46
47	void *buf;
48	int cnt;
49
50top:
51	/* Allocate enough space to accomodate a few new entries */
52	cnt = entry_cnt + 5;
53	buf = malloc(cnt * sizeof (struct ledger_template_info));
54	if (buf == NULL) {
55		fprintf(stderr, "Out of memory\n");
56		exit (1);
57	}
58
59	if (ledger(LEDGER_TEMPLATE_INFO, (caddr_t)buf, (caddr_t)&cnt, NULL) < 0) {
60        perror("ledger() system call failed");
61		exit(1);
62	}
63
64	/* We underestimated how many entries we needed.  Let's try again */
65	if (cnt == entry_cnt + 5) {
66		entry_cnt += 5;
67		free(buf);
68		goto top;
69	}
70	entry_cnt = cnt;
71	template = buf;
72}
73
74/*
75 * Note - this is a destructive operation.  Unless we're about to exit, this
76 * needs to be followed by another call to get_template_info().
77 */
78static void
79dump_template_info()
80{
81	int i, j;
82	const char *group = NULL;
83
84	printf("Resources being tracked:\n");
85	printf("\t%10s  %15s  %8s\n", "GROUP", "RESOURCE", "UNITS");
86	for (i = 0; i < entry_cnt; i++) {
87		if (strlen(template[i].lti_name) == 0)
88			continue;
89
90		group = template[i].lti_group;
91		for (j = i; j < entry_cnt; j++) {
92			if (strcmp(template[j].lti_group, group))
93				continue;
94			printf("\t%10s  %15s  %8s\n", template[j].lti_group,
95			    template[j].lti_name, template[j].lti_units);
96			template[j].lti_name[0] = '\0';
97		}
98	}
99}
100
101static void
102validate_group()
103{
104	int i;
105
106	if (template == NULL)
107		get_template_info();
108
109	for (i = 0; i < entry_cnt; i++)
110		if (!strcmp(group_print, template[i].lti_group))
111			return;
112
113	fprintf(stderr, "No such group: %s\n", group_print);
114	exit (1);
115}
116
117static void
118validate_resource()
119{
120	int i;
121
122	if (template == NULL)
123		get_template_info();
124
125	for (i = 0; i < entry_cnt; i++)
126		if (!strcmp(resource_print, template[i].lti_name))
127			return;
128
129	fprintf(stderr, "No such resource: %s\n", resource_print);
130	exit (1);
131}
132
133static size_t
134get_kern_max_proc(void)
135{
136	int mib[] = { CTL_KERN, KERN_MAXPROC };
137	int max;
138	size_t max_sz = sizeof (max);
139
140	if (sysctl(mib, 2, &max, &max_sz, NULL, 0) < 0) {
141		perror("Failed to get max proc count");
142		exit (1);
143	}
144
145	return (max);
146}
147
148static int
149get_proc_kinfo(pid_t pid, struct kinfo_proc *kinfo)
150{
151	int mib[] = { CTL_KERN, KERN_PROC, KERN_PROC_PID, pid };
152	size_t len;
153
154	len = sizeof(struct kinfo_proc);
155	return (sysctl(mib, 4, kinfo, &len, NULL, 0) < 0);
156}
157
158static struct ledger *
159ledger_find(struct ledger_info *li)
160{
161	struct ledger *l;
162
163	for (l = ledgers; l && (li->li_id != l->id); l = l->next)
164		;
165
166	if (l == NULL) {
167		l = (struct ledger *)malloc(sizeof (*l));
168		if (l == NULL) {
169			fprintf(stderr, "Out of memory");
170			exit (1);
171		}
172		l->id = li->li_id;
173		l->entries = li->li_entries;
174		l->next = ledgers;
175		l->seen = 0;
176		l->info = NULL;
177		l->old_info = NULL;
178		ledgers = l;
179	}
180	return (l);
181
182}
183
184static void
185ledger_update(pid_t pid, struct ledger *l)
186{
187	void *arg;
188	struct ledger_entry_info *lei;
189	int64_t cnt;
190
191	cnt = l->entries;
192	if (cnt > entry_cnt)
193		cnt = entry_cnt;
194	arg = (void *)(long)pid;
195    lei = (struct ledger_entry_info *)malloc((size_t)(cnt * sizeof (*lei)));
196	if (ledger(LEDGER_ENTRY_INFO, arg, (caddr_t)lei, (caddr_t)&cnt) < 0) {
197    	perror("ledger_info() failed: ");
198		exit (1);
199	}
200	l->info = lei;
201}
202
203static void
204get_proc_info(int pid)
205{
206	struct ledger_info li;
207	struct ledger *ledgerp;
208	struct proc_list *proc;
209	struct kinfo_proc kinfo;
210	void *arg;
211
212	if (pid == 0)
213		return;
214
215	arg = (void *)(long)pid;
216    errno = 0;
217    if (ledger(LEDGER_INFO, arg, (caddr_t)&li, NULL) < 0) {
218
219		if (errno == ENOENT || errno == ESRCH)
220			return;
221
222		perror("ledger_info() failed: ");
223		exit (1);
224	}
225
226	ledgerp = ledger_find(&li);
227	ledger_update(pid, ledgerp);
228	ledgerp->seen = 1;
229
230	for (proc = procs; proc; proc = proc->next)
231		if (proc->pid == pid)
232			break;
233	if (proc == NULL) {
234		proc = (struct proc_list *)malloc(sizeof (*proc));
235		if (proc == NULL) {
236			fprintf(stderr, "Out of memory\n");
237			exit (1);
238		}
239
240		if (get_proc_kinfo(pid, &kinfo))
241			strlcpy(proc->command, "Error", sizeof (proc->command));
242		else
243			strlcpy(proc->command, kinfo.kp_proc.p_comm,
244			    sizeof (proc->command));
245
246		proc->pid = pid;
247		proc->ledger = ledgerp;
248		proc->next = procs;
249		procs = proc;
250	}
251	proc->seen = 1;
252}
253
254static int
255pid_compare(const void *a, const void *b)
256{
257	pid_t *pid_a = (pid_t *)a;
258	pid_t *pid_b = (pid_t *)b;
259
260	return (*pid_b - *pid_a);
261}
262
263static void
264get_all_info()
265{
266	pid_t *pids;
267	int sz, cnt, i;
268
269	if (pid < 0)
270		cnt = (int) get_kern_max_proc();
271	else
272		cnt = 1;
273
274	sz = cnt * sizeof(pid_t);
275	pids = (pid_t *)malloc(sz);
276	if (pids == NULL) {
277		perror("can't allocate memory for proc buffer\n");
278		exit (1);
279	}
280
281	if (pid < 0) {
282		cnt = proc_listallpids(pids, sz);
283		if (cnt < 0) {
284			perror("failed to get list of active pids");
285			exit (1);
286		}
287		qsort(pids, cnt, sizeof (pid_t), pid_compare);
288	} else {
289		pids[0] = pid;
290	}
291
292	for (i = 0; i < cnt; i++)
293		get_proc_info(pids[i]);
294	free(pids);
295}
296
297static void
298print_num(int64_t num, int64_t delta)
299{
300	char suf = ' ';
301	char posneg = ' ';
302
303	if (diff_mode) {
304		num = delta;
305	}
306
307	if (num == LEDGER_LIMIT_INFINITY) {
308		printf("%10s ", "-  ");
309		return;
310	}
311
312	if (llabs(num) > 10000000000) {
313		num /= 1000000000;
314		suf = 'G';
315	} else if (llabs(num) > 10000000) {
316		num /= 1000000;
317		suf = 'M';
318	} else if (llabs(num) > 100000) {
319		num /= 1000;
320		suf = 'K';
321	}
322	posneg = (delta < 0) ? '-' : ((delta > 0) ? '+' : ' ');
323
324	if (suf == ' ') {
325		suf = posneg;
326		posneg = ' ';
327	}
328	printf("%8lld%c%c ", num, suf, posneg);
329}
330
331static void
332dump_all_info()
333{
334	struct ledger_entry_info *info, *old;
335	struct proc_list *p;
336	int line, i;
337	int64_t d;
338
339	printf("\n%5s %10s %15s %10s %10s %10s %10s %10s\n", "PID", "COMMAND",
340	    "RESOURCE", "CREDITS", "DEBITS", "BALANCE", "LIMIT", "PERIOD");
341
342	for (p = procs; p; p = p->next) {
343		if (p->seen == 0)
344			continue;
345
346		printf("%5d %10.10s ", p->pid, p->command);
347		line = 0;
348
349		info = p->ledger->info;
350		old = p->ledger->old_info;
351		for (i = 0; i < p->ledger->entries; i++) {
352			if (group_print &&
353			    strcmp(group_print, template[i].lti_group))
354				continue;
355
356			if (resource_print &&
357			    strcmp(resource_print, template[i].lti_name))
358				continue;
359
360			if (line++)
361				printf("                 ");
362			printf("%15s ", template[i].lti_name);
363
364			d = old ? info[i].lei_credit - old[i].lei_credit : 0;
365			print_num(info[i].lei_credit, d);
366
367			d = old ? info[i].lei_debit - old[i].lei_debit : 0;
368			print_num(info[i].lei_debit, d);
369
370			d = old ? info[i].lei_balance - old[i].lei_balance : 0;
371			print_num(info[i].lei_balance, d);
372
373			if (info[i].lei_limit == LEDGER_LIMIT_INFINITY) {
374				printf("%10s %10s", "none", "-  ");
375			} else {
376				print_num(info[i].lei_limit, 0);
377				print_num(info[i].lei_refill_period, 0);
378			}
379			printf("\n");
380		}
381	}
382
383	if (line == 0)
384		exit (0);
385}
386
387static void
388cleanup()
389{
390	struct proc_list *p, *pnext, *plast;
391	struct ledger *l, *lnext, *llast;
392
393	plast = NULL;
394	for (p = procs; p; p = pnext) {
395		pnext = p->next;
396		if (p->seen == 0) {
397			if (plast)
398				plast->next = pnext;
399			else
400				procs = pnext;
401
402			free(p);
403		} else {
404			p->seen = 0;
405		}
406	}
407
408	llast = NULL;
409	for (l = ledgers; l; l = lnext) {
410		lnext = l->next;
411		if (l->seen == 0) {
412			if (llast)
413				llast->next = lnext;
414			else
415				ledgers = lnext;
416			free(l->info);
417			if (l->old_info)
418				free(l->old_info);
419			free(l);
420		} else {
421			l->seen = 0;
422			free(l->old_info);
423			l->old_info = l->info;
424			l->info = NULL;
425		}
426	}
427
428	free(template);
429	template = NULL;
430}
431
432const char *pname;
433
434static void
435usage()
436{
437	printf("%s [-hdL] [-g group] [-p pid] [-r resource] [interval]\n", pname);
438}
439
440int
441main(int argc, char **argv)
442{
443	int c;
444	int interval = 0;
445
446	pname = argv[0];
447
448	while ((c = getopt(argc, argv, "g:hdLp:r:")) != -1) {
449		switch (c) {
450		case 'g':
451			group_print = optarg;
452			break;
453
454		case 'd':
455			diff_mode = 1;
456			break;
457
458		case 'h':
459			usage();
460			exit(0);
461
462		case 'L':
463			get_template_info();
464			dump_template_info();
465			exit(0);
466
467		case 'p':
468			pid = atoi(optarg);
469			break;
470
471		case 'r':
472			resource_print = optarg;
473			break;
474
475		default:
476			usage();
477			exit(1);
478		}
479
480	}
481	argc -= optind;
482	argv += optind;
483
484	if (argc)
485		interval = atoi(argv[0]);
486
487	if (group_print && resource_print) {
488		fprintf(stderr, "Cannot specify both a resource and a group\n");
489		exit (1);
490	}
491
492	if (group_print)
493		validate_group();
494	if (resource_print)
495		validate_resource();
496
497	do {
498		get_template_info();
499		get_all_info();
500		dump_all_info();
501		cleanup();
502		sleep(interval);
503	} while (interval);
504}
505