script: add the -o/--output-limit option. Fix race test.

When script is used on a host with a relatively small free disk space, it
is sometimes desirable to limit the size of the captured output. This
can now be enforced with the --output-limit option.

The --output-limit option lets the user specify a maximum size. The program
uses the size parsing from strutils and thus supports the usual
multiplicative suffixes (kiB, KB, MiB, MB, etc.). After the specified
number of bytes have been written to the output file, the script program
will terminate the child process.

Due to buffering, the size of the output file might exceed the specified
limit. This limit also does not include the start and done messages.

The race test was throwing an error dur to a variable being "" in some cases.
Quoting the variable in the equal test took care of that test.

[kzak@redhat.com: - use done() to stop script
                  - count also timing file
                  - remove unnamed member initialization in ctl struct
                  - add to bash-completion]

Signed-off-by: Fred Mora <fmora@datto.com>
Signed-off-by: Karel Zak <kzak@redhat.com>

2 files changed, 40 insertions, 6 deletions

diff --git a/term-utils/script.1 b/term-utils/script.1
index 6c50bc621..041b7620b 100644
--- a/term-utils/script.1
+++ b/term-utils/script.1
@@ -54,6 +54,10 @@ saves the dialogue in this
 If no filename is given, the dialogue is saved in the file
 .BR typescript .
 .SH OPTIONS
+Below, the \fIsize\fR argument may be followed by the multiplicative
+suffixes KiB (=1024), MiB (=1024*1024), and so on for GiB, TiB, PiB, EiB, ZiB and YiB
+(the "iB" is optional, e.g. "K" has the same meaning as "KiB"), or the suffixes
+KB (=1000), MB (=1000*1000), and so on for GB, TB, PB, EB, ZB and YB.
 .TP
 \fB\-a\fR, \fB\-\-append\fR
 Append the output to
