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