lib/subprocess: add callback feature to the subprocess library

For the libflashrom-compatible interface I'm working on, I needed the
ability to process data from the flashrom subprocess in a callback
function.

This adds a new type of subprocess_target, TARGET_CALLBACK, which can
read and write to/from a callback function.

BUG=chromium:478356
BRANCH=none
TEST=provided unit tests

Change-Id: I20b71000fc2b6b297a8617d2b03d0e91813007d1
Signed-off-by: Jack Rosenthal <jrosenth@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/1959944
Reviewed-by: Julius Werner <jwerner@chromium.org>

3 files changed, 251 insertions, 34 deletions

diff --git a/host/lib/include/subprocess.h b/host/lib/include/subprocess.h
index 61a214a2..eaf75502 100644
--- a/host/lib/include/subprocess.h
+++ b/host/lib/include/subprocess.h
@@ -36,6 +36,18 @@
  *   When writing to a buffer, subprocess_run will reserve one byte of
  *   the size for a null terminator and guarantee that the output is
  *   always NULL terminated.
+ *
+ * - TARGET_CALLBACK: when the target is provided as an input to a
+ *   process, the callback will be called occasionally to provide
+ *   input to the process. The callback should fill buf with up to
+ *   buf_sz bytes of data, and return the number of bytes
+ *   written, or negative values on error. When the target is provided
+ *   as an output to a process, the callback will be called
+ *   occasionally with buf_sz bytes of data from the output put into
+ *   buf. In this case, the return value from the callback is
+ *   ignored except for errors. The data field is for application use
+ *   and will always be passed to the data parameter of the callback
+ *   function.
  */
 struct subprocess_target {
 	enum {
@@ -44,6 +56,7 @@ struct subprocess_target {
 		TARGET_FILE,
 		TARGET_BUFFER,
 		TARGET_BUFFER_NULL_TERMINATED,
+		TARGET_CALLBACK,
 	} type;
 	union {
 		int fd;
@@ -53,19 +66,20 @@ struct subprocess_target {
 			size_t size;
 
 			/*
-			 * This variable is used internally by subprocess_run
-			 * and shouldn't be operated on by the caller.
-			 */
-			int _pipefd[2];
-
-			/*
 			 * This variable is the output of the number of bytes
 			 * read or written. It should be read by the caller, not
 			 * set.
 			 */
 			size_t bytes_consumed;
 		} buffer;
+		struct {
+			ssize_t (*cb)(char *buf, size_t buf_sz, void *data);
+			void *data;
+		} callback;
 	};
+	struct {
+		int pipefd[2];
+	} priv;
 };
 
 /**
diff --git a/host/lib/subprocess.c b/host/lib/subprocess.c
index 95a6e4d2..be535838 100644
--- a/host/lib/subprocess.c
+++ b/host/lib/subprocess.c
@@ -11,6 +11,8 @@
 
 #include "subprocess.h"
 
+#define MAX_CB_BUF_SIZE 2048
+
 static char *program_name;
 
 __attribute__((constructor, used))
@@ -25,7 +27,8 @@ static int init_target_private(struct subprocess_target *target)
 	switch (target->type) {
 	case TARGET_BUFFER:
 	case TARGET_BUFFER_NULL_TERMINATED:
-		return pipe(target->buffer._pipefd);
+	case TARGET_CALLBACK:
+		return pipe(target->priv.pipefd);
 	default:
 		return 0;
 	}
@@ -60,15 +63,16 @@ static int connect_process_target(struct subprocess_target *target, int fd)
 		break;
 	case TARGET_BUFFER:
 	case TARGET_BUFFER_NULL_TERMINATED:
+	case TARGET_CALLBACK:
 		switch (fd) {
 		case STDIN_FILENO:
-			target_fd = target->buffer._pipefd[0];
-			close(target->buffer._pipefd[1]);
+			target_fd = target->priv.pipefd[0];
+			close(target->priv.pipefd[1]);
 			break;
 		case STDOUT_FILENO:
 		case STDERR_FILENO:
-			target_fd = target->buffer._pipefd[1];
-			close(target->buffer._pipefd[0]);
+			target_fd = target->priv.pipefd[1];
+			close(target->priv.pipefd[0]);
 			break;
 		default:
 			return -1;
@@ -81,9 +85,8 @@ static int connect_process_target(struct subprocess_target *target, int fd)
 	return dup2(target_fd, fd);
 }
 
-static int process_target_input(struct subprocess_target *target)
+static int process_target_input_buffer(struct subprocess_target *target)
 {
-	int rv = 0;
 	ssize_t write_rv;
 	size_t bytes_to_write;
 	char *buf;
@@ -96,30 +99,79 @@ static int process_target_input(struct subprocess_target *target)
 		bytes_to_write = strlen(target->buffer.buf);
 		break;
 	default:
-		return 0;
+		return -1;
 	}
 
-	close(target->buffer._pipefd[0]);
 	buf = target->buffer.buf;
 	while (bytes_to_write) {
-		write_rv =
-			write(target->buffer._pipefd[1], buf, bytes_to_write);
-		if (write_rv <= 0) {
-			rv = -1;
-			goto cleanup;
-		}
+		write_rv = write(target->priv.pipefd[1], buf, bytes_to_write);
+		if (write_rv <= 0)
+			return -1;
 		buf += write_rv;
 		bytes_to_write -= write_rv;
 	}
 
- cleanup:
-	close(target->buffer._pipefd[1]);
+	return 0;
+}
+
+static int process_target_input_cb(struct subprocess_target *target)
+{
+	ssize_t write_rv, bytes_to_write;
+	char buf[MAX_CB_BUF_SIZE];
+	char *bufptr;
+
+	for (;;) {
+		bytes_to_write = target->callback.cb(buf, MAX_CB_BUF_SIZE,
+						     target->callback.data);
+		if (bytes_to_write < 0 || bytes_to_write > MAX_CB_BUF_SIZE)
+			return -1;
+		if (bytes_to_write == 0)
+			return 0;
+
+		bufptr = buf;
+		while (bytes_to_write) {
+			write_rv = write(target->priv.pipefd[1], bufptr,
+					 bytes_to_write);
+			if (write_rv <= 0)
+				return -1;
+			bufptr += write_rv;
+			bytes_to_write -= write_rv;
+		}
+	}
+}
+
+static int process_target_input(struct subprocess_target *target)
+{
+	int rv;
+
+	switch (target->type) {
+	case TARGET_BUFFER:
+	case TARGET_BUFFER_NULL_TERMINATED:
+	case TARGET_CALLBACK:
+		break;
+	default:
+		return 0;
+	}
+
+	close(target->priv.pipefd[0]);
+	switch (target->type) {
+	case TARGET_BUFFER:
+	case TARGET_BUFFER_NULL_TERMINATED:
+		rv = process_target_input_buffer(target);
+		break;
+	case TARGET_CALLBACK:
+		rv = process_target_input_cb(target);
+		break;
+	default:
+		return -1;
+	}
+
+	close(target->priv.pipefd[1]);
 	return rv;
 }
 
-static int process_target_output(struct subprocess_target *target)
+static int process_target_output_buffer(struct subprocess_target *target)
 {
-	int rv = 0;
 	ssize_t read_rv;
 	size_t bytes_remaining;
 
@@ -136,17 +188,14 @@ static int process_target_output(struct subprocess_target *target)
 		return 0;
 	}
 
-	close(target->buffer._pipefd[1]);
 	target->buffer.bytes_consumed = 0;
 	while (bytes_remaining) {
 		read_rv = read(
-			target->buffer._pipefd[0],
+			target->priv.pipefd[0],
 			target->buffer.buf + target->buffer.bytes_consumed,
 			bytes_remaining);
-		if (read_rv < 0) {
-			rv = -1;
-			goto cleanup;
-		}
+		if (read_rv < 0)
+			return -1;
 		if (read_rv == 0)
 			break;
 		target->buffer.bytes_consumed += read_rv;
@@ -155,9 +204,54 @@ static int process_target_output(struct subprocess_target *target)
 
 	if (target->type == TARGET_BUFFER_NULL_TERMINATED)
 		target->buffer.buf[target->buffer.bytes_consumed] = '\0';
+	return 0;
+}
+
+static int process_target_output_cb(struct subprocess_target *target)
+{
+	char buf[MAX_CB_BUF_SIZE];
+	ssize_t rv;
+
+	for (;;) {
+		rv = read(target->priv.pipefd[0], buf, MAX_CB_BUF_SIZE);
+		if (rv < 0)
+			return -1;
+		if (rv == 0)
+			break;
+		if (target->callback.cb(buf, rv, target->callback.data) < 0)
+			return -1;
+	}
+
+	return 0;
+}
+
+static int process_target_output(struct subprocess_target *target)
+{
+	int rv;
+
+	switch (target->type) {
+	case TARGET_BUFFER:
+	case TARGET_BUFFER_NULL_TERMINATED:
+	case TARGET_CALLBACK:
+		break;
+	default:
+		return 0;
+	}
+
+	close(target->priv.pipefd[1]);
+	switch (target->type) {
+	case TARGET_BUFFER:
+	case TARGET_BUFFER_NULL_TERMINATED:
+		rv = process_target_output_buffer(target);
+		break;
+	case TARGET_CALLBACK:
+		rv = process_target_output_cb(target);
+		break;
+	default:
+		return -1;
+	}
 
- cleanup:
-	close(target->buffer._pipefd[0]);
+	close(target->priv.pipefd[0]);
 	return rv;
 }
 
diff --git a/tests/subprocess_tests.c b/tests/subprocess_tests.c
index 58c8f235..75d6e493 100644
--- a/tests/subprocess_tests.c
+++ b/tests/subprocess_tests.c
@@ -164,12 +164,119 @@ static void test_subprocess_small_output_buffer(void)
 
 static void test_subprocess_return_code_failure(void)
 {
-	const char *const argv[] = {"false"};
+	const char *const argv[] = {"false", NULL};
 
 	TEST_NEQ(subprocess_run(argv, NULL, NULL, NULL), 0,
 		 "Return value of \"false\" is nonzero");
 }
 
+struct cb_ctx {
+	char buffer[49 * 1024];
+	char *ptr;
+};
+
+static ssize_t input_cb(char *buf, size_t buf_sz, void *data)
+{
+	struct cb_ctx *ctx = (struct cb_ctx *)data;
+	size_t len = (ctx->buffer + sizeof(ctx->buffer)) - ctx->ptr;
+	if (len > buf_sz)
+		len = buf_sz;
+	memcpy(buf, ctx->ptr, len);
+	ctx->ptr += len;
+	return len;
+}
+
+static void test_subprocess_input_from_cb(void)
+{
+	struct cb_ctx ctx;
+	char output_buffer[sizeof(ctx.buffer)];
+	const char *const argv[] = {"cat", NULL};
+
+	/* Initialize the input buffer with some data */
+	for (size_t i = 0; i < sizeof(ctx.buffer); i++)
+		ctx.buffer[i] = (char)i;
+	ctx.ptr = ctx.buffer;
+
+	struct subprocess_target output = {
+		.type = TARGET_BUFFER,
+		.buffer = {
+			.buf = output_buffer,
+			.size = sizeof(output_buffer),
+		},
+	};
+
+	struct subprocess_target input = {
+		.type = TARGET_CALLBACK,
+		.callback = {
+			.cb = input_cb,
+			.data = &ctx,
+		},
+	};
+	TEST_EQ(subprocess_run(argv, &input, &output, NULL), 0,
+		"Return value of \"cat\" is zero.");
+	TEST_EQ(memcmp(ctx.buffer, output_buffer, sizeof(output_buffer)), 0,
+		"The input buffer is equal to the output buffer.");
+	TEST_EQ(output.buffer.bytes_consumed, sizeof(output_buffer),
+		"The entire output buffer should have been used.");
+}
+
+static ssize_t output_cb(char *buf, size_t buf_sz, void *data)
+{
+	struct cb_ctx *ctx = (struct cb_ctx *)data;
+	if (ctx->ptr + buf_sz > ctx->buffer + sizeof(ctx->buffer)) {
+		TEST_TRUE(0, "Test failed as there is not enough space in the "
+			  "output buffer.");
+		return -1;
+	}
+	memcpy(ctx->ptr, buf, buf_sz);
+	ctx->ptr += buf_sz;
+	return 0;
+}
+
+static void test_subprocess_output_to_cb(void)
+{
+	struct cb_ctx ctx;
+	char output_buffer[sizeof(ctx.buffer)];
+	const char *const argv[] = {
+		"bc", "-l", NULL
+	};
+
+	ctx.ptr = ctx.buffer;
+
+	struct subprocess_target input = {
+		.type = TARGET_BUFFER_NULL_TERMINATED,
+		.buffer = {
+			.buf = (char *)"for (i = 0; i <= 10000; i += 1) i\n",
+		},
+	};
+
+	struct subprocess_target target_via_buffer = {
+		.type = TARGET_BUFFER,
+		.buffer = {
+			.buf = output_buffer,
+			.size = sizeof(output_buffer),
+		},
+	};
+
+	struct subprocess_target target_via_cb = {
+		.type = TARGET_CALLBACK,
+		.callback = {
+			.cb = output_cb,
+			.data = &ctx,
+		},
+	};
+
+	TEST_EQ(subprocess_run(argv, &input, &target_via_buffer, NULL), 0,
+		"Return value is zero when using buffer.");
+	TEST_EQ(subprocess_run(argv, &input, &target_via_cb, NULL), 0,
+		"Return value is zero when using callback.");
+	TEST_EQ(ctx.ptr - ctx.buffer, target_via_buffer.buffer.bytes_consumed,
+		"Both commmand invocations used the same number of bytes.");
+	TEST_EQ(memcmp(output_buffer, ctx.buffer,
+		       target_via_buffer.buffer.bytes_consumed),
+		0, "Both output buffers are equivalent.");
+}
+
 int main(int argc, char *argv[])
 {
 	test_subprocess_output_to_buffer();
@@ -178,6 +285,8 @@ int main(int argc, char *argv[])
 	test_subprocess_input_null_terminated();
 	test_subprocess_small_output_buffer();
 	test_subprocess_return_code_failure();
+	test_subprocess_input_from_cb();
+	test_subprocess_output_to_cb();
 
 	return gTestSuccess ? 0 : 255;
 }