@@ -82,6 +86,15 @@ being done using `cat foo'.
 Allow the default output destination, i.e. the typescript file, to be a hard
 or symbolic link.  The command will follow a symbolic link.
 .TP
+\fB\-o\fR, \fB\-\-output-limit\fR \fIsize\fR
+Limit the size of the typescript and timing files to
+.I size
+and stop the child process after this size is exceeded.  The calculated
+file size does not include the start and done messages that the
+.B script
+command prepends and appends to the child process output.
+Due to buffering, the resulting output file might be larger than the specified value.
+.TP
 \fB\-q\fR, \fB\-\-quiet\fR
 Be quiet (do not write start and done messages to standard output).
 .TP
diff --git a/term-utils/script.c b/term-utils/script.c
index d5ffa27f1..102f7ec32 100644
--- a/term-utils/script.c
+++ b/term-utils/script.c
@@ -71,6 +71,7 @@
 #include "all-io.h"
 #include "monotonic.h"
 #include "timeutils.h"
+#include "strutils.h"
 
 #include "debug.h"
 
@@ -104,6 +105,8 @@ struct script_control {
 	FILE *typescriptfp;	/* output file pointer */
 	char *tname;		/* timing file path */
 	FILE *timingfp;		/* timing file pointer */
+	ssize_t outsz;          /* current output file size */
+	ssize_t maxsz;		/* maximum output file size */
 	struct timeval oldtime;	/* previous write or command start time */
 	int master;		/* pseudoterminal master file descriptor */
 	int slave;		/* pseudoterminal slave file descriptor */
@@ -169,6 +172,7 @@ static void __attribute__((__noreturn__)) usage(void)
 		" -e, --return                  return exit code of the child process\n"
 		" -f, --flush                   run flush after each write\n"
 		"     --force                   use output file even when it is a link\n"
+		" -o, --output-limit <size>     terminate if output files exceed size\n"
 		" -q, --quiet                   be quiet\n"
 		" -t[<file>], --timing[=<file>] output timing data to stderr or to FILE\n"
 		), out);
@@ -266,6 +270,8 @@ static void wait_for_child(struct script_control *ctl, int wait)
 static void write_output(struct script_control *ctl, char *obuf,
 			    ssize_t bytes)
 {
+	int timing_bytes = 0;
+
 	DBG(IO, ul_debug(" writing output"));
 
 	if (ctl->timing && ctl->timingfp) {
@@ -275,11 +281,13 @@ static void write_output(struct script_control *ctl, char *obuf,
 
 		gettime_monotonic(&now);
 		timersub(&now, &ctl->oldtime, &delta);
-		fprintf(ctl->timingfp, "%ld.%06ld %zd\n",
+		timing_bytes = fprintf(ctl->timingfp, "%ld.%06ld %zd\n",
 		        (long)delta.tv_sec, (long)delta.tv_usec, bytes);
 		if (ctl->flush)
 			fflush(ctl->timingfp);
 		ctl->oldtime = now;
+		if (timing_bytes < 0)
+			timing_bytes = 0;
 	}
 
 	DBG(IO, ul_debug("  writing to script file"));
@@ -299,6 +307,9 @@ static void write_output(struct script_control *ctl, char *obuf,
 		fail(ctl);
 	}
 
+	if (ctl->maxsz != 0)
+		ctl->outsz += bytes + timing_bytes;
+
 	DBG(IO, ul_debug("  writing output *done*"));
 }
 
@@ -350,7 +361,6 @@ static void handle_io(struct script_control *ctl, int fd, int *eof)
 {
 	char buf[BUFSIZ];
 	ssize_t bytes;
-
 	DBG(IO, ul_debug("%d FD active", fd));
 	*eof = 0;
 
@@ -379,10 +389,18 @@ static void handle_io(struct script_control *ctl, int fd, int *eof)
 		 * shell output that looks like double echoing */
 		fdatasync(ctl->master);
 
-	/* from command (master) to stdout */
+	/* from command (master) to stdout and log */
 	} else if (fd == ctl->master) {
 		DBG(IO, ul_debug(" master --> stdout %zd bytes", bytes));
 		write_output(ctl, buf, bytes);
+
+		/* check output limit */
+		if (ctl->maxsz != 0 && ctl->outsz >= ctl->maxsz) {
+			if (!ctl->quiet)
+				printf(_("Script terminated, max output file size %zd exceeded.\n"), ctl->maxsz);
+			DBG(IO, ul_debug("output size %zd, exceeded limit %zd", ctl->outsz, ctl->maxsz));
+			done(ctl);
+		}
 	}
 }
 
@@ -688,8 +706,7 @@ int main(int argc, char **argv)
 		.line = "/dev/ptyXX",
 #endif
 		.master = -1,
-		.poll_timeout = -1,
-		0
+		.poll_timeout = -1
 	};
 	int ch;
 
@@ -701,6 +718,7 @@ int main(int argc, char **argv)
 		{"return", no_argument, NULL, 'e'},
 		{"flush", no_argument, NULL, 'f'},
 		{"force", no_argument, NULL, FORCE_OPTION,},
+		{"output-limit", required_argument, NULL, 'o'},
 		{"quiet", no_argument, NULL, 'q'},
 		{"timing", optional_argument, NULL, 't'},
 		{"version", no_argument, NULL, 'V'},
@@ -723,7 +741,7 @@ int main(int argc, char **argv)
 
 	script_init_debug();
 
-	while ((ch = getopt_long(argc, argv, "ac:efqt::Vh", longopts, NULL)) != -1)
+	while ((ch = getopt_long(argc, argv, "ac:efo:qt::Vh", longopts, NULL)) != -1)
 		switch (ch) {
 		case 'a':
 			ctl.append = 1;
@@ -740,6 +758,9 @@ int main(int argc, char **argv)
 		case FORCE_OPTION:
 			ctl.force = 1;
 			break;
+		case 'o':
+			ctl.maxsz = strtosize_or_err(optarg, _("failed to parse output limit size"));
+			break;
 		case 'q':
 			ctl.quiet = 1;
 			break;