1/* stdbuf -- setup the standard streams for a command
2   Copyright (C) 2009-2010 Free Software Foundation, Inc.
3
4   This program is free software: you can redistribute it and/or modify
5   it under the terms of the GNU General Public License as published by
6   the Free Software Foundation, either version 3 of the License, or
7   (at your option) any later version.
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, see <http://www.gnu.org/licenses/>.  */
16
17/* Written by Pádraig Brady.  */
18
19#include <config.h>
20#include <stdio.h>
21#include <getopt.h>
22#include <sys/types.h>
23#include <assert.h>
24
25#include "system.h"
26#include "error.h"
27#include "filenamecat.h"
28#include "posixver.h"
29#include "quote.h"
30#include "xreadlink.h"
31#include "xstrtol.h"
32#include "c-ctype.h"
33
34/* The official name of this program (e.g., no `g' prefix).  */
35#define PROGRAM_NAME "stdbuf"
36#define LIB_NAME "libstdbuf.so" /* FIXME: don't hardcode  */
37
38#define AUTHORS proper_name_utf8 ("Padraig Brady", "P\303\241draig Brady")
39
40static char *program_path;
41
42static struct
43{
44  size_t size;
45  int optc;
46  char *optarg;
47} stdbuf[3];
48
49static struct option const longopts[] =
50{
51  {"input", required_argument, NULL, 'i'},
52  {"output", required_argument, NULL, 'o'},
53  {"error", required_argument, NULL, 'e'},
54  {GETOPT_HELP_OPTION_DECL},
55  {GETOPT_VERSION_OPTION_DECL},
56  {NULL, 0, NULL, 0}
57};
58
59/* Set size to the value of STR, interpreted as a decimal integer,
60   optionally multiplied by various values.
61   Return -1 on error, 0 on success.
62
63   This supports dd BLOCK size suffixes.
64   Note we don't support dd's b=512, c=1, w=2 or 21x512MiB formats.  */
65static int
66parse_size (char const *str, size_t *size)
67{
68  uintmax_t tmp_size;
69  enum strtol_error e = xstrtoumax (str, NULL, 10, &tmp_size, "EGkKMPTYZ0");
70  if (e == LONGINT_OK && tmp_size > SIZE_MAX)
71    e = LONGINT_OVERFLOW;
72
73  if (e == LONGINT_OK)
74    {
75      errno = 0;
76      *size = tmp_size;
77      return 0;
78    }
79
80  errno = (e == LONGINT_OVERFLOW ? EOVERFLOW : 0);
81  return -1;
82}
83
84void
85usage (int status)
86{
87  if (status != EXIT_SUCCESS)
88    fprintf (stderr, _("Try `%s --help' for more information.\n"),
89             program_name);
90  else
91    {
92      printf (_("Usage: %s OPTION... COMMAND\n"), program_name);
93      fputs (_("\
94Run COMMAND, with modified buffering operations for its standard streams.\n\
95\n\
96"), stdout);
97      fputs (_("\
98Mandatory arguments to long options are mandatory for short options too.\n\
99"), stdout);
100      fputs (_("\
101  -i, --input=MODE   Adjust standard input stream buffering\n\
102  -o, --output=MODE  Adjust standard output stream buffering\n\
103  -e, --error=MODE   Adjust standard error stream buffering\n\
104"), stdout);
105      fputs (HELP_OPTION_DESCRIPTION, stdout);
106      fputs (VERSION_OPTION_DESCRIPTION, stdout);
107      fputs (_("\n\
108If MODE is `L' the corresponding stream will be line buffered.\n\
109This option is invalid with standard input.\n"), stdout);
110      fputs (_("\n\
111If MODE is `0' the corresponding stream will be unbuffered.\n\
112"), stdout);
113      fputs (_("\n\
114Otherwise MODE is a number which may be followed by one of the following:\n\
115KB 1000, K 1024, MB 1000*1000, M 1024*1024, and so on for G, T, P, E, Z, Y.\n\
116In this case the corresponding stream will be fully buffered with the buffer\n\
117size set to MODE bytes.\n\
118"), stdout);
119      fputs (_("\n\
120NOTE: If COMMAND adjusts the buffering of its standard streams (`tee' does\n\
121for e.g.) then that will override corresponding settings changed by `stdbuf'.\n\
122Also some filters (like `dd' and `cat' etc.) don't use streams for I/O,\n\
123and are thus unaffected by `stdbuf' settings.\n\
124"), stdout);
125      emit_ancillary_info ();
126    }
127  exit (status);
128}
129
130/* argv[0] can be anything really, but generally it contains
131   the path to the executable or just a name if it was executed
132   using $PATH. In the latter case to get the path we can:
133   search getenv("PATH"), readlink("/prof/self/exe"), getenv("_"),
134   dladdr(), pstat_getpathname(), etc.  */
135
136static void
137set_program_path (const char *arg)
138{
139  if (strchr (arg, '/'))        /* Use absolute or relative paths directly.  */
140    {
141      program_path = dir_name (arg);
142    }
143  else
144    {
145      char *path = xreadlink ("/proc/self/exe");
146      if (path)
147        program_path = dir_name (path);
148      else if ((path = getenv ("PATH")))
149        {
150          char *dir;
151          path = xstrdup (path);
152          for (dir = strtok (path, ":"); dir != NULL; dir = strtok (NULL, ":"))
153            {
154              char *candidate = file_name_concat (dir, arg, NULL);
155              if (access (candidate, X_OK) == 0)
156                {
157                  program_path = dir_name (candidate);
158                  free (candidate);
159                  break;
160                }
161              free (candidate);
162            }
163        }
164      free (path);
165    }
166}
167
168static int
169optc_to_fileno (int c)
170{
171  int ret = -1;
172
173  switch (c)
174    {
175    case 'e':
176      ret = STDERR_FILENO;
177      break;
178    case 'i':
179      ret = STDIN_FILENO;
180      break;
181    case 'o':
182      ret = STDOUT_FILENO;
183      break;
184    }
185
186  return ret;
187}
188
189static void
190set_LD_PRELOAD (void)
191{
192  int ret;
193  char *old_libs = getenv ("LD_PRELOAD");
194  char *LD_PRELOAD;
195
196  /* Note this would auto add the appropriate search path for "libstdbuf.so":
197     gcc stdbuf.c -Wl,-rpath,'$ORIGIN' -Wl,-rpath,$PKGLIBDIR
198     However we want the lookup done for the exec'd command not stdbuf.
199
200     Since we don't link against libstdbuf.so add it to LIBDIR rather than
201     LIBEXECDIR, as we'll search for it in the "sys default" case below.  */
202  char const *const search_path[] = {
203    program_path,
204    PKGLIBDIR,
205    "",                         /* sys default */
206    NULL
207  };
208
209  char const *const *path = search_path;
210  char *libstdbuf;
211
212  do
213    {
214      struct stat sb;
215
216      if (!**path)              /* system default  */
217        {
218          libstdbuf = xstrdup (LIB_NAME);
219          break;
220        }
221      ret = asprintf (&libstdbuf, "%s/%s", *path, LIB_NAME);
222      if (ret < 0)
223        xalloc_die ();
224      if (stat (libstdbuf, &sb) == 0)   /* file_exists  */
225        break;
226      free (libstdbuf);
227    }
228  while (*++path);
229
230  /* FIXME: Do we need to support libstdbuf.dll, c:, '\' separators etc?  */
231
232  if (old_libs)
233    ret = asprintf (&LD_PRELOAD, "LD_PRELOAD=%s:%s", old_libs, libstdbuf);
234  else
235    ret = asprintf (&LD_PRELOAD, "LD_PRELOAD=%s", libstdbuf);
236
237  if (ret < 0)
238    xalloc_die ();
239
240  free (libstdbuf);
241
242  ret = putenv (LD_PRELOAD);
243
244  if (ret != 0)
245    {
246      error (EXIT_CANCELED, errno,
247             _("failed to update the environment with %s"),
248             quote (LD_PRELOAD));
249    }
250}
251
252/* Populate environ with _STDBUF_I=$MODE _STDBUF_O=$MODE _STDBUF_E=$MODE  */
253
254static void
255set_libstdbuf_options (void)
256{
257  int i;
258
259  for (i = 0; i < ARRAY_CARDINALITY (stdbuf); i++)
260    {
261      if (stdbuf[i].optarg)
262        {
263          char *var;
264          int ret;
265
266          if (*stdbuf[i].optarg == 'L')
267            ret = asprintf (&var, "%s%c=L", "_STDBUF_",
268                            toupper (stdbuf[i].optc));
269          else
270            ret = asprintf (&var, "%s%c=%" PRIuMAX, "_STDBUF_",
271                            toupper (stdbuf[i].optc),
272                            (uintmax_t) stdbuf[i].size);
273          if (ret < 0)
274            xalloc_die ();
275
276          if (putenv (var) != 0)
277            {
278              error (EXIT_CANCELED, errno,
279                     _("failed to update the environment with %s"),
280                     quote (var));
281            }
282        }
283    }
284}
285
286int
287main (int argc, char **argv)
288{
289  int c;
290
291  initialize_main (&argc, &argv);
292  set_program_name (argv[0]);
293  setlocale (LC_ALL, "");
294  bindtextdomain (PACKAGE, LOCALEDIR);
295  textdomain (PACKAGE);
296
297  initialize_exit_failure (EXIT_CANCELED);
298  atexit (close_stdout);
299
300  while ((c = getopt_long (argc, argv, "+i:o:e:", longopts, NULL)) != -1)
301    {
302      int opt_fileno;
303
304      switch (c)
305        {
306        /* Old McDonald had a farm ei...  */
307        case 'e':
308        case 'i':
309        case 'o':
310          opt_fileno = optc_to_fileno (c);
311          assert (0 <= opt_fileno && opt_fileno < ARRAY_CARDINALITY (stdbuf));
312          stdbuf[opt_fileno].optc = c;
313          while (c_isspace (*optarg))
314            optarg++;
315          stdbuf[opt_fileno].optarg = optarg;
316          if (c == 'i' && *optarg == 'L')
317            {
318              /* -oL will be by far the most common use of this utility,
319                 but one could easily think -iL might have the same affect,
320                 so disallow it as it could be confusing.  */
321              error (0, 0, _("line buffering stdin is meaningless"));
322              usage (EXIT_CANCELED);
323            }
324
325          if (!STREQ (optarg, "L")
326              && parse_size (optarg, &stdbuf[opt_fileno].size) == -1)
327            error (EXIT_CANCELED, errno, _("invalid mode %s"), quote (optarg));
328
329          break;
330
331        case_GETOPT_HELP_CHAR;
332
333        case_GETOPT_VERSION_CHAR (PROGRAM_NAME, AUTHORS);
334
335        default:
336          usage (EXIT_CANCELED);
337        }
338    }
339
340  argv += optind;
341  argc -= optind;
342
343  /* must specify at least 1 command.  */
344  if (argc < 1)
345    {
346      error (0, 0, _("missing operand"));
347      usage (EXIT_CANCELED);
348    }
349
350  /* FIXME: Should we mandate at least one option?  */
351
352  set_libstdbuf_options ();
353
354  /* Try to preload libstdbuf first from the same path as
355     stdbuf is running from.  */
356  set_program_path (argv[0]);
357  if (!program_path)
358    program_path = xstrdup (PKGLIBDIR);  /* Need to init to non NULL.  */
359  set_LD_PRELOAD ();
360  free (program_path);
361
362  execvp (*argv, argv);
363
364  {
365    int exit_status = (errno == ENOENT ? EXIT_ENOENT : EXIT_CANNOT_INVOKE);
366    error (0, errno, _("failed to run command %s"), quote (argv[0]));
367    exit (exit_status);
368  }
369}
370