util: introduce memstream-util

There is many pitfalls in using memstream.
Let's introduce a wrapper to make us safely use it.

5 files changed, 164 insertions, 0 deletions

diff --git a/src/basic/memstream-util.c b/src/basic/memstream-util.c
new file mode 100644
index 0000000000..4e147fd78f
--- /dev/null
+++ b/src/basic/memstream-util.c
@@ -0,0 +1,75 @@
+/* SPDX-License-Identifier: LGPL-2.1-or-later */
+
+#include "alloc-util.h"
+#include "fd-util.h"
+#include "fileio.h"
+#include "log.h"
+#include "memstream-util.h"
+
+void memstream_done(MemStream *m) {
+        assert(m);
+
+        /* First, close file stream, as the buffer may be reallocated on close. */
+        safe_fclose(m->f);
+
+        /* Then, free buffer. */
+        free(m->buf);
+}
+
+FILE* memstream_init(MemStream *m) {
+        assert(m);
+        assert(!m->f);
+
+        m->f = open_memstream_unlocked(&m->buf, &m->sz);
+        return m->f;
+}
+
+int memstream_finalize(MemStream *m, char **ret_buf, size_t *ret_size) {
+        int r;
+
+        assert(m);
+        assert(m->f);
+        assert(ret_buf);
+
+        /* Add terminating NUL, so that the output buffer is a valid string. */
+        fputc('\0', m->f);
+
+        r = fflush_and_check(m->f);
+        if (r < 0)
+                return r;
+
+        m->f = safe_fclose(m->f);
+
+        /* On fclose(), the buffer may be reallocated, and may trigger OOM. */
+        if (!m->buf)
+                return -ENOMEM;
+
+        assert(m->sz > 0);
+
+        *ret_buf = TAKE_PTR(m->buf);
+        if (ret_size)
+                *ret_size = m->sz - 1;
+
+        m->sz = 0; /* For safety when the MemStream object will be reused later. */
+        return 0;
+}
+
+int memstream_dump_internal(
+                int level,
+                int error,
+                const char *file,
+                int line,
+                const char *func,
+                MemStream *m) {
+
+        _cleanup_free_ char *buf = NULL;
+        int r;
+
+        assert(m);
+
+        r = memstream_finalize(m, &buf, NULL);
+        if (r < 0)
+                return log_full_errno(level, r, "Failed to flush memstream: %m: %m");
+
+        return log_dump_internal(level, error, file, line, func, buf);
+}
diff --git a/src/basic/memstream-util.h b/src/basic/memstream-util.h
new file mode 100644
index 0000000000..1aa5651bcf
--- /dev/null
+++ b/src/basic/memstream-util.h
@@ -0,0 +1,27 @@
+/* SPDX-License-Identifier: LGPL-2.1-or-later */
+#pragma once
+
+#include <stdio.h>
+
+#include "macro.h"
+
+typedef struct MemStream {
+        FILE *f;
+        char *buf;
+        size_t sz;
+} MemStream;
+
+void memstream_done(MemStream *m);
+FILE* memstream_init(MemStream *m);
+int memstream_finalize(MemStream *m, char **ret_buf, size_t *ret_size);
+
+/* This finalizes the passed memstream. */
+int memstream_dump_internal(
+                int level,
+                int error,
+                const char *file,
+                int line,
+                const char *func,
+                MemStream *m);
+#define memstream_dump(level, m)                                        \
+        memstream_dump_internal(level, 0, PROJECT_FILE, __LINE__, __func__, m)
diff --git a/src/basic/meson.build b/src/basic/meson.build
index 9358a40001..5577c6bc78 100644
--- a/src/basic/meson.build
+++ b/src/basic/meson.build
@@ -53,6 +53,7 @@ basic_sources = files(
         'memfd-util.c',
         'memory-util.c',
         'mempool.c',
+        'memstream-util.c',
         'mkdir.c',
         'mountpoint-util.c',
         'namespace-util.c',
diff --git a/src/test/meson.build b/src/test/meson.build
index d7eb2a857d..84d642bf6e 100644
--- a/src/test/meson.build
+++ b/src/test/meson.build
@@ -115,6 +115,7 @@ simple_tests += files(
         'test-memfd-util.c',
         'test-memory-util.c',
         'test-mempool.c',
+        'test-memstream-util.c',
         'test-mkdir.c',
         'test-modhex.c',
         'test-mountpoint-util.c',
diff --git a/src/test/test-memstream-util.c b/src/test/test-memstream-util.c
new file mode 100644
index 0000000000..254bdcaa15
--- /dev/null
+++ b/src/test/test-memstream-util.c
@@ -0,0 +1,60 @@
+/* SPDX-License-Identifier: LGPL-2.1-or-later */
+
+#include "memstream-util.h"
+#include "string-util.h"
+#include "tests.h"
+
+TEST(memstream_done) {
+        _cleanup_(memstream_done) MemStream m = {};
+
+        assert_se(memstream_init(&m));
+}
+
+TEST(memstream_empty) {
+        _cleanup_(memstream_done) MemStream m = {};
+        _cleanup_free_ char *buf = NULL;
+        size_t sz;
+
+        assert_se(memstream_init(&m));
+        assert_se(memstream_finalize(&m, &buf, &sz) >= 0);
+        assert_se(streq(buf, ""));
+        assert_se(sz == 0);
+}
+
+TEST(memstream) {
+        _cleanup_(memstream_done) MemStream m = {};
+        _cleanup_free_ char *buf = NULL;
+        size_t sz;
+        FILE *f;
+
+        assert_se(f = memstream_init(&m));
+        fputs("hoge", f);
+        fputs("おはよう！", f);
+        fputs(u8"😀😀😀", f);
+        assert_se(memstream_finalize(&m, &buf, &sz) >= 0);
+        assert_se(streq(buf, u8"hogeおはよう！😀😀😀"));
+        assert_se(sz == strlen(u8"hogeおはよう！😀😀😀"));
+
+        buf = mfree(buf);
+
+        assert_se(f = memstream_init(&m));
+        fputs("second", f);
+        assert_se(memstream_finalize(&m, &buf, &sz) >= 0);
+        assert_se(streq(buf, "second"));
+        assert_se(sz == strlen("second"));
+}
+
+TEST(memstream_dump) {
+        _cleanup_(memstream_done) MemStream m = {};
+        FILE *f;
+
+        assert_se(f = memstream_init(&m));
+        fputs("first", f);
+        assert_se(memstream_dump(LOG_DEBUG, &m) >= 0);
+
+        assert_se(f = memstream_init(&m));
+        fputs("second", f);
+        assert_se(memstream_dump(LOG_DEBUG, &m) >= 0);
+}
+
+DEFINE_TEST_MAIN(LOG_DEBUG);