iovec-wrapper: introduce iovec_split() and iovw_merge()

In many network protocols, the length-prefixed data format is often
used. Let's add a simple parser and builder for the format.

src/basic/iovec-wrapper.c

@@ -6,6 +6,7 @@
 #include "iovec-util.h"
 #include "iovec-wrapper.h"
 #include "string-util.h"
+#include "unaligned.h"
 
 void iovw_done(struct iovec_wrapper *iovw) {
         assert(iovw);
@@ -302,3 +303,121 @@ char* iovw_to_cstring(const struct iovec_wrapper *iovw) {
         assert(!memchr(iov.iov_base, 0, iov.iov_len));
         return TAKE_PTR(iov.iov_base);
 }
+
+int iovec_split(const struct iovec *iov, size_t length_size, struct iovec_wrapper *ret) {
+        int r;
+
+        assert(IN_SET(length_size, 1, 2, 4));
+        assert(ret);
+
+        /* This parses the input iovec as length-prefixed data, and stores the result as iovec_wrapper.
+         * Note, zero-length entries are silently dropped. */
+
+        if (!iovec_is_set(iov)) {
+                *ret = (struct iovec_wrapper) {};
+                return 0;
+        }
+
+        _cleanup_(iovw_done_free) struct iovec_wrapper iovw = {};
+        for (struct iovec i = *iov; iovec_is_set(&i); ) {
+                if (i.iov_len < length_size)
+                        return -EBADMSG;
+
+                size_t len;
+                switch (length_size) {
+                case 1:
+                        len = *(uint8_t*) i.iov_base;
+                        break;
+                case 2:
+                        len = unaligned_read_be16(i.iov_base);
+                        break;
+                case 4:
+                        len = unaligned_read_be32(i.iov_base);
+                        break;
+                default:
+                        assert_not_reached();
+                }
+
+                iovec_inc(&i, length_size);
+
+                if (len == 0)
+                        continue;
+
+                if (i.iov_len < len)
+                        return -EBADMSG;
+
+                r = iovw_extend(&iovw, i.iov_base, len);
+                if (r < 0)
+                        return r;
+
+                iovec_inc(&i, len);
+        }
+
+        *ret = TAKE_STRUCT(iovw);
+        return 0;
+}
+
+int iovw_merge(const struct iovec_wrapper *iovw, size_t length_size, struct iovec *ret) {
+        assert(IN_SET(length_size, 1, 2, 4));
+        assert(ret);
+
+        /* This is the inverse of iovec_split(), and builds a length-prefixed data from iovec_wrapper.
+         * Note, zero-length entries are silently dropped. */
+
+        size_t sz = iovw_size(iovw);
+        if (sz == 0) {
+                *ret = (struct iovec) {};
+                return 0;
+        }
+        if (sz == SIZE_MAX)
+                return -E2BIG;
+
+        if (size_multiply_overflow(length_size, iovw->count))
+                return -E2BIG;
+
+        sz = size_add(sz, iovw->count * length_size);
+        if (sz == SIZE_MAX)
+                return -E2BIG;
+
+        _cleanup_free_ uint8_t *buf = new(uint8_t, sz);
+        if (!buf)
+                return -ENOMEM;
+
+        uint8_t *p = buf;
+        FOREACH_ARRAY(iov, iovw->iovec, iovw->count) {
+                if (iov->iov_len == 0)
+                        continue;
+
+                switch (length_size) {
+                case 1:
+                        if (iov->iov_len > UINT8_MAX)
+                                return -ERANGE;
+
+                        *p = iov->iov_len;
+                        break;
+                case 2:
+                        if (iov->iov_len > UINT16_MAX)
+                                return -ERANGE;
+
+                        unaligned_write_be16(p, iov->iov_len);
+                        break;
+                case 4:
+                        if (iov->iov_len > UINT32_MAX)
+                                return -ERANGE;
+
+                        unaligned_write_be32(p, iov->iov_len);
+                        break;
+                default:
+                        assert_not_reached();
+                }
+                p += length_size;
+
+                p = mempcpy(p, iov->iov_base, iov->iov_len);
+        }
+
+        assert(sz >= (size_t) (p - buf));
+        sz = p - buf;
+
+        *ret = IOVEC_MAKE(TAKE_PTR(buf), sz);
+        return 0;
+}

src/basic/iovec-wrapper.h

@@ -68,3 +68,6 @@ void iovw_rebase(struct iovec_wrapper *iovw, void *old, void *new);
 size_t iovw_size(const struct iovec_wrapper *iovw);
 int iovw_concat(const struct iovec_wrapper *iovw, struct iovec *ret);
 char* iovw_to_cstring(const struct iovec_wrapper *iovw);
+
+int iovec_split(const struct iovec *iov, size_t length_size, struct iovec_wrapper *ret);
+int iovw_merge(const struct iovec_wrapper *iovw, size_t length_size, struct iovec *ret);

src/test/test-iovec-wrapper.c

@@ -5,6 +5,7 @@
 #include "alloc-util.h"
 #include "iovec-util.h"
 #include "iovec-wrapper.h"
+#include "random-util.h"
 #include "tests.h"
 
 TEST(iovw_compare) {
@@ -506,4 +507,196 @@ TEST(iovw_to_cstring) {
         ASSERT_STREQ(s, "foo/bar");
 }
 
+TEST(iovw_merge_and_iovec_split) {
+        _cleanup_(iovw_done_free) struct iovec_wrapper iovw = {}, iovw2 = {};
+        _cleanup_(iovec_done) struct iovec v = {}, v2 = {};
+        uint8_t *p;
+
+        struct iovec
+                a = IOVEC_MAKE_STRING("aaa"),
+                b = IOVEC_MAKE_STRING("bbbb"),
+                c = IOVEC_MAKE_STRING("ccccc");
+
+        /* single entry */
+        ASSERT_OK(iovw_extend_iov(&iovw, &a));
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint8_t), &v));
+        ASSERT_EQ(v.iov_len, 1 + a.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_OK(iovec_split(&v, sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint16_t), &v));
+        ASSERT_EQ(v.iov_len, sizeof(uint16_t) + a.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_OK(iovec_split(&v, sizeof(uint16_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint32_t), &v));
+        ASSERT_EQ(v.iov_len, sizeof(uint32_t) + a.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_OK(iovec_split(&v, sizeof(uint32_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        /* multiple entries */
+        ASSERT_OK(iovw_extend_iov(&iovw, &b));
+        ASSERT_OK(iovw_extend_iov(&iovw, &c));
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint8_t), &v));
+        ASSERT_EQ(v.iov_len, 3 + a.iov_len + b.iov_len + c.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_EQ(*p++, b.iov_len);
+        ASSERT_EQ(memcmp(p, b.iov_base, b.iov_len), 0);
+        p += b.iov_len;
+        ASSERT_EQ(*p++, c.iov_len);
+        ASSERT_EQ(memcmp(p, c.iov_base, c.iov_len), 0);
+        ASSERT_OK(iovec_split(&v, sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint16_t), &v));
+        ASSERT_EQ(v.iov_len, 3 * sizeof(uint16_t) + a.iov_len + b.iov_len + c.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, b.iov_len);
+        ASSERT_EQ(memcmp(p, b.iov_base, b.iov_len), 0);
+        p += b.iov_len;
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, c.iov_len);
+        ASSERT_EQ(memcmp(p, c.iov_base, c.iov_len), 0);
+        ASSERT_OK(iovec_split(&v, sizeof(uint16_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint32_t), &v));
+        ASSERT_EQ(v.iov_len, 3 * sizeof(uint32_t) + a.iov_len + b.iov_len + c.iov_len);
+        p = ASSERT_NOT_NULL(v.iov_base);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, a.iov_len);
+        ASSERT_EQ(memcmp(p, a.iov_base, a.iov_len), 0);
+        p += a.iov_len;
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, b.iov_len);
+        ASSERT_EQ(memcmp(p, b.iov_base, b.iov_len), 0);
+        p += b.iov_len;
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, 0);
+        ASSERT_EQ(*p++, c.iov_len);
+        ASSERT_EQ(memcmp(p, c.iov_base, c.iov_len), 0);
+        ASSERT_OK(iovec_split(&v, sizeof(uint32_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        /* with empty entries */
+        _cleanup_(iovw_done) struct iovec_wrapper with_empty = {
+                .iovec = ASSERT_PTR(new0(struct iovec, 6)),
+                .count = 6,
+        };
+        with_empty.iovec[0] = a;
+        with_empty.iovec[2] = b;
+        with_empty.iovec[4] = c;
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint8_t), &v));
+        ASSERT_OK(iovw_merge(&with_empty, sizeof(uint8_t), &v2));
+        ASSERT_TRUE(iovec_equal(&v, &v2));
+
+        iovec_done(&v);
+        iovec_done(&v2);
+
+        size_t sz = 6 + a.iov_len + b.iov_len + c.iov_len;
+        _cleanup_free_ uint8_t *buf = ASSERT_PTR(new(uint8_t, sz));
+        p = buf;
+        *p++ = a.iov_len;
+        p = mempcpy(p, a.iov_base, a.iov_len);
+        *p++ = 0;
+        *p++ = b.iov_len;
+        p = mempcpy(p, b.iov_base, b.iov_len);
+        *p++ = 0;
+        *p++ = c.iov_len;
+        p = mempcpy(p, c.iov_base, c.iov_len);
+        *p++ = 0;
+        ASSERT_OK(iovec_split(&IOVEC_MAKE(buf, sz), sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_equal(&iovw, &iovw2));
+
+        iovw_done_free(&iovw2);
+
+        /* truncated */
+        ASSERT_OK(iovw_merge(&iovw, sizeof(uint8_t), &v));
+        ASSERT_ERROR(iovec_split(&IOVEC_MAKE(v.iov_base, v.iov_len - 1), sizeof(uint8_t), &iovw2), EBADMSG);
+
+        iovec_done(&v);
+
+        /* too long */
+        _cleanup_(iovec_done) struct iovec large = {};
+        ASSERT_OK(random_bytes_allocate_iovec(256, &large));
+        ASSERT_ERROR(iovw_merge(&(struct iovec_wrapper) { .iovec = &large, .count = 1, }, sizeof(uint8_t), &v), ERANGE);
+        ASSERT_OK(iovw_merge(&(struct iovec_wrapper) { .iovec = &large, .count = 1, }, sizeof(uint16_t), &v));
+        ASSERT_OK(iovec_split(&v, sizeof(uint16_t), &iovw2));
+        ASSERT_EQ(iovw2.count, 1u);
+        ASSERT_TRUE(iovec_equal(&iovw2.iovec[0], &large));
+
+        iovec_done(&v);
+        iovw_done_free(&iovw2);
+
+        /* No entry */
+        ASSERT_OK(iovw_merge(&(struct iovec_wrapper) {}, sizeof(uint8_t), &v));
+        ASSERT_FALSE(iovec_is_set(&v));
+
+        ASSERT_OK(iovw_merge(NULL, sizeof(uint8_t), &v));
+        ASSERT_FALSE(iovec_is_set(&v));
+
+        ASSERT_OK(iovec_split(&(struct iovec) {}, sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_isempty(&iovw2));
+
+        ASSERT_OK(iovec_split(NULL, sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_isempty(&iovw2));
+
+        /* empty entry only */
+        ASSERT_OK(iovw_merge(&(struct iovec_wrapper) { .iovec = &(struct iovec) {}, .count = 1, }, sizeof(uint8_t), &v));
+        ASSERT_FALSE(iovec_is_set(&v));
+
+        ASSERT_OK(iovec_split(&IOVEC_MAKE("", 1), sizeof(uint8_t), &iovw2));
+        ASSERT_TRUE(iovw_isempty(&iovw2));
+
+}
+
 DEFINE_TEST_MAIN(LOG_INFO);