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