xref: /plan9-contrib/sys/src/cmd/ptrace/tron.c (revision af198995f3f1f1d87a386835b2bcd15ca90a6fed)
1 #include <u.h>
2 #include <libc.h>
3 
4 static void
usage(void)5 usage(void)
6 {
7 	fprint(2, "usage: %s [-s nkevs] [-f dev] [-o file] [-p pid] cmd [args...]\n", argv0);
8 	exits("usage");
9 }
10 
11 char buf[1*1024*1024];
12 int done;
13 int pids[128];
14 int npids;
15 
16 void
main(int argc,char * argv[])17 main(int argc, char *argv[])
18 {
19 	int fd, cfd, pid, ofd, i;
20 	char file[128], *f;
21 	char *dev, *devctl, *out;
22 	long tot, nr;
23 	int sz;
24 
25 	sz = 8;
26 	dev = "/dev/ptrace";
27 	out = nil;
28 	ARGBEGIN{
29 	case 'f':
30 		dev = EARGF(usage());
31 		break;
32 	case 'o':
33 		out = EARGF(usage());
34 		break;
35 	case 's':
36 		sz = strtoul(EARGF(usage()), 0, 0);
37 		break;
38 	case 'p':
39 		if(npids > nelem(pids))
40 			sysfatal("too many pids");
41 		pids[npids++] = strtoul(EARGF(usage()), 0, 0);
42 		break;
43 	default:
44 		usage();
45 	}ARGEND
46 	if(argc == 0)
47 		usage();
48 
49 	sz *= 1024;
50 	if(sz < 1024)
51 		sz = 1024;
52 	if(sz > 64*1024)	/* perhaps we shouldn't limit it */
53 		sz = 64*1024;
54 
55 	devctl = smprint("%sctl", dev);
56 	if(devctl == nil)
57 		sysfatal("no memory");
58 
59 	if(access(argv[0], AEXIST) < 0){
60 		seprint(file, file+sizeof file, "/bin/%s", argv[0]);
61 		f = file;
62 	}else
63 		f = argv[0];
64 	if(access(f, AEXIST) < 0)
65 		sysfatal("%s: %r", argv[0]);
66 
67 
68 	/* prepage; don't want interference while tracing */
69 	for(i = 0; i < sizeof buf; i += 4096)
70 		buf[i] = 0;
71 
72 	rfork(RFREND);
73 	pid = 0;
74 
75 	/* play a dance with the traced command so it starts when we want,
76 	 * or issue the trace command for the traced pids.
77 	 */
78 	cfd = open(devctl, OWRITE);
79 	if(cfd < 0)
80 		sysfatal("can't trace: %r");
81 	if(npids == 0){
82 		pid = rfork(RFPROC|RFFDG);
83 		switch(pid){
84 		case -1:
85 			sysfatal("fork: %r");
86 		case 0:
87 			/* wait until traced */
88 			if(rendezvous(usage, 0) == 0)
89 				exits("no trace");
90 			rfork(RFREND);
91 			exec(f, argv);
92 			sysfatal("exec: %r");
93 		}
94 		if(fprint(cfd, "size %d\n", sz) < 0 || fprint(cfd, "trace %d 1\n", pid) < 0){
95 			rendezvous(usage, 0);
96 			sysfatal("can't trace");
97 		}
98 
99 		if(out == nil){
100 			rendezvous(usage, usage);
101 			waitpid();
102 			exits(nil);
103 		}
104 	}else{
105 		for(i = 0; i < npids; i++)
106 			if(fprint(cfd, "trace %d 1\n", pids[i]) < 0)
107 				fprint(2, "%s: trace %d: %r\n", argv0, pids[i]);
108 	}
109 	close(cfd);
110 
111 	/* ready: open the trace and collect what we can. */
112 	fd = open(dev, OREAD);
113 	ofd = create(out, OWRITE, 0664);
114 	if(fd < 0 || ofd < 0){
115 		if(npids == 0)
116 			rendezvous(usage, 0);
117 		sysfatal("can't trace: %r");
118 	}
119 
120 	done = 0;
121 	switch(rfork(RFPROC|RFMEM)){
122 	case -1:
123 		sysfatal("fork");
124 	default:
125 		if(npids == 0){
126 			/* let the traced go! */
127 			rendezvous(usage, usage);
128 		}
129 
130 		while(waitpid() != pid)
131 			;
132 		done = 1;
133 		exits(nil);
134 	case 0:
135 		/* The damn device won't block us. We must insist upon eof.
136 		 * Besides, when done, we must give time for the events to
137 		 * arrive.
138 		 */
139 		for(;;){
140 			for(tot = 0; tot < sizeof buf - 128; tot += nr){
141 				nr = read(fd, buf+tot, sizeof buf - tot);
142 				if(nr < 0)
143 					sysfatal("read: %r");
144 				if(nr == 0){
145 					sleep(100);
146 					if(done)
147 						break;
148 				}
149 			}
150 			if(tot > 0 && write(ofd, buf, tot) != tot)
151 				sysfatal("write: %r");
152 			/* try to read more if we were done; but just once. */
153 			if(done && tot == 0 && done++ == 2)
154 				break;
155 		}
156 		close(fd);
157 		close(ofd);
158 		rendezvous(main, nil);
159 		exits(nil);
160 	}
161 }
162