From dcae9c19bb412d726751154ad8c94aa4753c6d37 Mon Sep 17 00:00:00 2001
From: Jerzy Kasenberg <jerzy.kasenberg@codecoup.pl>
Date: Fri, 24 May 2024 14:11:20 +0200
Subject: [PATCH] util/stream: Add pumping functions

This extends in_stream and out_stream to provide pumping
functions

Additionally stream_pump() function is added to allow easy
pumping data between streams.

Signed-off-by: Jerzy Kasenberg <jerzy.kasenberg@codecoup.pl>
---
 util/stream/include/stream/stream.h | 52 +++++++++++++++++++++++++++--
 util/stream/src/stream.c            | 29 ++++++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/util/stream/include/stream/stream.h b/util/stream/include/stream/stream.h
index 209c660ce1..61137951f7 100644
--- a/util/stream/include/stream/stream.h
+++ b/util/stream/include/stream/stream.h
@@ -61,6 +61,17 @@ struct in_stream_vft {
      *         negative error code
      */
     int (*available)(struct in_stream *istream);
+
+    /**
+     * Pump data from this in_stream to other out_stream
+     *
+     * @param istream - this input stream to get data from
+     * @param ostream - output stream to put data to
+     * @param count - number of bytes to pump
+     * @return non-negative number of bytes written
+     *         negative error code
+     */
+    int (*pump_to)(struct in_stream *istream, struct out_stream *ostream, uint32_t count);
 };
 
 /* Output stream functions */
@@ -83,6 +94,17 @@ struct out_stream_vft {
      *         negative error code
      */
     int (*flush)(struct out_stream *ostream);
+
+    /**
+     * Pump data from in_stream to this out_stream
+     *
+     * @param ostream - this output stream to write data to
+     * @param istream - other input stream to get data from
+     * @param count - number for bytes to pump
+     * @return non-negative number of bytes written
+     *         negative error code
+     */
+    int (*pump_from)(struct out_stream *ostream, struct in_stream *istream, uint32_t count);
 };
 
 /* Plain input stream */
@@ -102,15 +124,39 @@ struct mem_in_stream {
     uint32_t read_ptr;
 };
 
+#define OSTREAM_DEF(type) \
+    static int type ## _write(struct out_stream *ostream, const uint8_t *buf, uint32_t count); \
+    static int type ## _flush(struct out_stream *ostream); \
+    static int type ## _pump_from(struct out_stream *ostream, struct in_stream *istream, uint32_t count); \
+    const struct out_stream_vft type ## _vft = { \
+        .write = type ## _write, \
+        .flush = type ## _flush, \
+        .pump_from = type ## _pump_from, \
+    }
+
+#define OSTREAM_VFT(type, _write, _flush, _pump) \
+    static int _write(struct out_stream *ostream, const uint8_t *buf, uint32_t count); \
+    static int _flush(struct out_stream *ostream); \
+    static int _pump_from(struct out_stream *ostream, struct in_stream *istream, uint32_t count); \
+    const struct out_stream_vft type ## _vft = { \
+        .write = _write, \
+        .flush = _flush, \
+        .pump_from = _pump_from, \
+    }
+
 #define OSTREAM(type, name)     \
     struct type name = {        \
         .vft = &type ## _vft    \
     }
 
+#define OSTREAM_INIT(type, field)     \
+    .field.vft = &type ## _vft
+
 #define ISTREAM_DEF(type)                       \
     const struct in_stream_vft type ## _vft = { \
         .available = type ## _available,        \
         .read = type ## _read,                  \
+        .pump_to = type ## _pump_to,            \
     }
 
 #define ISTREAM(type, name)     \
@@ -118,8 +164,8 @@ struct mem_in_stream {
         .vft = &type ## _vft    \
     }
 
-#define ISTREAM_INIT(type)     \
-    .type.vft = &type ## _vft  \
+#define ISTREAM_INIT(type, field)     \
+    .field.vft = &type ## _vft  \
 
 /**
  * Check how many bytes can be read out of stream.
@@ -175,6 +221,8 @@ int ostream_flush(struct out_stream *ostream);
  */
 int ostream_write(struct out_stream *ostream, const uint8_t *buf, uint32_t count, bool flush);
 
+int stream_pump(struct in_stream *istream, struct out_stream *ostream, uint32_t count);
+
 /**
  * Initialize memory input stream
  *
diff --git a/util/stream/src/stream.c b/util/stream/src/stream.c
index e0ce166fec..fe76051579 100644
--- a/util/stream/src/stream.c
+++ b/util/stream/src/stream.c
@@ -20,6 +20,7 @@
 #include <string.h>
 
 #include <stream/stream.h>
+#include "os/os.h"
 
 int
 istream_flush(struct in_stream *istream)
@@ -107,3 +108,31 @@ istream_read(struct in_stream *istream, uint8_t *buf, uint32_t count)
     return istream->vft->read(istream, buf, count);
 }
 
+int
+stream_pump(struct in_stream *istream, struct out_stream *ostream, uint32_t count)
+{
+    int pumped;
+
+    /* If output stream has pump_from function used it */
+    if (ostream->vft->pump_from) {
+        pumped = ostream->vft->pump_from(ostream, istream, count);
+    } else if (istream->vft->pump_to) {
+        /* When input stream has pump_to used it */
+        pumped = istream->vft->pump_to(istream, ostream, count);
+    } else {
+        /* Both stream don't provide pump functions, use local buffer for pumping data */
+        uint8_t buf[16];
+        pumped = 0;
+        while (count) {
+            uint32_t n = min(count, sizeof(buf));
+            uint32_t r = istream_read(istream, buf, n);
+            pumped += r;
+            count -= r;
+            ostream_write(ostream, buf, r, false);
+            if (r == 0) {
+                break;
+            }
+        }
+    }
+    return pumped;
+}