From df7b68eea97da5948087b65a84392b77e52ce61b Mon Sep 17 00:00:00 2001 From: Rusty Russell Date: Fri, 10 Mar 2017 21:15:42 +1030 Subject: [PATCH] wire/wire_io: support take() arg to io_write_wire(). This simplifies memory management, but means we have to keep the original pointer. Signed-off-by: Rusty Russell --- wire/wire_io.c | 21 ++++++++++++++------- wire/wire_io.h | 2 +- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/wire/wire_io.c b/wire/wire_io.c index bb02780cb..5e5fa020c 100644 --- a/wire/wire_io.c +++ b/wire/wire_io.c @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include /* @@ -90,27 +92,32 @@ static int do_write_wire_header(int fd, struct io_plan_arg *arg) /* Both bytes written? Set up for normal write of data. */ if (arg->u2.s == INSIDE_HEADER_BIT + HEADER_LEN) - arg->u2.s = be16_to_cpu(hdr); + arg->u2.s = 0; - return arg->u2.s == 0; + return 0; } static int do_write_wire(int fd, struct io_plan_arg *arg) { ssize_t ret; + size_t totlen = tal_len(arg->u1.cp); /* Still writing header? */ if (arg->u2.s & INSIDE_HEADER_BIT) return do_write_wire_header(fd, arg); /* Normal write */ - ret = write(fd, arg->u1.cp, arg->u2.s); - if (ret <= 0) + ret = write(fd, arg->u1.cp + arg->u2.s, totlen - arg->u2.s); + if (ret < 0) return -1; - arg->u1.cp += ret; - arg->u2.s -= ret; - return arg->u2.s == 0; + arg->u2.s += ret; + if (arg->u2.s != totlen) + return 0; + + if (taken(arg->u1.cp)) + tal_free(arg->u1.cp); + return 1; } /* Write message from data (tal_count(data) gives length). */ diff --git a/wire/wire_io.h b/wire/wire_io.h index e2b085b13..de9eb13b7 100644 --- a/wire/wire_io.h +++ b/wire/wire_io.h @@ -17,7 +17,7 @@ struct io_plan *io_read_wire_(struct io_conn *conn, (next), (arg), struct io_conn *), \ (arg)) -/* Write message from data (tal_count(data) gives length). */ +/* Write message from data (tal_count(data) gives length). data can be take() */ struct io_plan *io_write_wire_(struct io_conn *conn, const u8 *data, struct io_plan *(*next)(struct io_conn *, void *),