lib-test: Add ostream-final-trickle

diff --git a/src/lib-test/Makefile.am b/src/lib-test/Makefile.am
index ce7417d6428546880593a75a9553bbd900ca7cb9..cb6c66edb9dd96edacbe69dc726d24a6bfc7916c 100644 (file)
--- a/src/lib-test/Makefile.am
+++ b/src/lib-test/Makefile.am
@@ -6,6 +6,7 @@ AM_CPPFLAGS = \
 
 libtest_la_SOURCES = \
        fuzzer.c \
+       ostream-final-trickle.c \
        test-common.c \
        test-istream.c \
        test-ostream.c \
@@ -13,6 +14,7 @@ libtest_la_SOURCES = \
 
 headers = \
        fuzzer.h \
+       ostream-final-trickle.h \
        test-common.h \
        test-subprocess.h
 

diff --git a/src/lib-test/ostream-final-trickle.c b/src/lib-test/ostream-final-trickle.c
new file mode 100644 (file)
index 0000000..fc8a00f
--- /dev/null
+++ b/src/lib-test/ostream-final-trickle.c
@@ -0,0 +1,148 @@
+/* Copyright (c) 2023 Dovecot authors, see the included COPYING file */
+
+#include "lib.h"
+#include "ostream-private.h"
+#include "ostream-final-trickle.h"
+
+struct final_trickle_ostream {
+       struct ostream_private ostream;
+       struct timeout *to;
+
+       unsigned char buffer_char;
+       bool buffer_used;
+};
+
+static void
+o_stream_final_trickle_close(struct iostream_private *stream, bool close_parent)
+{
+       struct final_trickle_ostream *dstream =
+               container_of(stream, struct final_trickle_ostream,
+                            ostream.iostream);
+
+       timeout_remove(&dstream->to);
+       if (close_parent)
+               o_stream_close(dstream->ostream.parent);
+}
+
+static int
+o_stream_final_trickle_flush_buffer(struct final_trickle_ostream *dstream)
+{
+       int ret = 1;
+
+       if (dstream->buffer_used) {
+               if ((ret = o_stream_send(dstream->ostream.parent,
+                                        &dstream->buffer_char, 1)) < 0)
+                       o_stream_copy_error_from_parent(&dstream->ostream);
+               else if (ret > 0)
+                       dstream->buffer_used = FALSE;
+               if (ret != 0)
+                       timeout_remove(&dstream->to);
+       }
+       return ret;
+}
+
+static void
+o_stream_final_trickle_timeout(struct final_trickle_ostream *dstream)
+{
+       i_assert(dstream->buffer_used);
+
+       (void)o_stream_final_trickle_flush_buffer(dstream);
+}
+
+static int o_stream_final_trickle_flush(struct ostream_private *stream)
+{
+       struct final_trickle_ostream *dstream =
+               container_of(stream, struct final_trickle_ostream, ostream);
+
+       if (dstream->buffer_used)
+               return 0;
+       return o_stream_flush_parent(stream);
+}
+
+static ssize_t
+o_stream_final_trickle_sendv(struct ostream_private *stream,
+                            const struct const_iovec *iov,
+                            unsigned int iov_count)
+{
+       struct final_trickle_ostream *dstream =
+               container_of(stream, struct final_trickle_ostream, ostream);
+       ssize_t ret;
+
+       if ((ret = o_stream_final_trickle_flush_buffer(dstream)) <= 0)
+               return ret;
+       i_assert(!dstream->buffer_used);
+
+       /* send all but the last byte */
+       struct const_iovec iov_copy[iov_count];
+       memcpy(iov_copy, iov, iov_count * sizeof(*iov));
+       struct const_iovec *last_iov = &iov_copy[iov_count-1];
+
+       i_assert(last_iov->iov_len > 0);
+       last_iov->iov_len--;
+       const unsigned char *last_iov_data = last_iov->iov_base;
+       dstream->buffer_char = last_iov_data[last_iov->iov_len];
+       dstream->buffer_used = TRUE;
+       if (dstream->to == NULL) {
+               dstream->to = timeout_add_short(0,
+                       o_stream_final_trickle_timeout, dstream);
+       }
+       if (last_iov->iov_len == 0)
+               iov_count--;
+
+       if (iov_count > 0) {
+               size_t full_size = 0;
+               for (unsigned int i = 0; i < iov_count; i++)
+                       full_size += iov_copy[i].iov_len;
+               if ((ret = o_stream_sendv(stream->parent, iov_copy, iov_count)) < 0) {
+                       o_stream_copy_error_from_parent(stream);
+                       return -1;
+               }
+               if ((size_t)ret < full_size) {
+                       dstream->buffer_used = FALSE;
+                       timeout_remove(&dstream->to);
+               }
+       }
+       if (dstream->buffer_used)
+               ret++;
+
+       stream->ostream.offset += ret;
+       return ret;
+}
+
+static size_t
+o_stream_final_trickle_get_buffer_used_size(const struct ostream_private *stream)
+{
+       const struct final_trickle_ostream *dstream =
+               container_of(stream, const struct final_trickle_ostream, ostream);
+
+       return (dstream->buffer_used ? 1 : 0) +
+               o_stream_get_buffer_used_size(stream->parent);
+}
+
+static void
+o_stream_final_trickle_switch_ioloop_to(struct ostream_private *stream,
+                                       struct ioloop *ioloop)
+{
+       struct final_trickle_ostream *dstream =
+               container_of(stream, struct final_trickle_ostream, ostream);
+
+       if (dstream->to != NULL)
+               dstream->to = io_loop_move_timeout_to(ioloop, &dstream->to);
+       if (stream->parent != NULL)
+               o_stream_switch_ioloop_to(stream->parent, ioloop);
+}
+
+struct ostream *o_stream_create_final_trickle(struct ostream *output)
+{
+       struct final_trickle_ostream *dstream;
+
+       dstream = i_new(struct final_trickle_ostream, 1);
+       dstream->ostream.iostream.close = o_stream_final_trickle_close;
+       dstream->ostream.sendv = o_stream_final_trickle_sendv;
+       dstream->ostream.flush = o_stream_final_trickle_flush;
+       dstream->ostream.get_buffer_used_size = o_stream_final_trickle_get_buffer_used_size;
+       dstream->ostream.switch_ioloop_to = o_stream_final_trickle_switch_ioloop_to;
+
+       return o_stream_create(&dstream->ostream, output,
+                              o_stream_get_fd(output));
+}

diff --git a/src/lib-test/ostream-final-trickle.h b/src/lib-test/ostream-final-trickle.h
new file mode 100644 (file)
index 0000000..c97faf9
--- /dev/null
+++ b/src/lib-test/ostream-final-trickle.h
@@ -0,0 +1,9 @@
+#ifndef OSTREAM_FINAL_TRICKLE_H
+#define OSTREAM_FINAL_TRICKLE_H
+
+/* Creates a wrapper istream that delays sending the final byte until the next
+   ioloop run. This can catch bugs where caller doesn't expect the final bytes
+   to be delayed. */
+struct ostream *o_stream_create_final_trickle(struct ostream *output);
+
+#endif