Refactor project, offer uio variant of region-rw

11 files changed, 241 insertions, 85 deletions

diff --git a/Makefile b/Makefile
index c9dd6b5..71ce55b 100644
--- a/Makefile
+++ b/Makefile
@@ -7,17 +7,20 @@ WARNINGS := -Wall -Wextra -Wpedantic -Wformat=2 -Wstrict-aliasing=3 -Wstrict-ove
 	-Wfloat-equal -Wcast-align -Wpointer-arith -Wchar-subscripts -Warray-bounds=2
 
 override CFLAGS ?= -g
-override CFLAGS += -std=c99 -D_DEFAULT_SOURCE $(WARNINGS)
+override CFLAGS += -std=c99 $(WARNINGS)
+override CPPFLAGS += -Isrc
 
-bins = proc-region-rw binsearch bintrim
+bins = ptrace-region-rw uio-region-rw binsearch bintrim
 all: $(bins)
 
 $(bins): %:
 	$(LINK.c) $^ $(LDLIBS) -o $@
 
-proc-region-rw: proc-region-rw.c
-binsearch: binsearch.c
-bintrim: bintrim.c
+ptrace-region-rw: src/ptrace-region-rw.c src/cli/proc-region-rw.c src/io/io-ptrace.c
+uio-region-rw: private CPPFLAGS += -D_GNU_SOURCE
+uio-region-rw: src/uio-region-rw.c src/cli/proc-region-rw.c src/io/io-uio.c
+binsearch: src/binsearch.c
+bintrim: src/bintrim.c
 
 install-bin: $(bins)
 	install -Dm755 $^ -t "$(DESTDIR)$(PREFIX)$(bindir)"
diff --git a/contrib/brute-map.bash b/contrib/brute-map.bash
index f37581f..4a5a811 100755
--- a/contrib/brute-map.bash
+++ b/contrib/brute-map.bash
@@ -5,7 +5,7 @@
 while read -r region; do
    offset=$(printf '%d' "0x$(awk '{print $3}' <<<"$region")")
    if ((offset == 0)); then
-      offset=$(binsearch <(proc-region-rw "$1" read <<<"$region" 2>/dev/null | bintrim) $3 < "$2")
+      offset=$(binsearch <(ptrace-region-rw "$1" read <<<"$region" 2>/dev/null | bintrim) $3 < "$2")
       if [[ -n "$offset" ]]; then
          hex=$(printf '%.8x' "$offset")
          awk '{printf "%s %s %s %s %s %s\n", $1, $2, "'"$hex"'", $4, $5, $6, $7}' <<<"$region"
diff --git a/binsearch.c b/src/binsearch.c
index a4195ba..a4195ba 100644
--- a/binsearch.c
+++ b/src/binsearch.c
diff --git a/bintrim.c b/src/bintrim.c
index 2a2018e..2a2018e 100644
--- a/bintrim.c
+++ b/src/bintrim.c
diff --git a/proc-region-rw.c b/src/cli/proc-region-rw.c
index f25fe2b..fa8f343 100644
--- a/proc-region-rw.c
+++ b/src/cli/proc-region-rw.c
@@ -1,12 +1,19 @@
 #include <stdio.h>
 #include <stdlib.h>
-#include <stdbool.h>
 #include <string.h>
-#include <sys/ptrace.h>
-#include <sys/wait.h>
 #include <err.h>
 
-// It's recommended to `setcap cap_sys_ptrace=eip proc-region-rw` to run this tool without sudo
+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;
@@ -27,31 +34,14 @@ struct context {
       } mode;
    } op;
 
-   struct {
-      FILE *mem;
-      pid_t pid;
-   } proc;
+   struct io io;
 };
 
-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);
-}
-
-static void
+static inline void
 context_init(struct context *ctx, size_t argc, const char *argv[])
 {
-   if (argc < 3)
-      usage(argv[0]);
-
-   size_t arg = 1;
+   size_t arg = 0;
    *ctx = (struct context){0};
-   ctx->proc.pid = strtoull(argv[arg++], NULL, 10);
 
    {
       bool m = false, w = false, r = false;
@@ -89,11 +79,6 @@ context_init(struct context *ctx, size_t argc, const char *argv[])
 
       ctx->op.wm.size = ftell(ctx->op.wm.src);
    }
-
-   char path[128];
-   snprintf(path, sizeof(path), "/proc/%u/mem", ctx->proc.pid);
-   if (!(ctx->proc.mem = fopen(path, "w+b")))
-      err(EXIT_FAILURE, "fopen(%s)", path);
 }
 
 static void
@@ -102,32 +87,11 @@ context_release(struct context *ctx)
    if (ctx->op.wm.src)
       fclose(ctx->op.wm.src);
 
-   fclose(ctx->proc.mem);
    free(ctx->buf);
    *ctx = (struct context){0};
 }
 
 static 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);
-}
-
-static void
 region_cb(const char *line, void *data)
 {
    struct context *ctx = data;
@@ -159,22 +123,12 @@ region_cb(const char *line, void *data)
    if (!(ctx->buf = realloc(ctx->buf, len)))
       err(EXIT_FAILURE, "realloc");
 
-   clearerr(ctx->proc.mem);
    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);
-
-      if (fseek(ctx->proc.mem, start, SEEK_SET) != 0)
-         err(EXIT_FAILURE, "fseek");
-
-      const size_t wd = fwrite(ctx->buf, 1, rd, ctx->proc.mem);
-
-      if (ferror(ctx->proc.mem)) {
-         warn("fread(/proc/%u/mem)", ctx->proc.pid);
-         return;
-      }
+      const size_t wd = ctx->io.write(&ctx->io, ctx->buf, start, rd);
 
       if (ctx->op.mode == MODE_WRITE) {
          if (rlen > wd) {
@@ -190,36 +144,51 @@ region_cb(const char *line, void *data)
          }
       }
    } else {
-      if (fseek(ctx->proc.mem, start, SEEK_SET) != 0)
-         err(EXIT_FAILURE, "fseek");
-
-      const size_t rd = fread(ctx->buf, 1, len, ctx->proc.mem);
-
-      if (ferror(ctx->proc.mem))
-         warn("fread(/proc/%u/mem)", ctx->proc.pid);
+      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
-main(int argc, const char *argv[])
+proc_region_rw(int argc, const char *argv[], bool (*io_init)(struct io*, const pid_t))
 {
-   struct context ctx;
-   context_init(&ctx, argc, argv);
+   if (argc < 3)
+      usage(argv[0]);
 
-   if (ptrace(PTRACE_ATTACH, ctx.proc.pid, NULL, NULL) == -1L)
-      err(EXIT_FAILURE, "ptrace(PTRACE_ATTACH, %u, NULL, NULL)", ctx.proc.pid);
+   const pid_t pid = strtoull(argv[1], NULL, 10);
+   struct context ctx;
+   context_init(&ctx, argc - 2, argv + 2);
 
-   {
-      int status;
-      if (waitpid(ctx.proc.pid, &status, 0) == -1 || !WIFSTOPPED(status))
-         err(EXIT_FAILURE, "waitpid");
-   }
+   if (!io_init(&ctx.io, pid))
+       return EXIT_FAILURE;
 
    for_each_line_in_file(stdin, region_cb, &ctx);
-   ptrace(PTRACE_DETACH, ctx.proc.pid, 1, 0);
+   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);
+}