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_out(char* buff, size_t length, FILE *fd) {
 194         if (fwrite(buff, length, 1, fd) < 1) {
 195                 perror("error writing buffer to output file");
 196                 fclose(fd);
 197                 exit(1);
 198         }
 199 }
 200
 201 static void copy_tmpfile(FILE *tmpfile, FILE *outfile, char *buf, size_t size) {
 202         if (fseek(tmpfile, 0, SEEK_SET)) {
 203                 perror("could to seek to start of temporary file");
 204                 fclose(tmpfile);
 205                 exit(1);
 206         }
 207         while (fread(buf, size, 1, tmpfile) > 0) {
 208                 write_buff_out(buf, size, outfile);
 209         }
 210         if (ferror(tmpfile)) {
 211                 perror("read temporary file");
 212                 fclose(tmpfile);
 213                 exit(1);
 214         }
 215         fclose(tmpfile);
 216         fclose(outfile);
 217 }
 218
 219 FILE *open_tmpfile(void) {
 220         struct cs_status cs;
 221         int tmpfd;
 222         FILE *tmpfile;
 223         mode_t mask;
 224         char *tmpdir;
 225         char const * const template="%s/sponge.XXXXXX";
 226
 227         trapsignals();
 228         cs = cs_enter();
 229         tmpdir = getenv("TMPDIR");
 230         if (tmpdir == NULL)
 231                 tmpdir = "/tmp";
 232         /* Subtract 2 for `%s' and add 1 for the trailing NULL. */
 233         tmpname=malloc(strlen(tmpdir) + strlen(template) - 2 + 1);
 234         if (! tmpname) {
 235                 perror("failed to allocate memory");
 236                 exit(1);
 237         }
 238         sprintf(tmpname, template, tmpdir);
 239         mask=umask(077);
 240         tmpfd = mkstemp(tmpname);
 241         umask(mask);
 242         atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
 243         cs_leave(cs);
 244
 245         if (tmpfd < 0) {
 246                 perror("mkstemp failed");
 247                 exit(1);
 248         }
 249         tmpfile = fdopen(tmpfd, "w+");
 250         if (! tmpfile) {
 251                 perror("fdopen");
 252                 exit(1);
 253         }
 254         return tmpfile;
 255 }
 256
 257 int main (int argc, char **argv) {
 258         char *buf, *bufstart, *outname = NULL;
 259         size_t bufsize = BUFF_SIZE;
 260         size_t bufused = 0;
 261         FILE *outfile, *tmpfile = 0;
 262         ssize_t i = 0;
 263         size_t mem_available = default_sponge_size();
 264
 265         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 266                 usage();
 267         }
 268         if (argc == 2) {
 269                 outname = argv[1];
 270         }
 271
 272         bufstart = buf = malloc(bufsize);
 273         if (!buf) {
 274                 perror("failed to allocate memory");
 275                 exit(1);
 276         }
 277         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 278                 bufused = bufused+i;
 279                 if (bufused == bufsize) {
 280                         if ((bufsize*2) >= mem_available) {
 281                                 if (!tmpfile) {
 282                                         tmpfile=open_tmpfile();
 283                                 }
 284                                 write_buff_tmp(bufstart, bufused, tmpfile);
 285                                 bufused = 0;
 286                         }
 287                         else {
 288                                 bufsize *= 2;
 289                                 bufstart = realloc(bufstart, bufsize);
 290                                 if (!bufstart) {
 291                                         perror("failed to realloc memory");
 292                                         exit(1);
 293                                 }
 294                         }
 295                 }
 296                 buf = bufstart + bufused;
 297         }
 298         if (i < 0) {
 299                 perror("failed to read from stdin");
 300                 exit(1);
 301         }
 302         if (tmpfile) {
 303                 struct stat statbuf;
 304
 305                 /* write whatever we have in memory to tmpfile */
 306                 if (bufused)
 307                         write_buff_tmp(bufstart, bufused, tmpfile);
 308                 if (fflush(tmpfile) != 0) {
 309                         perror("fflush");
 310                         exit(1);
 311                 }
 312
 313                 if (outname) {
 314                         /* If it's a regular file, or does not yet exist,
 315                          * attempt a fast rename of the temp file. */
 316                         if (((lstat(outname, &statbuf) == 0 &&
 317                               S_ISREG(statbuf.st_mode) &&
 318                               ! S_ISLNK(statbuf.st_mode)
 319                              ) || errno == ENOENT) &&
 320                             rename(tmpname, outname) == 0) {
 321                                 /* Fix renamed file mode to match either
 322                                  * the old file mode, or the default file
 323                                  * mode for a newly created file. */
 324                                 mode_t mode;
 325                                 if (errno != ENOENT) {
 326                                         mode = statbuf.st_mode;
 327                                 }
 328                                 else {
 329                                         mode_t mask = umask(0);
 330                                         umask(mask);
 331                                         mode = 0666 & ~mask;
 332                                 }
 333                                 if (chmod(outname, mode) != 0) {
 334                                         perror("chmod");
 335                                         exit(1);
 336                                 }
 337                                 return(0);
 338                         }
 339
 340                         /* Fall back to slow copy. */
 341                         outfile = fopen(outname, "w");
 342                         if (!outfile) {
 343                                 perror("error opening output file");
 344                                 exit(1);
 345                         }
 346                         copy_tmpfile(tmpfile, outfile, bufstart, bufsize);
 347                 }
 348                 else {
 349                         copy_tmpfile(tmpfile, stdout, bufstart, bufsize);
 350                 }
 351         }
 352         else {
 353                 if (outname) {
 354                         outfile = fopen(outname, "w");
 355                         if (!outfile) {
 356                                 perror("error opening output file");
 357                                 exit(1);
 358                         }
 359                 }
 360                 else {
 361                         outfile = stdout;
 362                 }
 363                 if (bufused)
 364                         write_buff_out(bufstart, bufused, outfile);
 365                 fclose(outfile);
 366         }
 367
 368         return 0;
 369 }