From 9f69edabd9138aa2ca6c359a57b09f490e344664 Mon Sep 17 00:00:00 2001 From: Adrian Muzyka Date: Thu, 25 Apr 2024 11:32:36 +0200 Subject: [PATCH] Support zstd compression and option to set dump dir and compression per PID --- Makefile | 5 + README.md | 3 +- memcr-client.c | 171 ++++++++++++++++++++++---- memcr.c | 322 +++++++++++++++++++++++++++++++++++++++++-------- memcr.h | 23 +++- 5 files changed, 451 insertions(+), 73 deletions(-) diff --git a/Makefile b/Makefile index 2dcf394..4361c36 100644 --- a/Makefile +++ b/Makefile @@ -38,6 +38,11 @@ ifeq ($(COMPRESS_LZ4), 1) LDFLAGS += -llz4 endif +ifeq ($(COMPRESS_ZSTD), 1) + MCFLAGS += -DCOMPRESS_ZSTD + LDFLAGS += -lzstd +endif + ifeq ($(CHECKSUM_MD5), 1) MCFLAGS += -DCHECKSUM_MD5 LDFLAGS += -lcrypto diff --git a/README.md b/README.md index dd6f326..c70d95a 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ make ##### compilation options You can enable support for compression and checksumming of memory dump file: - `COMPRESS_LZ4=1` - requires liblz4 + - `COMPRESS_ZSTD=1` - requires libzstd - `CHECKSUM_MD5=1` - requires libcrypto and openssl headers There is also `ENCRYPT` option for building `libencrypt.so` that provides sample implementation of encryption layer based on libcrypto API. memcr is not linked with libencrypt.so, but it can be preloaded with `LD_PRELOAD`. @@ -65,7 +66,7 @@ options: -n --no-wait no wait for key press -m --proc-mem get pages from /proc/pid/mem -f --rss-file include file mapped memory - -z --compress compress memory dump + -z --compress compress memory dump with alg: lz4 or zstd -c --checksum enable md5 checksum for memory dump -e --encrypt enable encryption of memory dump ``` diff --git a/memcr-client.c b/memcr-client.c index c1f73d5..4cb1867 100644 --- a/memcr-client.c +++ b/memcr-client.c @@ -24,9 +24,13 @@ #include #include #include +#include #include "memcr.h" +#define CMD_LEN_MAX (sizeof(struct service_command) + (2*sizeof(memcr_svc_checkpoint_options)) \ + + MEMCR_DUMPDIR_LEN_MAX + sizeof(memcr_compress_alg)) + static int xconnect(struct sockaddr *addr, socklen_t addrlen) { int cd, ret; @@ -74,8 +78,61 @@ static int send_cmd(int cd, struct service_command cmd) int ret; struct service_response resp = {0}; - ret = write(cd, &cmd, sizeof(struct service_command)); - if (ret != sizeof(struct service_command)) { + ret = write(cd, &cmd, sizeof(cmd)); + if (ret != sizeof(cmd)) { + fprintf(stderr, "%s() write request failed: ret %d, errno %m\n", __func__, ret); + return -1; + } + + ret = read(cd, &resp, sizeof(struct service_response)); + if (ret != sizeof(struct service_response)) { + fprintf(stderr, "%s() read response failed: ret %d, errno %m\n", __func__, ret); + return -1; + } + + fprintf(stdout, "Procedure finished with %s status.\n", MEMCR_OK == resp.resp_code ? "OK" : "ERROR"); + + return resp.resp_code; +} + +static int send_cmd_v2(int cd, memcr_svc_cmd cmd, pid_t pid, const struct service_options *opt) +{ + int ret; + struct service_response resp = {0}; + + size_t cmd_size = 0; + unsigned char cmd_buf[CMD_LEN_MAX]; + + struct service_command cmd_cr = {.cmd = cmd, .pid = pid}; + memcpy(cmd_buf, &cmd_cr, sizeof(struct service_command)); + cmd_size += sizeof(struct service_command); + + if (opt && opt->is_dump_dir) { + memcr_svc_checkpoint_options opt_id = MEMCR_CHECKPOINT_DUMPDIR; + memcpy(cmd_buf+cmd_size, &opt_id, sizeof(memcr_svc_checkpoint_options)); + cmd_size += sizeof(memcr_svc_checkpoint_options); + strncpy((char*)cmd_buf + cmd_size, opt->dump_dir, MEMCR_DUMPDIR_LEN_MAX); + cmd_size += strlen(opt->dump_dir) + 1; + } + + if (opt && opt->is_compress_alg) { + memcr_svc_checkpoint_options opt_id = MEMCR_CHECKPOINT_COMPRESS_ALG; + memcpy(cmd_buf+cmd_size, &opt_id, sizeof(memcr_svc_checkpoint_options)); + cmd_size += sizeof(memcr_svc_checkpoint_options); + memcpy(cmd_buf+cmd_size, &opt->compress_alg, sizeof(memcr_compress_alg)); + cmd_size += sizeof(memcr_compress_alg); + } + + struct service_command cmd_v2 = {.cmd = MEMCR_CMDS_V2, .pid = cmd_size}; + + ret = write(cd, &cmd_v2, sizeof(cmd_v2)); + if (ret != sizeof(cmd_v2)) { + fprintf(stderr, "%s() write request failed: ret %d, errno %m\n", __func__, ret); + return -1; + } + + ret = write(cd, cmd_buf, cmd_size); + if (ret != cmd_size) { fprintf(stderr, "%s() write request failed: ret %d, errno %m\n", __func__, ret); return -1; } @@ -94,15 +151,18 @@ static int send_cmd(int cd, struct service_command cmd) static void usage(const char *name, int status) { fprintf(status ? stderr : stdout, - "%s -l PORT|PATH -p PID [-c -r]\n" \ + "%s -l PORT|PATH -p PID [-c [-d DIR] [-z ALG] -r]\n" \ "options: \n" \ " -h --help\t\thelp\n" \ " -l --location\t\tTCP port number of localhost memcr service\n" \ "\t\t\t or filesystem path to memcr service UNIX socket\n" \ " -p --pid\t\tprocess ID to be checkpointed / restored\n" \ " -c --checkpoint\tsend checkpoint command to memcr service\n" \ - " -r --restore\t\tsend restore command to memcr service\n", - name); + " -d --dir\tdir where memory dump is stored of max length %d\n" \ + " -z --compress\tcompress memory dump with selected algorithm: 'lz4', 'zstd' or disable with 'none'\n" \ + " -r --restore\t\tsend restore command to memcr service\n" \ + " -v --v1\t\tforce using old protocol\n", + name, MEMCR_DUMPDIR_LEN_MAX); exit(status); } @@ -113,20 +173,26 @@ int main(int argc, char *argv[]) int restore = 0; int port = -1; int option_index; - struct service_command cmd = {0}; char *comm_location = NULL; int pid = 0; + const char *dump_dir_s = 0; + const char *compress_alg_s = 0; + struct service_options checkpoint_options = {0}; + int v1 = 0; struct option long_options[] = { { "help", 0, 0, 'h'}, { "location", 1, 0, 'l'}, { "pid", 1, 0, 'p'}, { "checkpoint", 0, 0, 'c'}, + { "dir", 1, 0, 'd'}, + { "compress", 1, 0, 'z'}, { "restore", 0, 0, 'r'}, + { "v1", 0, 0, 'v'}, { NULL, 0, 0, 0 } }; - while ((opt = getopt_long(argc, argv, "hl:p:cr", long_options, &option_index)) != -1) { + while ((opt = getopt_long(argc, argv, "hl:p:cd:z:rv", long_options, &option_index)) != -1) { switch (opt) { case 'h': usage(argv[0], 0); @@ -140,9 +206,18 @@ int main(int argc, char *argv[]) case 'c': checkpoint = 1; break; + case 'd': + dump_dir_s = optarg; + break; + case 'z': + compress_alg_s = optarg; + break; case 'r': restore = 1; break; + case 'v': + v1 = 1; + break; default: /* '?' */ usage(argv[0], 1); break; @@ -161,35 +236,85 @@ int main(int argc, char *argv[]) return -1; } - port = atoi(comm_location); + if (!checkpoint && (dump_dir_s || compress_alg_s)) { + fprintf(stderr, "Dir dump and compression is available only for checkpoint!\n"); + usage(argv[0], 1); + return -1; + } - if (port > 0) - cd = connect_tcp(port); - else - cd = connect_unix(comm_location); + if (dump_dir_s) { + if (strlen(dump_dir_s) >= MEMCR_DUMPDIR_LEN_MAX) { + fprintf(stderr, "Dir dump too long!\n"); + usage(argv[0], 1); + return -1; + } - if (cd < 0) { - fprintf(stderr, "Connection creation failed!\n"); - return cd; + strcpy(checkpoint_options.dump_dir, dump_dir_s); + checkpoint_options.is_dump_dir = 1; } + if (compress_alg_s) { + checkpoint_options.is_compress_alg = 1; + if (strcmp(compress_alg_s, "none") == 0) + checkpoint_options.compress_alg = MEMCR_COMPRESS_NONE; + else if (strcmp(compress_alg_s, "lz4") == 0) + checkpoint_options.compress_alg = MEMCR_COMPRESS_LZ4; + else if (strcmp(compress_alg_s, "zstd") == 0) + checkpoint_options.compress_alg = MEMCR_COMPRESS_ZSTD; + else { + fprintf(stderr, "Incorrect compression algorithm provided!\n"); + usage(argv[0], 1); + return -1; + } + } + + port = atoi(comm_location); + if (checkpoint) { fprintf(stdout, "Will checkpoint %d.\n", pid); - cmd.cmd = MEMCR_CHECKPOINT; - cmd.pid = pid; - ret = send_cmd(cd, cmd); + + if (port > 0) + cd = connect_tcp(port); + else + cd = connect_unix(comm_location); + + if (cd < 0) { + fprintf(stderr, "Connection creation failed!\n"); + return cd; + } + + if (v1) { + struct service_command cmd = {.cmd = MEMCR_CHECKPOINT, pid = pid}; + ret = send_cmd(cd, cmd); + } else + ret = send_cmd_v2(cd, MEMCR_CHECKPOINT, pid, &checkpoint_options); + + close(cd); } if (restore) { fprintf(stdout, "Will restore %d.\n", pid); - cmd.cmd = MEMCR_RESTORE; - cmd.pid = pid; - ret = send_cmd(cd, cmd); + + if (port > 0) + cd = connect_tcp(port); + else + cd = connect_unix(comm_location); + + if (cd < 0) { + fprintf(stderr, "Connection creation failed!\n"); + return cd; + } + + if (v1) { + struct service_command cmd = {.cmd = MEMCR_RESTORE, pid = pid}; + ret = send_cmd(cd, cmd); + } else + ret = send_cmd_v2(cd, MEMCR_RESTORE, pid, NULL); + + close(cd); } fprintf(stdout, "Command executed, exiting.\n"); - close(cd); - return ret; } diff --git a/memcr.c b/memcr.c index 6009b03..f2a1619 100644 --- a/memcr.c +++ b/memcr.c @@ -51,6 +51,10 @@ #include #endif +#ifdef COMPRESS_ZSTD +#include +#endif + #ifdef CHECKSUM_MD5 #include #if OPENSSL_VERSION_NUMBER >= 0x30000000L @@ -110,13 +114,13 @@ struct vm_area { unsigned long flags; }; -static char *dump_dir; +static char *dfl_dump_dir; static char *parasite_socket_dir; static int parasite_socket_use_netns; static int no_wait; static int proc_mem; static int rss_file; -static int compress; +static memcr_compress_alg dfl_compress_alg; static int checksum; static int service; @@ -150,6 +154,13 @@ static int nr_vmas; #define MAX_LZ4_DST_SIZE LZ4_compressBound(MAX_VM_REGION_SIZE) #endif +#ifdef COMPRESS_ZSTD +#define MAX_ZSTD_DST_SIZE ZSTD_compressBound(MAX_VM_REGION_SIZE) +#ifndef COMPRESS_ZSTD_LVL +#define COMPRESS_ZSTD_LVL 3 +#endif +#endif + static pid_t parasite_pid; static pid_t parasite_pid_clone; static struct target_context ctx; @@ -197,6 +208,7 @@ static struct { int state; int checkpoint_abort; int checkpoint_cmd_sd; + struct service_options options; } checkpoint_service_data[CHECKPOINTED_PIDS_LIMIT]; #define SOCKET_INVALID (-1) @@ -347,6 +359,17 @@ static void md5_final(unsigned char *md, unsigned int *len, void *ctx) } #endif +static void die(const char *fmt, ...) +{ + va_list ap; + + va_start(ap, fmt); + vfprintf(stderr, fmt, ap); + va_end(ap); + + exit(1); +} + static void parasite_status_signal(pid_t pid, int status) { pthread_mutex_lock(¶site_watch.lock); @@ -421,7 +444,7 @@ static void parasite_socket_init(struct sockaddr_un *addr, pid_t pid) } } -static void cleanup_pid(pid_t pid) +static void cleanup_pid(pid_t pid, const char* dump_dir) { char path[PATH_MAX]; @@ -811,7 +834,15 @@ static int dump_write(int fd, const void *buf, size_t count) return ret; } -static void init_pid_checkpoint_data(pid_t pid) +static void clear_checkpoint_options(struct service_options *options) +{ + options->is_dump_dir = FALSE; + options->dump_dir[0] = 0; + options->is_compress_alg = FALSE; + options->compress_alg = MEMCR_COMPRESS_NONE; +} + +static void init_pid_checkpoint_data(pid_t pid, struct service_options *options) { pthread_mutex_lock(&checkpoint_service_data_lock); for (int i=0; iis_dump_dir; + strncpy(checkpoint_service_data[i].options.dump_dir, options->dump_dir, + MEMCR_DUMPDIR_LEN_MAX); + checkpoint_service_data[i].options.is_compress_alg = options->is_compress_alg; + checkpoint_service_data[i].options.compress_alg = options->compress_alg; + } else { + clear_checkpoint_options(&checkpoint_service_data[i].options); + } pthread_mutex_unlock(&checkpoint_service_data_lock); return; } @@ -836,11 +876,14 @@ static void cleanup_checkpointed_pids(void) if (checkpoint_service_data[i].pid != PID_INVALID) { fprintf(stdout, "[i] Killing PID %d\n", checkpoint_service_data[i].pid); kill(checkpoint_service_data[i].pid, SIGKILL); - cleanup_pid(checkpoint_service_data[i].pid); + const char *dir = checkpoint_service_data[i].options.is_dump_dir ? + checkpoint_service_data[i].options.dump_dir : dfl_dump_dir; + cleanup_pid(checkpoint_service_data[i].pid, dir); checkpoint_service_data[i].pid = PID_INVALID; checkpoint_service_data[i].worker = PID_INVALID; checkpoint_service_data[i].state = STATE_RESTORED; checkpoint_service_data[i].checkpoint_cmd_sd = SOCKET_INVALID; + clear_checkpoint_options(&checkpoint_service_data[i].options); } } pthread_mutex_unlock(&checkpoint_service_data_lock); @@ -903,6 +946,7 @@ static void clear_pid_checkpoint_data(pid_t pid) checkpoint_service_data[i].worker = PID_INVALID; checkpoint_service_data[i].checkpoint_cmd_sd = SOCKET_INVALID; checkpoint_service_data[i].state = STATE_RESTORED; + clear_checkpoint_options(&checkpoint_service_data[i].options); } } pthread_mutex_unlock(&checkpoint_service_data_lock); @@ -914,11 +958,14 @@ static void clear_pid_on_worker_exit_non_blocking(pid_t worker) if (checkpoint_service_data[i].worker == worker) { fprintf(stdout, "[+] Clearing pid: %d with worker: %d ...\n", checkpoint_service_data[i].pid, worker); - cleanup_pid(checkpoint_service_data[i].pid); + const char *dir = checkpoint_service_data[i].options.is_dump_dir ? + checkpoint_service_data[i].options.dump_dir : dfl_dump_dir; + cleanup_pid(checkpoint_service_data[i].pid, dir); checkpoint_service_data[i].pid = PID_INVALID; checkpoint_service_data[i].worker = PID_INVALID; checkpoint_service_data[i].checkpoint_cmd_sd = SOCKET_INVALID; checkpoint_service_data[i].state = STATE_RESTORED; + clear_checkpoint_options(&checkpoint_service_data[i].options); } } } @@ -1008,10 +1055,27 @@ static int read_vm_region(int fd, struct vm_region *vmr, char *buf) if (!vm_region_valid(vmr)) return -1; -#ifdef COMPRESS_LZ4 - if (compress) { + + if (dfl_compress_alg != MEMCR_COMPRESS_NONE) { int src_size; - char src[MAX_LZ4_DST_SIZE]; + int src_size_max = 0; + + switch (dfl_compress_alg) { +#ifdef COMPRESS_LZ4 + case MEMCR_COMPRESS_LZ4: + src_size_max = MAX_LZ4_DST_SIZE; + break; +#endif +#ifdef COMPRESS_ZSTD + case MEMCR_COMPRESS_ZSTD: + src_size_max = MAX_ZSTD_DST_SIZE; + break; +#endif + default: + die("compression set but not enabled, recompile with COMPRESS_LZ4=1 and/or COMPRESS_ZSTD=1"); + } + + char src[src_size_max]; ret = dump_read(fd, &src_size, sizeof(src_size)); if (ret != sizeof(src_size)) @@ -1021,13 +1085,24 @@ static int read_vm_region(int fd, struct vm_region *vmr, char *buf) if (ret != src_size) return -1; - ret = LZ4_decompress_safe(src, buf, src_size, MAX_VM_REGION_SIZE); + switch (dfl_compress_alg) { +#ifdef COMPRESS_LZ4 + case MEMCR_COMPRESS_LZ4: + ret = LZ4_decompress_safe(src, buf, src_size, MAX_VM_REGION_SIZE); + break; +#endif +#ifdef COMPRESS_ZSTD + case MEMCR_COMPRESS_ZSTD: + ret = ZSTD_decompress(buf, MAX_VM_REGION_SIZE, src, src_size); + break; +#endif + default: + die("compression set but not enabled, recompile with COMPRESS_LZ4=1 and/or COMPRESS_ZSTD=1"); + } /* fprintf(stdout, "[+] Decompressed %d Bytes back into %d.\n", srcSize, ret); */ if (ret <= 0) return -1; - } else -#endif - { + } else { ret = dump_read(fd, buf, vmr->len); if (ret != vmr->len) return -1; @@ -1047,12 +1122,39 @@ static int write_vm_region(int fd, const struct vm_region *vmr, const void *buf) if (ret != sizeof(struct vm_region)) return -1; + if (dfl_compress_alg != MEMCR_COMPRESS_NONE) { + int dst_size_max = 0; + switch (dfl_compress_alg) { #ifdef COMPRESS_LZ4 - if (compress) { - char dst[MAX_LZ4_DST_SIZE]; - int dst_size; + case MEMCR_COMPRESS_LZ4: + dst_size_max = MAX_LZ4_DST_SIZE; + break; +#endif +#ifdef COMPRESS_ZSTD + case MEMCR_COMPRESS_ZSTD: + dst_size_max = MAX_ZSTD_DST_SIZE; + break; +#endif + default: + die("compression set but not enabled, recompile with COMPRESS_LZ4=1 and/or COMPRESS_ZSTD=1"); + } + char dst[dst_size_max]; + int dst_size = 0; - dst_size = LZ4_compress_default(buf, dst, vmr->len, MAX_LZ4_DST_SIZE); + switch (dfl_compress_alg) { +#ifdef COMPRESS_LZ4 + case MEMCR_COMPRESS_LZ4: + dst_size = LZ4_compress_default(buf, dst, vmr->len, MAX_LZ4_DST_SIZE); + break; +#endif +#ifdef COMPRESS_ZSTD + case MEMCR_COMPRESS_ZSTD: + dst_size = ZSTD_compress(dst, MAX_ZSTD_DST_SIZE, buf, vmr->len, COMPRESS_ZSTD_LVL); + break; +#endif + default: + die("compression set but not enabled, recompile with COMPRESS_LZ4=1 and/or COMPRESS_ZSTD=1"); + } /* fprintf(stdout, "[+] Compressed %lu Bytes into %d.\n", len, dstSize); */ if (dst_size <= 0) return -1; @@ -1065,9 +1167,7 @@ static int write_vm_region(int fd, const struct vm_region *vmr, const void *buf) if (ret != dst_size) return -1; - } else -#endif - { + } else { ret = dump_write(fd, buf, vmr->len); if (ret != vmr->len) return -1; @@ -1540,7 +1640,7 @@ static int get_target_pages(int pid, struct vm_area vmas[], int nr_vmas) if (pd < 0) goto out; - snprintf(path, sizeof(path), "%s/pages-%d.img", dump_dir, pid); + snprintf(path, sizeof(path), "%s/pages-%d.img", dfl_dump_dir, pid); fd = dump_open(path, O_CREAT | O_TRUNC | O_WRONLY, S_IRUSR | S_IWUSR); if (fd < 0) { @@ -1659,7 +1759,7 @@ static int target_set_pages(pid_t pid) int cd = -1; int fd = -1; - snprintf(path, sizeof(path), "%s/pages-%d.img", dump_dir, pid); + snprintf(path, sizeof(path), "%s/pages-%d.img", dfl_dump_dir, pid); fd = dump_open(path, O_RDONLY, 0); if (fd < 0) { @@ -1780,7 +1880,7 @@ static int cmd_checkpoint(pid_t pid) } fprintf(stdout, "[i] download took %lu ms\n", diff_ms(&ts)); - fprintf(stdout, "[i] stored at %s/pages-%d.img\n", dump_dir, pid); + fprintf(stdout, "[i] stored at %s/pages-%d.img\n", dfl_dump_dir, pid); get_target_rss(pid, &vms_b); @@ -2270,6 +2370,112 @@ static int read_command(int cd, struct service_command *svc_cmd) return ret; } +static int read_command_v2(int cd, struct service_command *svc_cmd, struct service_options *options, size_t len) +{ + /* There must be at least service_command that can be followed by service_checkpoint_options */ + if (len < sizeof(struct service_command)) { + fprintf(stderr, "[-] %s(): cmds len to short: %d\n", __func__, (unsigned int)len); + return -1; + } + + int ret = read_command(cd, svc_cmd); + if (ret < 0) { + fprintf(stderr, "%s(): Error reading a command!\n", __func__); + return ret; + } + + len -= sizeof(struct service_command); + + switch (svc_cmd->cmd) { + case MEMCR_CHECKPOINT: { + fprintf(stdout, "[+] read MEMCR_CHECKPOINT for %d\n", svc_cmd->pid); + /* try to read checkpoint options */ + memcr_svc_checkpoint_options option; + while (len >= sizeof(option) && _read(cd, &option, sizeof(option)) > 0 ) { + len -= sizeof(option); + + switch (option) { + case MEMCR_CHECKPOINT_DUMPDIR: { + /* read string till NULL */ + unsigned int pos = 0; + while (len-- > 0 && (_read(cd, &options->dump_dir[pos], sizeof(char)) == sizeof(char)) && + (options->dump_dir[pos] != 0) && (++pos < MEMCR_DUMPDIR_LEN_MAX)); + + if (pos >= MEMCR_DUMPDIR_LEN_MAX || options->dump_dir[pos] != 0) + { + fprintf(stderr, "[-] %s(): dump dir path too long or not terminated with NULL\n", __func__); + return -1; + } + + options->is_dump_dir = TRUE; + fprintf(stdout, "[+] read dump dir path for this checkpoint: %s\n", options->dump_dir); + break; + } + case MEMCR_CHECKPOINT_COMPRESS_ALG: + { + if (len < sizeof(options->compress_alg) || + _read(cd, &options->compress_alg, sizeof(options->compress_alg)) != sizeof(options->compress_alg)) { + fprintf(stderr, "[-] %s(): compression algorithm invalid\n", __func__); + return -1; + } + + len -= sizeof(options->compress_alg); + + if (options->compress_alg != MEMCR_COMPRESS_NONE +#ifdef COMPRESS_LZ4 + && options->compress_alg != MEMCR_COMPRESS_LZ4 +#endif +#ifdef COMPRESS_ZSTD + && options->compress_alg != MEMCR_COMPRESS_ZSTD +#endif + ) { + /* skip if not support */ + fprintf(stderr, "[-] %s(): compression algorithm not supported: %d\n", __func__, options->compress_alg); + break; + } + + options->is_compress_alg = TRUE; + const char *caToS[] = { + [MEMCR_COMPRESS_NONE] = "none", + [MEMCR_COMPRESS_LZ4] = "lz4", + [MEMCR_COMPRESS_ZSTD] = "zstd"}; + fprintf(stdout, "[+] read compress alg for this checkpoint: %s\n", caToS[options->compress_alg]); + break; + } + default: + fprintf(stderr, "[-] %s(): checkpoint option invalid: %d\n", __func__, option); + } + } + break; + } + case MEMCR_RESTORE: { + /* nothing more to read for RESTORE*/ + fprintf(stdout, "[+] read MEMCR_RESTORE for %d\n", svc_cmd->pid); + break; + } + default: + fprintf(stderr, "%s(): Error command not expected or invalid: %d!\n", __func__, svc_cmd->cmd); + return -1; + } + + return 0; +} + +static void set_checkpoint_options_dfl(pid_t pid) +{ + for (int i=0; isvc_cmd.cmd) @@ -2616,7 +2823,7 @@ static void service_command(struct service_command_ctx *svc_ctx) break; } - init_pid_checkpoint_data(svc_ctx->svc_cmd.pid); + init_pid_checkpoint_data(svc_ctx->svc_cmd.pid, checkpoint_options); ret = service_cmds_push_back(svc_ctx); if (!ret) fprintf(stdout, "[+] Checkpoint request scheduled...\n"); @@ -2669,6 +2876,7 @@ static int service_mode(const char *listen_location) struct timeval tv; int errsv; pthread_t svc_cmd_thread_id; + struct service_options checkpoint_options; if (listen_port > 0) csd = setup_listen_tcp_socket(listen_port); @@ -2716,7 +2924,19 @@ static int service_mode(const char *listen_location) continue; } - service_command(&svc_ctx); + clear_checkpoint_options(&checkpoint_options); + + if (svc_ctx.svc_cmd.cmd == MEMCR_CMDS_V2) { + size_t cmds_len = svc_ctx.svc_cmd.pid; + ret = read_command_v2(cd, &svc_ctx.svc_cmd, &checkpoint_options, cmds_len); + if (ret < 0) { + fprintf(stderr, "%s(): Error reading a command!\n", __func__); + close(cd); + continue; + } + } + + service_command(&svc_ctx, &checkpoint_options); continue; } @@ -2775,7 +2995,7 @@ static int user_interactive_mode(pid_t pid) out: unseize_target(); - cleanup_pid(pid); + cleanup_pid(pid, dfl_dump_dir); return ret; } @@ -2783,7 +3003,7 @@ static int user_interactive_mode(pid_t pid) static void usage(const char *name, int status) { fprintf(status ? stderr : stdout, - "%s [-h] [-p PID] [-d DIR] [-S DIR] [-l PORT|PATH] [-n] [-m] [-f] [-z] [-c] [-e]\n" \ + "%s [-h] [-p PID] [-d DIR] [-S DIR] [-l PORT|PATH] [-n] [-m] [-f] [-z ALG] [-c] [-e]\n" \ "options:\n" \ " -h --help help\n" \ " -p --pid target processs pid\n" \ @@ -2798,7 +3018,7 @@ static void usage(const char *name, int status) " -n --no-wait no wait for key press\n" \ " -m --proc-mem get pages from /proc/pid/mem\n" \ " -f --rss-file include file mapped memory\n" \ - " -z --compress compress memory dump\n" \ + " -z --compress compress memory dump with selected algorithm: lz4, zstd\n" \ " -c --checksum enable md5 checksum for memory dump\n" \ " -e --encrypt enable encryption of memory dump\n", name); @@ -2806,17 +3026,6 @@ static void usage(const char *name, int status) exit(status); } -static void die(const char *fmt, ...) -{ - va_list ap; - - va_start(ap, fmt); - vfprintf(stderr, fmt, ap); - va_end(ap); - - exit(1); -} - int main(int argc, char *argv[]) { int ret; @@ -2837,17 +3046,17 @@ int main(int argc, char *argv[]) { "no-wait", 0, NULL, 'n'}, { "proc-mem", 0, NULL, 'm'}, { "rss-file", 0, NULL, 'f'}, - { "compress", 0, NULL, 'z'}, + { "compress", 1, NULL, 'z'}, { "checksum", 0, NULL, 'c'}, { "encrypt", 2, 0, 'e'}, { NULL, 0, NULL, 0 } }; - dump_dir = "/tmp"; + dfl_dump_dir = "/tmp"; parasite_socket_dir = NULL; parasite_socket_use_netns = 0; - while ((opt = getopt_long(argc, argv, "hp:d:S:Nl:nmfzce::", long_options, &option_index)) != -1) { + while ((opt = getopt_long(argc, argv, "hp:d:S:Nl:nmfz:ce::", long_options, &option_index)) != -1) { switch (opt) { case 'h': usage(argv[0], 0); @@ -2856,7 +3065,7 @@ int main(int argc, char *argv[]) pid = atoi(optarg); break; case 'd': - dump_dir = optarg; + dfl_dump_dir = optarg; break; case 'S': parasite_socket_dir = optarg; @@ -2878,10 +3087,27 @@ int main(int argc, char *argv[]) rss_file = 1; break; case 'z': + if (optarg) { + if (strcmp(optarg, "lz4") == 0) { #ifndef COMPRESS_LZ4 - die("not enabled, recompile with COMPRESS_LZ4=1\n"); + die("not enabled, recompile with COMPRESS_LZ4=1\n"); +#endif + dfl_compress_alg = MEMCR_COMPRESS_LZ4; + } else if (strcmp(optarg, "zstd") == 0) { +#ifndef COMPRESS_ZSTD + die("not enabled, recompile with COMPRESS_ZSTD=1\n"); #endif - compress = 1; + dfl_compress_alg = MEMCR_COMPRESS_ZSTD; + } + } else { +#ifdef COMPRESS_LZ4 + dfl_compress_alg = MEMCR_COMPRESS_LZ4; +#elif defined COMPRESS_ZSTD + dfl_compress_alg = MEMCR_COMPRESS_ZSTD; +#else + die("not enabled, recompile with COMPRESS_LZ4=1 and/or COMPRESS_ZSTD=1\n"); +#endif + } break; case 'c': #ifndef CHECKSUM_MD5 diff --git a/memcr.h b/memcr.h index 05dce41..aab6aa1 100644 --- a/memcr.h +++ b/memcr.h @@ -43,14 +43,35 @@ typedef enum { typedef enum { MEMCR_CHECKPOINT = 100, - MEMCR_RESTORE + MEMCR_RESTORE, + MEMCR_CMDS_V2 } memcr_svc_cmd; +typedef enum { + MEMCR_CHECKPOINT_DUMPDIR = 200, + MEMCR_CHECKPOINT_COMPRESS_ALG, +} memcr_svc_checkpoint_options; + +#define MEMCR_DUMPDIR_LEN_MAX 1024 + +typedef enum { + MEMCR_COMPRESS_NONE = 0, + MEMCR_COMPRESS_LZ4, + MEMCR_COMPRESS_ZSTD +} memcr_compress_alg; + struct service_command { memcr_svc_cmd cmd; pid_t pid; } __attribute__((packed)); +struct service_options { + int is_dump_dir; + char dump_dir[MEMCR_DUMPDIR_LEN_MAX]; + int is_compress_alg; + memcr_compress_alg compress_alg; +}; + typedef enum { MEMCR_OK = 0, MEMCR_ERROR_GENERAL = -1,