Refactor project, offer uio variant of region-rw

9 files changed, 506 insertions, 0 deletions

diff --git a/src/binsearch.c b/src/binsearch.c
new file mode 100644
index 0000000..a4195ba
--- /dev/null
+++ b/src/binsearch.c
@@ -0,0 +1,84 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <stdbool.h>
+#include <err.h>
+
+static void
+usage(const char *argv0)
+{
+   fprintf(stderr, "usage: %s needle [window-size] < haystack\n", argv0);
+   exit(EXIT_FAILURE);
+}
+
+static const char*
+search(const char *haystack, const char *needle, const size_t window_size)
+{
+   for (const char *s = haystack; s < haystack + window_size; ++s) {
+      if (memcmp(s, needle, window_size))
+         continue;
+      return s;
+   }
+   return NULL;
+}
+
+static void
+search_and_exit_if_match(const char *haystack, const char *needle, const size_t window_size, const size_t offset)
+{
+   const char *match;
+   if ((match = search(haystack, needle, window_size))) {
+      printf("%zu\n", offset + match - haystack);
+      free((void*)needle);
+      free((void*)haystack);
+      exit(EXIT_SUCCESS);
+   }
+}
+
+int
+main(int argc, const char *argv[])
+{
+   // default incase failure, or cant get size of file
+   size_t window_size = 4096 * 1024;
+   bool has_window_size = false;
+
+   if (argc < 2)
+      usage(argv[0]);
+   else if (argc > 2) {
+      window_size = strtoull(argv[2], NULL, 10);
+      has_window_size = true;
+   }
+
+   FILE *f;
+   if (!(f = fopen(argv[1], "rb")))
+      err(EXIT_FAILURE, "fopen(%s)", argv[1]);
+
+   if (!has_window_size) {
+      fseek(f, 0, SEEK_END);
+      const long tell = ftell(f);
+      window_size = (tell > 0 ? (size_t)tell : window_size);
+      fseek(f, 0, SEEK_SET);
+   }
+
+   char *needle;
+   if (!(needle = malloc(window_size)))
+      err(EXIT_FAILURE, "malloc");
+
+   window_size = fread(needle, 1, window_size, f);
+   fclose(f);
+
+   char *haystack;
+   if (!(haystack = calloc(2, window_size)))
+         err(EXIT_FAILURE, "calloc");
+
+   size_t rd = 0, offset = 0;
+   while ((rd = fread(haystack + window_size * !!offset, 1 + !offset, window_size, stdin))) {
+      search_and_exit_if_match(haystack, needle, (rd >= window_size ? window_size : 0), offset);
+      offset += window_size;
+      memmove(haystack, haystack + window_size, window_size);
+   }
+
+   search_and_exit_if_match(haystack, needle, (rd >= window_size ? window_size : 0), offset);
+   free(needle);
+   free(haystack);
+   return EXIT_FAILURE;
+}
diff --git a/src/bintrim.c b/src/bintrim.c
new file mode 100644
index 0000000..2a2018e
--- /dev/null
+++ b/src/bintrim.c
@@ -0,0 +1,44 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdbool.h>
+#include <string.h>
+#include <err.h>
+
+int
+main(int argc, const char *argv[])
+{
+   unsigned char trim = 0;
+
+   if (argc > 1)
+      trim = strtoul(argv[1], NULL, 10);
+
+   bool leading = true;
+   size_t rd, out_sz = 0, out_allocated = 0;
+   char buf[4096], *out = NULL;
+   while ((rd = fread(buf, 1, sizeof(buf), stdin))) {
+      for (const char *s = buf; s < buf + rd; ++s) {
+         if (*s == trim && leading)
+            continue;
+
+         if (out_sz >= out_allocated) {
+            if (!(out = realloc(out, out_allocated += sizeof(buf))))
+               err(EXIT_FAILURE, "realloc");
+         }
+
+         out[out_sz++] = *s;
+         leading = false;
+      }
+
+      const char *s;
+      for (s = out + (out_sz ? out_sz - 1 : 0); s > out && *s == trim; --s);
+
+      const size_t to_write = (size_t)(s - out);
+      if (fwrite(out, 1, to_write, stdout) != to_write)
+         err(EXIT_FAILURE, "fwrite");
+
+      memmove(out, s, (out_sz = out_sz - to_write));
+   }
+
+   free(out);
+   return EXIT_SUCCESS;
+}
diff --git a/src/cli/proc-region-rw.c b/src/cli/proc-region-rw.c
new file mode 100644
index 0000000..fa8f343
--- /dev/null
+++ b/src/cli/proc-region-rw.c
@@ -0,0 +1,194 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <err.h>
+
+static void
+usage(const char *argv0)
+{
+   fprintf(stderr, "usage: %s pid map file [offset] [len] < regions\n"
+                   "       %s pid write file [offset] [len] < regions\n"
+                   "       %s pid read [offset] [len] < regions\n"
+                   "       regions must be in /proc/<pid>/maps format", argv0, argv0, argv0);
+   exit(EXIT_FAILURE);
+}
+
+#include "io/io.h"
+
+struct context {
+   void *buf;
+
+   struct {
+      size_t offset, len;
+      bool has_offset, has_len;
+
+      struct {
+         FILE *src;
+         size_t size;
+      } wm; // write/map
+
+      enum {
+         MODE_MAP,
+         MODE_WRITE,
+         MODE_READ
+      } mode;
+   } op;
+
+   struct io io;
+};
+
+static inline void
+context_init(struct context *ctx, size_t argc, const char *argv[])
+{
+   size_t arg = 0;
+   *ctx = (struct context){0};
+
+   {
+      bool m = false, w = false, r = false;
+      const char *mode = argv[arg++];
+      if (!(m = !strcmp(mode, "map")) && !(w = !strcmp(mode, "write")) && !(r = !strcmp(mode, "read")))
+         errx(EXIT_FAILURE, "mode must be write or read");
+
+      ctx->op.mode = (m ? MODE_MAP : (w ? MODE_WRITE : MODE_READ));
+   }
+
+   const char *wmname = NULL;
+   if (ctx->op.mode == MODE_MAP || ctx->op.mode == MODE_WRITE) {
+      if (argc < arg + 1)
+         usage(argv[0]);
+
+      wmname = argv[arg++];
+   }
+
+   if (argc >= arg + 1) {
+      ctx->op.offset = strtoull(argv[arg++], NULL, 10);
+      ctx->op.has_offset = true;
+   }
+
+   if (argc >= arg + 1) {
+      ctx->op.len = strtoull(argv[arg++], NULL, 10);
+      ctx->op.has_len = true;
+   }
+
+   if (wmname) {
+      if (!(ctx->op.wm.src = fopen(wmname, "rb")))
+         err(EXIT_FAILURE, "fopen(%s)", wmname);
+
+      if (fseek(ctx->op.wm.src, 0, SEEK_END) != 0)
+         err(EXIT_FAILURE, "fseek");
+
+      ctx->op.wm.size = ftell(ctx->op.wm.src);
+   }
+}
+
+static void
+context_release(struct context *ctx)
+{
+   if (ctx->op.wm.src)
+      fclose(ctx->op.wm.src);
+
+   free(ctx->buf);
+   *ctx = (struct context){0};
+}
+
+static void
+region_cb(const char *line, void *data)
+{
+   struct context *ctx = data;
+   unsigned long start, end, region_offset;
+   if (sscanf(line, "%lx-%lx %*s %lx", &start, &end, &region_offset) < 3) {
+      warnx("failed to parse mapping:\n%s", line);
+      return;
+   }
+
+   warnx("%s", line);
+   start += ctx->op.offset;
+
+   if (start > end) {
+      warnx("write offset %lx is out of bounds", start);
+      return;
+   }
+
+   region_offset = (ctx->op.mode == MODE_MAP ? region_offset : 0);
+
+   // requested write/read
+   const size_t rlen = (ctx->op.has_len ? ctx->op.len : (ctx->op.mode == MODE_READ ? end - start : ctx->op.wm.size));
+
+   // actual write/read
+   const size_t len = (rlen > end - start ? end - start : rlen);
+
+   if (!len)
+      return;
+
+   if (!(ctx->buf = realloc(ctx->buf, len)))
+      err(EXIT_FAILURE, "realloc");
+
+   if (ctx->op.mode == MODE_MAP || ctx->op.mode == MODE_WRITE) {
+      if (fseek(ctx->op.wm.src, region_offset, SEEK_SET) != 0)
+         err(EXIT_FAILURE, "fseek");
+
+      const size_t rd = fread(ctx->buf, 1, len, ctx->op.wm.src);
+      const size_t wd = ctx->io.write(&ctx->io, ctx->buf, start, rd);
+
+      if (ctx->op.mode == MODE_WRITE) {
+         if (rlen > wd) {
+            warnx("wrote %lu bytes (%lu bytes truncated) to offset 0x%lx", wd, rlen - wd, start);
+         } else {
+            warnx("wrote %lu bytes to offset 0x%lx", wd, start);
+         }
+      } else {
+         if (rlen > wd) {
+            warnx("mapped %lu bytes (%lu bytes truncated) from offset 0x%lx to offset 0x%lx", wd, rlen - wd, region_offset, start);
+         } else {
+            warnx("mapped %lu bytes from offset 0x%lx to offset 0x%lx", wd, region_offset, start);
+         }
+      }
+   } else {
+      const size_t rd = ctx->io.read(&ctx->io, ctx->buf, start, len);
+
+      if (fwrite(ctx->buf, 1, rd, stdout) != rd)
+         err(EXIT_FAILURE, "fwrite");
+   }
+}
+
+static inline void
+for_each_line_in_file(FILE *f, void (*cb)(const char *line, void *data), void *data)
+{
+   char *buffer = NULL;
+   size_t step = 1024, allocated = 0, written = 0, read = 0;
+   do {
+      if (written + read >= allocated && !(buffer = realloc(buffer, (allocated += step) + 1)))
+         err(EXIT_FAILURE, "realloc");
+
+      buffer[(written += read)] = 0;
+
+      size_t ate = 0;
+      for (char *line = buffer, *nl; (nl = strchr(line, '\n')); line = nl + 1) {
+         *nl = 0;
+         cb(line, data);
+         ate += nl + 1 - line;
+      }
+
+      memmove(buffer, buffer + ate, (written = written - ate));
+   } while ((read = fread(buffer + written, 1, allocated - written, f)));
+   free(buffer);
+}
+
+int
+proc_region_rw(int argc, const char *argv[], bool (*io_init)(struct io*, const pid_t))
+{
+   if (argc < 3)
+      usage(argv[0]);
+
+   const pid_t pid = strtoull(argv[1], NULL, 10);
+   struct context ctx;
+   context_init(&ctx, argc - 2, argv + 2);
+
+   if (!io_init(&ctx.io, pid))
+       return EXIT_FAILURE;
+
+   for_each_line_in_file(stdin, region_cb, &ctx);
+   io_release(&ctx.io);
+   context_release(&ctx);
+   return EXIT_SUCCESS;
+}
diff --git a/src/cli/proc-region-rw.h b/src/cli/proc-region-rw.h
new file mode 100644
index 0000000..7ebae67
--- /dev/null
+++ b/src/cli/proc-region-rw.h
@@ -0,0 +1,9 @@
+#pragma once
+
+#include <stdbool.h>
+#include <sys/types.h> // pid_t
+
+struct io;
+
+int
+proc_region_rw(int argc, const char *argv[], bool (*io_init)(struct io*, const pid_t));
diff --git a/src/io/io-ptrace.c b/src/io/io-ptrace.c
new file mode 100644
index 0000000..50d0766
--- /dev/null
+++ b/src/io/io-ptrace.c
@@ -0,0 +1,82 @@
+#include "io.h"
+#include <stdio.h>
+#include <err.h>
+#include <sys/ptrace.h>
+#include <sys/wait.h>
+
+static size_t
+io_ptrace_do(const struct io *io, void *ptr, const size_t offset, const size_t size, size_t (*iofun)(void*, size_t, size_t, FILE*))
+{
+   if (fseek(io->backing, offset, SEEK_SET) != 0) {
+      warn("fseek(/proc/%u/mem, %zu)", io->pid, offset);
+      return 0;
+   }
+
+   return iofun(ptr, 1, size, io->backing);
+}
+
+static size_t
+io_ptrace_write(const struct io *io, const void *ptr, const size_t offset, const size_t size)
+{
+   clearerr(io->backing);
+   const size_t ret = io_ptrace_do(io, (void*)ptr, offset, size, (size_t(*)())fwrite);
+
+   if (ferror(io->backing))
+      warn("fwrite(/proc/%u/mem)", io->pid);
+
+   return ret;
+}
+
+static size_t
+io_ptrace_read(const struct io *io, void *ptr, const size_t offset, const size_t size)
+{
+   clearerr(io->backing);
+   const size_t ret = io_ptrace_do(io, ptr, offset, size, fread);
+
+   if (ferror(io->backing))
+      warn("fread(/proc/%u/mem)", io->pid);
+
+   return ret;
+}
+
+static void
+io_ptrace_cleanup(struct io *io)
+{
+   if (io->backing)
+      fclose(io->backing);
+
+   if (io->pid)
+      ptrace(PTRACE_DETACH, io->pid, 1, 0);
+}
+
+bool
+io_ptrace_init(struct io *io, const pid_t pid)
+{
+   *io = (struct io){ .pid = pid, .read = io_ptrace_read, .write = io_ptrace_write, .cleanup = io_ptrace_cleanup };
+
+   if (ptrace(PTRACE_ATTACH, pid, NULL, NULL) == -1L) {
+      warn("ptrace(PTRACE_ATTACH, %u, NULL, NULL)", pid);
+      goto fail;
+   }
+
+   {
+      int status;
+      if (waitpid(pid, &status, 0) == -1 || !WIFSTOPPED(status)) {
+         warn("waitpid(%u) == %d", pid, status);
+         goto fail;
+      }
+   }
+
+   char path[128];
+   snprintf(path, sizeof(path), "/proc/%u/mem", pid);
+   if (!(io->backing = fopen(path, "w+b"))) {
+      warn("fopen(%s)", path);
+      goto fail;
+   }
+
+   return true;
+
+fail:
+   io->cleanup(io);
+   return false;
+}
diff --git a/src/io/io-uio.c b/src/io/io-uio.c
new file mode 100644
index 0000000..dd9f873
--- /dev/null
+++ b/src/io/io-uio.c
@@ -0,0 +1,41 @@
+#include "io.h"
+#include <stdint.h>
+#include <err.h>
+#include <sys/uio.h>
+
+static size_t
+io_uio_do(const struct io *io, const void *ptr, const size_t offset, const size_t size, ssize_t (*iofun)(pid_t, const struct iovec*, unsigned long, const struct iovec*, unsigned long, unsigned long))
+{
+   const struct iovec lio = { .iov_base = (void*)ptr, .iov_len = size };
+   const struct iovec rio = { .iov_base = (void*)(intptr_t)offset, .iov_len = size };
+   return iofun(io->pid, &lio, 1, &rio, 1, 0);
+}
+
+static size_t
+io_uio_write(const struct io *io, const void *ptr, const size_t offset, const size_t size)
+{
+   const size_t ret = io_uio_do(io, ptr, offset, size, process_vm_writev);
+
+   if (ret == (size_t)-1)
+      warn("process_vm_writev(%u)", io->pid);
+
+   return (ret == (size_t)-1 ? 0 : ret);
+}
+
+static size_t
+io_uio_read(const struct io *io, void *ptr, const size_t offset, const size_t size)
+{
+   const size_t ret = io_uio_do(io, ptr, offset, size, process_vm_readv);
+
+   if (ret == (size_t)-1)
+      warn("process_vm_readv(%u)", io->pid);
+
+   return (ret == (size_t)-1 ? 0 : ret);
+}
+
+bool
+io_uio_init(struct io *io, const pid_t pid)
+{
+   *io = (struct io){ .pid = pid, .read = io_uio_read, .write = io_uio_write };
+   return true;
+}
diff --git a/src/io/io.h b/src/io/io.h
new file mode 100644
index 0000000..8a0639b
--- /dev/null
+++ b/src/io/io.h
@@ -0,0 +1,27 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdbool.h>
+#include <sys/types.h> // pid_t
+
+struct io {
+   void *backing;
+   size_t (*read)(const struct io *io, void *ptr, const size_t offset, const size_t size);
+   size_t (*write)(const struct io *io, const void *ptr, const size_t offset, const size_t size);
+   void (*cleanup)(struct io *io);
+   pid_t pid;
+};
+
+static inline void
+io_release(struct io *io)
+{
+   if (io->cleanup)
+      io->cleanup(io);
+   *io = (struct io){0};
+}
+
+bool
+io_uio_init(struct io *io, const pid_t pid);
+
+bool
+io_ptrace_init(struct io *io, const pid_t pid);
diff --git a/src/ptrace-region-rw.c b/src/ptrace-region-rw.c
new file mode 100644
index 0000000..135c2a5
--- /dev/null
+++ b/src/ptrace-region-rw.c
@@ -0,0 +1,12 @@
+#include "cli/proc-region-rw.h"
+#include "io/io.h"
+
+// This region-rw uses ptrace
+// This works with older kernels, but it also ensures non racy read/writes, as it stops the process.
+// It's recommended to `setcap cap_sys_ptrace=eip ptrace-region-rw` to run this tool without sudo
+
+int
+main(int argc, const char *argv[])
+{
+   return proc_region_rw(argc, argv, io_ptrace_init);
+}
diff --git a/src/uio-region-rw.c b/src/uio-region-rw.c
new file mode 100644
index 0000000..96cac49
--- /dev/null
+++ b/src/uio-region-rw.c
@@ -0,0 +1,13 @@
+#include "cli/proc-region-rw.h"
+#include "io/io.h"
+
+// This region-rw uses uio
+// It needs recent kernel, but may be racy as it reads / writes while process is running.
+// Ideal for realtime memory tools.
+// It's recommended to `setcap cap_sys_ptrace=eip uio-region-rw` to run this tool without sudo
+
+int
+main(int argc, const char *argv[])
+{
+   return proc_region_rw(argc, argv, io_uio_init);
+}