fdgrab: add --grab (only one for now)HEADmaster

1 files changed, 68 insertions, 7 deletions

diff --git a/fdgrab.c b/fdgrab.c
index 86e3cbf..5099a75 100644
--- a/fdgrab.c
+++ b/fdgrab.c
@@ -1,6 +1,8 @@
 #include <argp.h>
 #include <assert.h>
 #include <string.h>
+#include <sys/stat.h>
+#include <fcntl.h>
 #include "fdgrab.h"
 
 const char *argp_program_bug_address =
@@ -13,9 +15,11 @@ static const char * const argp_args_doc = "[command [args...]]";
 static struct argp_option options[] = {
 	{ "socket",	's', "socket", 0, "Socket path to bind to.\n"
 	  "@ at start of string will be replaced by a null for abstract socket "
-	  "(defaults to @fdgrab)",
-	  0 },
-	{ "verbose",	'v', 0, 0, "Verbose output.", 1 },
+	  "(defaults to @fdgrab)", 0 },
+	{ "grab",	'g', "pid:fd:name", 0, "fd to forcefully grab from target pid. "
+	  "If given socket receive mode is disabled unless explicitely given, in which "
+	  "case both are processed. Can be given multiple times.", 1 },
+	{ "verbose",	'v', 0, 0, "Verbose output.", 2 },
 	/* help & usage:
 	 * default help is -? which I don't like, so redefine them
 	 * (ARGP_NO_HELP set below). usage is added because usage
@@ -29,6 +33,7 @@ static struct argp_option options[] = {
 
 struct arguments {
 	int verbose;
+	char *grab;
 	char sun_path[108]; /* XXX find a proper way to look this up? */
 };
 
@@ -46,6 +51,10 @@ static error_t parse_opt(int key, char *arg, struct argp_state *state) {
 		}
 		return 0;
 
+	case 'g':
+		arguments->grab = arg; // XXX
+		return 0;
+
 	case 'v':
 		arguments->verbose++;
 		return 0;
@@ -83,6 +92,48 @@ static void replace_str(const char *orig, char **arg, char *pattern, char *repla
 	}
 }
 
+static int handle_grab(struct arguments *arguments, int argc,
+		       char *orig_argv[], char *argv[]) {
+	int rc, i, newfd;
+	char *pid, *fd, *name;
+
+	pid = arguments->grab;
+	fd = strchr(pid, ':');
+	if (!fd || !fd[1])
+		return logerr(-EINVAL, "Invalid argument for grab, must be pid:fd[:name], got %s", pid);
+	fd[0] = 0;
+	fd++;
+	name = strchr(fd, ':');
+	if (name) {
+		name[0] = 0;
+		name++;
+	}
+	if (!name || !name[0])
+		name = fd;
+
+	char procfd[256];
+	rc = snprintf(procfd, sizeof(procfd), "/proc/%s/fd/%s", pid, fd);
+	assert((size_t)rc < sizeof(procfd));
+	rc = open(procfd, O_RDWR);
+	if (rc < 0)
+		return logerr(-errno, "Could not open %s", procfd);
+	newfd = rc;
+
+	char pattern[256], replace[11];
+
+	rc = snprintf(pattern, sizeof(pattern), "{%s}", name);
+	if ((size_t)rc >= sizeof(pattern))
+		return logerr(-EINVAL, "pattern too long in %s", name);
+	rc = snprintf(replace, sizeof(replace), "%d", newfd);
+	assert((size_t)rc < sizeof(replace));
+
+	for (i = 0; i < argc; i++) {
+		replace_str(orig_argv[i], argv + i, pattern, replace);
+	}
+	return 0;
+}
+
+
 static int handle_socket(struct arguments *arguments, int argc,
 			 char *orig_argv[], char *argv[]) {
 	int fd, rc, i;
@@ -114,6 +165,7 @@ static int handle_socket(struct arguments *arguments, int argc,
 		pattern[len] = '}';
 		pattern[len+1] = 0;
 		rc = snprintf(replace, sizeof(replace), "%d", rc);
+		assert((size_t)rc < sizeof(replace));
 
 		/* XXX fix replace_str and remove check */
 		if (strlen(replace) > strlen(pattern))
@@ -155,10 +207,19 @@ int main(int argc, char *argv[]) {
 		}
 	}
 
-	rc = handle_socket(&arguments, argc - firstarg,
-			   argv + firstarg, child_argv);
-	if (rc < 0)
-		return rc;
+	if (arguments.grab) {
+		rc = handle_grab(&arguments, argc - firstarg,
+				 argv + firstarg, child_argv);
+		if (rc < 0)
+			return rc;
+	}
+
+	if (!arguments.grab || arguments.sun_path[0]) {
+		rc = handle_socket(&arguments, argc - firstarg,
+				   argv + firstarg, child_argv);
+		if (rc < 0)
+			return rc;
+	}
 
 	return execvp(child_argv[0], child_argv);
 }