minor optimizations
[util-vserver.git] / util-vserver / src / vps.c
1 // $Id$    --*- c -*--
2
3 // Copyright (C) 2004 Enrico Scholz <enrico.scholz@informatik.tu-chemnitz.de>
4 //  
5 // This program is free software; you can redistribute it and/or modify
6 // it under the terms of the GNU General Public License as published by
7 // the Free Software Foundation; version 2 of the License.
8 //  
9 // This program is distributed in the hope that it will be useful,
10 // but WITHOUT ANY WARRANTY; without even the implied warranty of
11 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 // GNU General Public License for more details.
13 //  
14 // You should have received a copy of the GNU General Public License
15 // along with this program; if not, write to the Free Software
16 // Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
17
18
19 #ifdef HAVE_CONFIG_H
20 #  include <config.h>
21 #endif
22
23 #include "util.h"
24 #include "pathconfig.h"
25
26 #include <lib/vserver.h>
27 #include <lib/fmt.h>
28 #include <assert.h>
29 #include <fcntl.h>
30 #include <signal.h>
31
32 #define ENSC_WRAPPERS_VSERVER   1
33 #define ENSC_WRAPPERS_STDLIB    1
34 #define ENSC_WRAPPERS_UNISTD    1
35 #define ENSC_WRAPPERS_FCNTL     1
36 #include <wrappers.h>
37
38 #define CTXNR_WIDTH     5
39 #define HUNK_SIZE       0x4000
40 #define CONTEXT_WIDTH   20
41 #define CONTEXT_PLACE   "                    "
42
43 int wrapper_exit_code = 254;
44
45 struct ContextMapping {
46     xid_t               ctx;
47     char const *        id;
48 };
49
50 static struct ContextMapping            *mapping    = 0;
51 static size_t                           mapping_len = 0;
52
53
54 static void
55 showHelp(int fd, char const *cmd, int res)
56 {
57   WRITE_MSG(fd, "Usage:  ");
58   WRITE_STR(fd, cmd);
59   WRITE_MSG(fd,
60             " <ps-opts>*\n\n"
61             "Please report bugs to " PACKAGE_BUGREPORT "\n");
62   exit(res);
63 }
64
65 static void
66 showVersion()
67 {
68   WRITE_MSG(1,
69             "vps " VERSION " -- shows processes in vserver-contexts\n"
70             "This program is part of " PACKAGE_STRING "\n\n"
71             "Copyright (C) 2004 Enrico Scholz\n"
72             VERSION_COPYRIGHT_DISCLAIMER);
73   exit(0);
74 }
75
76
77 static size_t
78 writeContextInfo(xid_t ctx, char const *name)
79 {
80   size_t        l1  = name==0 ? 0 : strlen(name);
81   size_t        res = CTXNR_WIDTH + 1;
82   
83   if (ctx==VC_NOCTX) {
84     if (3<CTXNR_WIDTH) write(1, CONTEXT_PLACE, CTXNR_WIDTH-3);
85     write(1, "ERR ", 4);
86   }
87   else {
88     char        buf[sizeof(ctx)*3+1];
89     size_t      l = utilvserver_fmt_ulong(buf, ctx);
90
91     if (l<CTXNR_WIDTH) write(1, CONTEXT_PLACE, CTXNR_WIDTH-l);
92     write(1, buf, l);
93     write(1, " ", 1);
94   }
95
96   if (l1!=0) {
97     assert(name!=0);
98     write(1, name, l1);
99   }
100
101   return res+l1;
102 }
103
104 static xid_t
105 extractCtx(char *pid_str)
106 {
107   pid_t         pid;
108   
109   while (*pid_str==' ') ++pid_str;
110   pid = atoi(pid_str);
111
112   return vc_get_task_xid(pid);
113 }
114
115 static char const *
116 resolveCtx(xid_t ctx)
117 {
118   char const *  res;
119   size_t        i;
120
121   for (i=0; i<mapping_len; ++i)
122     if (mapping[i].ctx==ctx) return mapping[i].id;
123
124   ++mapping_len;
125   mapping = Erealloc(mapping, mapping_len * sizeof(mapping[0]));
126   
127   if (ctx==0)      res = strdup("MAIN");
128   else if (ctx==1) res = strdup("ALL_PROC");
129   else {
130     vcCfgStyle  style = vcCFG_AUTO;
131     char        *tmp  = vc_getVserverByCtx(ctx, &style,0);
132     if (tmp) res = vc_getVserverName(tmp, style);
133     else     res = 0;
134     free(tmp);
135   }
136
137   mapping[mapping_len-1].ctx = ctx;
138   mapping[mapping_len-1].id  = res;
139   return res;
140 }
141
142 static char *
143 readOutput(int fd, size_t *total_len)
144 {
145   size_t        len  = 2*HUNK_SIZE;
146   char          *buf = Emalloc(len+1);
147   size_t        offset = 0;
148
149   for (;;) {
150     size_t      l;
151
152     while (offset >= len) {
153       len += HUNK_SIZE;
154       buf  = Erealloc(buf, len+1);
155     }
156     
157     l = Eread(fd, buf+offset, len - offset);
158     if (l==0) break;
159
160     offset += l;
161   }
162
163   buf[offset] = '\0';
164
165   if (total_len)
166     *total_len = offset;
167
168   return buf;
169 }
170
171 static void
172 processOutput(char *data, size_t len)
173 {
174   size_t        pid_end;
175   char *        eol_pos = strchr(data, '\n');
176   char *        pos;
177
178   if (eol_pos==0) eol_pos  = data + len;
179   else            *eol_pos = '\0';
180   
181   pos = strstr(data, "PID");
182   if (pos==0) {
183     WRITE_MSG(2, "Failed to parse ps-output\n");
184     exit(wrapper_exit_code);
185   }
186
187   pid_end = pos-data + 4;
188
189   write(1, data, pid_end);
190   write(1, "CONTEXT" CONTEXT_PLACE, CONTEXT_WIDTH);
191   write(1, data+pid_end, eol_pos-(data+pid_end));
192   write(1, "\n", 1);
193
194   len -= eol_pos-data;
195   data = eol_pos+1;
196   
197   while (len > 1) {
198     char const  *vserver_name = 0;
199     xid_t       ctx;
200     size_t      l;
201
202     --len;
203     eol_pos = strchr(data, '\n');
204     
205     if (eol_pos==0) eol_pos  = data + len;
206
207     ctx          = extractCtx(data + pid_end - 6);
208     vserver_name = resolveCtx(ctx);
209
210     write(1, data, pid_end);
211     l = writeContextInfo(ctx, vserver_name);
212     if (l<CONTEXT_WIDTH) write(1, CONTEXT_PLACE, CONTEXT_WIDTH-l);
213     else                 write(1, " ", 1);
214     write(1, data+pid_end, eol_pos-(data+pid_end));
215     write(1, "\n", 1);
216     
217     len -= eol_pos-data;
218     data = eol_pos+1;
219   }
220 }
221
222 int main(int argc, char *argv[])
223 {
224   int           p[2];
225   pid_t         pid;
226   char *        data;
227   size_t        len;
228   char const *  errptr;
229
230   if (argc>1) {
231     if (strcmp(argv[1], "--help")   ==0) showHelp(1, argv[0], 0);
232     if (strcmp(argv[1], "--version")==0) showVersion();
233   }
234
235   signal(SIGCHLD, SIG_DFL);
236
237   if (!switchToWatchXid(&errptr)) {
238     perror(errptr);
239     exit(wrapper_exit_code);
240   }
241
242   Epipe(p);
243   pid = Efork();
244
245   if (pid==0) {
246     int         fd = Eopen("/dev/null", O_RDONLY, 0);
247     Edup2(fd,   0);
248     Edup2(p[1], 1);
249     Eclose(p[0]);
250     Eclose(p[1]);
251     Eclose(fd);
252
253     argv[0] = "ps";
254
255     Eexecv(PS_PROG, argv);
256   }
257
258   Eclose(p[1]);
259   data = readOutput(p[0], &len);
260   Eclose(p[0]);
261   
262   processOutput(data, len);
263   exitLikeProcess(pid, "ps", wrapper_exit_code);
264 }