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