sponge.c

   1 /*
   2  *  sponge.c - read in all available info from stdin, then output it to
   3  *  file named on the command line
   4  *
   5  *  Copyright ©  2006  Tollef Fog Heen
   6  *
   7  *  This program is free software; you can redistribute it and/or
   8  *  modify it under the terms of the GNU General Public License
   9  *  version 2 as published by the Free Software Foundation.
  10  *
  11  *  This program is distributed in the hope that it will be useful,
  12  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  *  General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, write to the Free Software
  18  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
  19  * USA
  20  *
  21  */
  22
  23 #include <stdio.h>
  24 #include <stdlib.h>
  25 #include <unistd.h>
  26 #include <sys/types.h>
  27 #include <sys/stat.h>
  28 /* MAX() */
  29 #include <sys/param.h>
  30 #include <fcntl.h>
  31 #include <errno.h>
  32 #include <string.h>
  33 #include <sys/resource.h>
  34 /* SIZE_MAX */
  35 #include <stdint.h>
  36 #include <signal.h>
  37
  38 #include "physmem.c"
  39
  40 #define BUFF_SIZE           8192
  41 #define MIN_SPONGE_SIZE     BUFF_SIZE
  42 char *tmpname = NULL;
  43
  44 void usage() {
  45         printf("sponge <file>: soak up all input from stdin and write it to <file>\n");
  46         exit(0);
  47 }
  48
  49 /* all the signal stuff copied from gnu sort */
  50
  51 /* The set of signals that are caught.  */
  52 static sigset_t caught_signals;
  53
  54 /* Critical section status.  */
  55 struct cs_status {
  56         int valid; // was bool
  57         sigset_t sigs;
  58 };
  59
  60 /* Enter a critical section.  */
  61 static struct cs_status cs_enter (void) {
  62         struct cs_status status;
  63         status.valid = (sigprocmask(SIG_BLOCK, &caught_signals, &status.sigs) == 0);
  64         return status;
  65 }
  66
  67 /* Leave a critical section.  */
  68 static void cs_leave (struct cs_status status) {
  69         if (status.valid) {
  70                 /* Ignore failure when restoring the signal mask. */
  71                 sigprocmask(SIG_SETMASK, &status.sigs, NULL);
  72         }
  73 }
  74
  75 static void cleanup () {
  76         if (tmpname) {
  77                 unlink(tmpname);
  78         }
  79 }
  80
  81 static void onexit_cleanup (void) {
  82         struct cs_status cs = cs_enter();
  83         cleanup();
  84         cs_leave(cs);
  85 }
  86
  87 static void sighandler (int sig) {
  88         if (! SA_NOCLDSTOP)
  89                 signal(sig, SIG_IGN);
  90
  91         cleanup();
  92
  93         signal(sig, SIG_DFL);
  94         raise(sig);
  95 }
  96
  97 /* taken from coreutils sort */
  98 static size_t default_sponge_size (void) {
  99         /* Let MEM be available memory or 1/8 of total memory, whichever
 100            is greater.  */
 101         double avail = physmem_available();
 102         double total = physmem_total();
 103         double mem = MAX(avail, total / 8);
 104         struct rlimit rlimit;
 105
 106         /* Let SIZE be MEM, but no more than the maximum object size or
 107            system resource limits.  Avoid the MIN macro here, as it is not
 108            quite right when only one argument is floating point.  Don't
 109            bother to check for values like RLIM_INFINITY since in practice
 110            they are not much less than SIZE_MAX.  */
 111         size_t size = SIZE_MAX;
 112         if (mem < size)
 113                 size = mem;
 114         if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size)
 115                 size = rlimit.rlim_cur;
 116 #ifdef RLIMIT_AS
 117         if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size)
 118                 size = rlimit.rlim_cur;
 119 #endif
 120
 121         /* Leave a large safety margin for the above limits, as failure can
 122            occur when they are exceeded.  */
 123         size /= 2;
 124
 125 #ifdef RLIMIT_RSS
 126         /* Leave a 1/16 margin for RSS to leave room for code, stack, etc.
 127            Exceeding RSS is not fatal, but can be quite slow.  */
 128         if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size)
 129                 size = rlimit.rlim_cur / 16 * 15;
 130 #endif
 131
 132         /* Use no less than the minimum. */
 133         return MAX (size, MIN_SPONGE_SIZE);
 134 }
 135
 136 void trapsignals (void) {
 137         ssize_t i = 0;
 138         static int const sig[] = {
 139                 /* The usual suspects.  */
 140                 SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM,
 141 #ifdef SIGPOLL
 142                 SIGPOLL,
 143 #endif
 144 #ifdef SIGPROF
 145                 SIGPROF,
 146 #endif
 147 #ifdef SIGVTALRM
 148                 SIGVTALRM,
 149 #endif
 150 #ifdef SIGXCPU
 151                 SIGXCPU,
 152 #endif
 153 #ifdef SIGXFSZ
 154                 SIGXFSZ,
 155 #endif
 156         };
 157         int nsigs = sizeof(sig) / sizeof(sig[0]);
 158
 159 #if SA_NOCLDSTOP
 160         struct sigaction act;
 161
 162         sigemptyset(&caught_signals);
 163         for (i = 0; i < nsigs; i++) {
 164                 sigaction(sig[i], NULL, &act);
 165                 if (act.sa_handler != SIG_IGN)
 166                         sigaddset(&caught_signals, sig[i]);
 167                 }
 168
 169                 act.sa_handler = sighandler;
 170                 act.sa_mask = caught_signals;
 171                 act.sa_flags = 0;
 172
 173                 for (i = 0; i < nsigs; i++)
 174                         if (sigismember(&caught_signals, sig[i]))
 175                                 sigaction(sig[i], &act, NULL);
 176 #else
 177         for (i = 0; i < nsigs; i++)
 178                 if (signal(sig[i], SIG_IGN) != SIG_IGN) {
 179                         signal(sig[i], sighandler);
 180                         siginterrupt (sig[i], 1);
 181                 }
 182 #endif
 183 }
 184
 185 static void write_buff_tmp(char* buff, size_t length, FILE *fd) {
 186         if (fwrite(buff, length, 1, fd) < 1) {
 187                 perror("error writing buffer to temporary file");
 188                 fclose(fd);
 189                 exit(1);
 190         }
 191 }
 192
 193 static void write_buff_tmp_finish (char* buff, size_t length, FILE *fd) {
 194         if (length)
 195                 write_buff_tmp(buff, length, fd);
 196         if (fflush(fd) != 0) {
 197                 perror("fflush");
 198                 exit(1);
 199         }
 200 }
 201
 202 static void write_buff_out (char* buff, size_t length, FILE *fd) {
 203         if (fwrite(buff, length, 1, fd) < 1) {
 204                 perror("error writing buffer to output file");
 205                 fclose(fd);
 206                 exit(1);
 207         }
 208 }
 209
 210 static void copy_tmpfile (FILE *tmpfile, FILE *outfile, char *buf, size_t size) {
 211         ssize_t i;
 212         if (lseek(fileno(tmpfile), 0, SEEK_SET)) {
 213                 perror("could to seek to start of temporary file");
 214                 fclose(tmpfile);
 215                 exit(1);
 216         }
 217         while ((i = read(fileno(tmpfile), buf, size)) > 0) {
 218                 write_buff_out(buf, i, outfile);
 219         }
 220         if (i == -1) {
 221                 perror("read temporary file");
 222                 fclose(tmpfile);
 223                 exit(1);
 224         }
 225         fclose(tmpfile);
 226         fclose(outfile);
 227 }
 228
 229 FILE *open_tmpfile (void) {
 230         struct cs_status cs;
 231         int tmpfd;
 232         FILE *tmpfile;
 233         mode_t mask;
 234         char *tmpdir;
 235         char const * const template="%s/sponge.XXXXXX";
 236
 237         trapsignals();
 238         cs = cs_enter();
 239         tmpdir = getenv("TMPDIR");
 240         if (tmpdir == NULL)
 241                 tmpdir = "/tmp";
 242         /* Subtract 2 for `%s' and add 1 for the trailing NULL. */
 243         tmpname=malloc(strlen(tmpdir) + strlen(template) - 2 + 1);
 244         if (! tmpname) {
 245                 perror("failed to allocate memory");
 246                 exit(1);
 247         }
 248         sprintf(tmpname, template, tmpdir);
 249         mask=umask(077);
 250         tmpfd = mkstemp(tmpname);
 251         umask(mask);
 252         atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
 253         cs_leave(cs);
 254
 255         if (tmpfd < 0) {
 256                 perror("mkstemp failed");
 257                 exit(1);
 258         }
 259         tmpfile = fdopen(tmpfd, "w+");
 260         if (! tmpfile) {
 261                 perror("fdopen");
 262                 exit(1);
 263         }
 264         return tmpfile;
 265 }
 266
 267 int main (int argc, char **argv) {
 268         char *buf, *bufstart, *outname = NULL;
 269         size_t bufsize = BUFF_SIZE;
 270         size_t bufused = 0;
 271         FILE *outfile, *tmpfile = 0;
 272         ssize_t i = 0;
 273         size_t mem_available = default_sponge_size();
 274         int tmpfile_used=0;
 275
 276         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 277                 usage();
 278         }
 279         if (argc == 2) {
 280                 outname = argv[1];
 281         }
 282
 283         tmpfile = open_tmpfile();
 284         bufstart = buf = malloc(bufsize);
 285         if (!buf) {
 286                 perror("failed to allocate memory");
 287                 exit(1);
 288         }
 289         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 290                 bufused = bufused+i;
 291                 if (bufused == bufsize) {
 292                         if ((bufsize*2) >= mem_available) {
 293                                 write_buff_tmp(bufstart, bufused, tmpfile);
 294                                 bufused = 0;
 295                                 tmpfile_used = 1;
 296                         }
 297                         else {
 298                                 bufsize *= 2;
 299                                 bufstart = realloc(bufstart, bufsize);
 300                                 if (!bufstart) {
 301                                         perror("failed to realloc memory");
 302                                         exit(1);
 303                                 }
 304                         }
 305                 }
 306                 buf = bufstart + bufused;
 307         }
 308         if (i < 0) {
 309                 perror("failed to read from stdin");
 310                 exit(1);
 311         }
 312
 313         if (outname) {
 314                 mode_t mode;
 315                 struct stat statbuf;
 316                 int exists = (lstat(outname, &statbuf) == 0);
 317
 318                 write_buff_tmp_finish(bufstart, bufused, tmpfile);
 319
 320                 /* Set temp file mode to match either
 321                  * the old file mode, or the default file
 322                  * mode for a newly created file. */
 323                 if (exists) {
 324                         mode = statbuf.st_mode;
 325                 }
 326                 else {
 327                         mode_t mask = umask(0);
 328                         umask(mask);
 329                         mode = 0666 & ~mask;
 330                 }
 331                 if (chmod(tmpname, mode) != 0) {
 332                         perror("chmod");
 333                         exit(1);
 334                 }
 335
 336                 /* If it's a regular file, or does not yet exist,
 337                  * attempt a fast rename of the temp file. */
 338                 if (((exists &&
 339                       S_ISREG(statbuf.st_mode) &&
 340                       ! S_ISLNK(statbuf.st_mode)
 341                      ) || ! exists) &&
 342                     rename(tmpname, outname) == 0) {
 343                         tmpname=NULL; /* don't try to cleanup tmpname */
 344                 }
 345                 else {
 346                         /* Fall back to slow copy. */
 347                         outfile = fopen(outname, "w");
 348                         if (!outfile) {
 349                                 perror("error opening output file");
 350                                 exit(1);
 351                         }
 352                         copy_tmpfile(tmpfile, outfile, bufstart, bufsize);
 353                 }
 354         }
 355         else {
 356                 if (tmpfile_used) {
 357                         write_buff_tmp_finish(bufstart, bufused, tmpfile);
 358                         copy_tmpfile(tmpfile, stdout, bufstart, bufsize);
 359                 }
 360                 else if (bufused) {
 361                         /* buffer direct to stdout, no tmpfile */
 362                         write_buff_out(bufstart, bufused, stdout);
 363                 }
 364         }
 365
 366         return 0;
 367 }