diff --git a/src/fstrm_capture.c b/src/fstrm_capture.c index dbeb720b8..810becfd4 100644 --- a/src/fstrm_capture.c +++ b/src/fstrm_capture.c @@ -28,6 +28,7 @@ #include #include #include +#include #include #include @@ -39,6 +40,7 @@ #include "libmy/argv.h" #include "libmy/my_alloc.h" #include "libmy/print_string.h" +#include "libmy/ubuf.h" #if HAVE_DECL_FFLUSH_UNLOCKED # define fflush fflush_unlocked @@ -53,6 +55,7 @@ #endif #define CAPTURE_HIGH_WATERMARK 262144 +#define FNAME_MAXLEN 1024 struct capture; struct capture_args; @@ -98,6 +101,8 @@ struct capture { struct event *ev_sighup; FILE *output_file; + char *output_filename; + time_t output_open_ts; size_t bytes_written; size_t count_written; @@ -109,6 +114,8 @@ struct capture_args { char *str_content_type; char *str_read_unix; char *str_write_fname; + int split_s; + char *str_suffix; }; static struct capture g_program_ctx; @@ -143,7 +150,19 @@ static argv_t g_args[] = { ARGV_CHAR_P, &g_program_args.str_write_fname, "", - "file path to write Frame Streams data to" }, + "filename to write, as a strftime(3) string" }, + + { 'x', "suffix", + ARGV_CHAR_P, + &g_program_args.str_suffix, + "", + "output filename suffix to add after closing" }, + + { 's', "split", + ARGV_INT, + &g_program_args.split_s, + "", + "Amount of seconds to capture into each file" }, { ARGV_LAST }, }; @@ -241,6 +260,10 @@ parse_args(const int argc, char **argv) { argv_version_string = PACKAGE_VERSION; + /* These are non-mandatory options */ + g_program_args.split_s = 0; + g_program_args.str_suffix = NULL; + if (argv_process(g_args, argc, argv) != 0) return false; @@ -253,7 +276,15 @@ parse_args(const int argc, char **argv) usage("Unix socket path to read from (--unix) is not set"); if (g_program_args.str_write_fname == NULL) usage("file path to write Frame Streams data to (--write) is not set"); - + if (strcmp(g_program_args.str_write_fname, "-") == 0) { + int fd = fileno(stdout); + if (isatty(fd) == 1) + usage("refusing binary output to terminal (stdout)"); + if (g_program_args.split_s != 0) + usage("cannot use output splitting with stdout"); + if (g_program_args.str_suffix != NULL) + usage("cannot use output suffix renaming with stdout"); + } return true; } @@ -373,30 +404,61 @@ open_write_start(struct capture *ctx) static bool open_write_file(struct capture *ctx) { - const char *fname = ctx->args->str_write_fname; + char *fstr = ctx->args->str_write_fname; + char *fname, *fn_buf; + time_t t; + struct tm the_time; + size_t fn_len; + int fd; mode_t open_mode = S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH; int open_flags = O_CREAT | O_WRONLY | O_TRUNC; #if defined(O_CLOEXEC) open_flags |= O_CLOEXEC; #endif - /* Open the file descriptor. */ - int fd = open(fname, open_flags, open_mode); - if (fd == -1) { - fprintf(stderr, "%s: failed to open output file %s\n", - argv_program, fname); - return false; + if (strcmp(fstr, "-") == 0) { + fd = fileno(stdout); + ctx->output_file = stdout; + fname = fstr; } + else { + /* Filename strftime expansion, useful when splitting output */ + t = time(NULL); + localtime_r(&t, &the_time); + fn_buf = (char*)my_calloc(FNAME_MAXLEN, 1); + fn_len = strftime(fn_buf, FNAME_MAXLEN, fstr, &the_time); + if (!fn_len) { + fprintf(stderr, "%s: strftime failed on filename \"%s\"\n", + argv_program, fstr); + return false; + } + fname = my_malloc(fn_len); + strcpy(fname, fn_buf); + my_free(fn_buf); + + /* Open the file descriptor. */ + fd = open(fname, open_flags, open_mode); + if (fd == -1) { + fprintf(stderr, "%s: failed to open output file %s\n", + argv_program, fname); + return false; + } - /* Open the FILE object. */ - ctx->output_file = fdopen(fd, "w"); - if (!ctx->output_file) { - close(fd); - fprintf(stderr, "%s: failed to fdopen output file %s\n", - argv_program, fname); - return false; + /* Open the FILE object. */ + ctx->output_file = fdopen(fd, "w"); + if (!ctx->output_file) { + close(fd); + fprintf(stderr, "%s: failed to fdopen output file %s\n", + argv_program, fname); + return false; + } } + ctx->count_written = 0; + ctx->bytes_written = 0; + ctx->output_open_ts = t; + ctx->output_filename = fname; + /* Write the START frame. */ if (!open_write_start(ctx)) { fclose(ctx->output_file); @@ -407,13 +469,16 @@ open_write_file(struct capture *ctx) } /* Success. */ - fprintf(stderr, "%s: opened output file %s\n", argv_program, fname); + fprintf(stderr, "%s: opened output file %s (fd=%d)\n", argv_program, fname, fd); return true; } static bool close_write_file(struct capture *ctx) { + char *fname = ctx->output_filename; + const char *suffix = ctx->args->str_suffix; + if (ctx->output_file != NULL) { /* Write the STOP frame. */ if (!close_write_stop(ctx)) @@ -424,10 +489,26 @@ close_write_file(struct capture *ctx) ctx->output_file = NULL; } + if (fname != NULL) { + /* Rename newly closed file if suffix is desired */ + if (suffix != NULL) { + ubuf *fn_buf = ubuf_dup_cstr(fname); + ubuf_add_cstr(fn_buf, suffix); + (void)rename(fname, (char*)fn_buf->_v); + ubuf_destroy(&fn_buf); + } + } + /* Success. */ - fprintf(stderr, "%s: closed output file %s (wrote %zd frames, %zd bytes)\n", - argv_program, ctx->args->str_write_fname, + fprintf(stderr, "%s: closed output file %s (suffix %s, wrote %zd frames, %zd bytes)\n", + argv_program, fname, suffix ? suffix : "(none)", ctx->count_written, ctx->bytes_written); + + if (fname != NULL) { + my_free(fname); + ctx->output_filename = NULL; + } + return true; } @@ -482,6 +563,21 @@ process_data_frame(struct conn *conn) conn->bytes_read += bytes_read; conn->ctx->count_written += 1; conn->ctx->bytes_written += bytes_read; + + /* Output file rotation requested? */ + if (conn->ctx->args->split_s > 0) { + time_t t_now = time(NULL); + /* Is it time to rotate? */ + if (t_now >= conn->ctx->output_open_ts + conn->ctx->args->split_s) { + /* Rotate output file, fail hard if unsuccessful */ + if (!close_write_file(conn->ctx)) { + exit(EXIT_FAILURE); + } + if (!open_write_file(conn->ctx)) { + exit(EXIT_FAILURE); + } + } + } } static bool