Add generic proto writing method

4 files changed, 68 insertions, 31 deletions

diff --git a/encode.c b/encode.c
index 9ee7ddd..5975efc 100644
--- a/encode.c
+++ b/encode.c
@@ -23,7 +23,6 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <sys/uio.h>
 #include <unistd.h>
 #include <va/va.h>
 #include <va/va_drm.h>
@@ -761,24 +760,6 @@ static void UpdatePicHeader(struct EncodeContext* encode_context, bool idr) {
   }
 }
 
-static bool DrainBuffers(int fd, struct iovec* iovec, int count) {
-  for (;;) {
-    ssize_t result = writev(fd, iovec, count);
-    if (result < 0) {
-      if (errno == EINTR) continue;
-      LOG("Failed to write (%s)", strerror(errno));
-      return false;
-    }
-    for (int i = 0; i < count; i++) {
-      size_t delta = MIN((size_t)result, iovec[i].iov_len);
-      iovec[i].iov_base = (uint8_t*)iovec[i].iov_base + delta;
-      iovec[i].iov_len -= delta;
-      result -= delta;
-    }
-    if (!result) return true;
-  }
-}
-
 bool EncodeContextEncodeFrame(struct EncodeContext* encode_context, int fd,
                               unsigned long long timestamp) {
   bool result = false;
@@ -931,12 +912,8 @@ bool EncodeContextEncodeFrame(struct EncodeContext* encode_context, int fd,
       .flags = idr ? PROTO_FLAG_KEYFRAME : 0,
       .latency = (uint16_t)(MicrosNow() - timestamp),
   };
-  struct iovec iovec[] = {
-      {.iov_base = &proto, .iov_len = sizeof(proto)},
-      {.iov_base = segment->buf, .iov_len = segment->size},
-  };
-  if (!DrainBuffers(fd, iovec, LENGTH(iovec))) {
-    LOG("Failed to drain encoded frame");
+  if (!WriteProto(fd, &proto, segment->buf)) {
+    LOG("Failed to write encoded frame");
     goto rollback_segment;
   }
 
diff --git a/input.c b/input.c
index e47f363..5a279cc 100644
--- a/input.c
+++ b/input.c
@@ -100,12 +100,11 @@ bool InputHandlerHandle(struct InputHandler* input_handler, int fd) {
         // mburakov: Payload of ping message is not yet available.
         return true;
       }
-      char buffer[sizeof(struct Proto) + sizeof(uint64_t)];
-      struct Proto* proto = (void*)buffer;
-      proto->size = sizeof(uint64_t);
-      proto->type = PROTO_TYPE_MISC;
-      memcpy(proto->data, &event->u, sizeof(uint64_t));
-      if (write(fd, buffer, sizeof(buffer)) != sizeof(buffer)) {
+      struct Proto proto = {
+          .size = sizeof(uint64_t),
+          .type = PROTO_TYPE_MISC,
+      };
+      if (!WriteProto(fd, &proto, &event->u)) {
         LOG("Failed to write pong message (%s)", strerror(errno));
         return false;
       }
diff --git a/proto.c b/proto.c
new file mode 100644
index 0000000..f1d2bc3
--- /dev/null
+++ b/proto.c
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2023 Mikhail Burakov. This file is part of streamer.
+ *
+ * streamer is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * streamer is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with streamer.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include "proto.h"
+
+#include <errno.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/uio.h>
+#include <unistd.h>
+
+#include "toolbox/utils.h"
+
+#define UNCONST(x) ((void*)(uintptr_t)(x))
+
+static bool DrainBuffers(int fd, struct iovec* iovec, int count) {
+  for (;;) {
+    ssize_t result = writev(fd, iovec, count);
+    if (result < 0) {
+      if (errno == EINTR) continue;
+      LOG("Failed to write (%s)", strerror(errno));
+      return false;
+    }
+    for (int i = 0; i < count; i++) {
+      size_t delta = MIN((size_t)result, iovec[i].iov_len);
+      iovec[i].iov_base = (uint8_t*)iovec[i].iov_base + delta;
+      iovec[i].iov_len -= delta;
+      result -= delta;
+    }
+    if (!result) return true;
+  }
+}
+
+bool WriteProto(int fd, const struct Proto* proto, const void* data) {
+  struct iovec iovec[] = {
+      {.iov_base = UNCONST(proto), .iov_len = sizeof(struct Proto)},
+      {.iov_base = UNCONST(data), .iov_len = proto->size},
+  };
+  if (!DrainBuffers(fd, iovec, LENGTH(iovec))) {
+    LOG("Failed to drain buffers");
+    return false;
+  }
+  return true;
+}
diff --git a/proto.h b/proto.h
index c8c4d0b..62f65d5 100644
--- a/proto.h
+++ b/proto.h
@@ -20,6 +20,7 @@
 
 #include <assert.h>
 #include <stdint.h>
+#include <stdbool.h>
 
 #define PROTO_TYPE_MISC 0
 #define PROTO_TYPE_VIDEO 1
@@ -38,4 +39,6 @@ struct Proto {
 static_assert(sizeof(struct Proto) == 8 * sizeof(uint8_t),
               "Suspicious proto struct size");
 
+bool WriteProto(int fd, const struct Proto* proto, const void* data);
+
 #endif  // STREAMER_PROTO_H_