Merge pull request #10 from widgetii/locks_n_fill

Fix race conditions in TCP streams and potential internal buffers overrun
OpenIPC · May 13, 2023 · 750fa54 · 750fa54
2 parents 4a73d94 + 91a5b91
commit 750fa54
Show file tree

Hide file tree

Showing 11 changed files with 104 additions and 4 deletions.
diff --git a/include/smolrtsp/nal_transport.h b/include/smolrtsp/nal_transport.h
@@ -110,3 +110,6 @@ int SmolRTSP_NalTransport_send_packet(
  * usage.
  */
 declImplExtern99(SmolRTSP_Droppable, SmolRTSP_NalTransport);
+
+bool
+SmolRTSP_NalTransport_is_full(SmolRTSP_NalTransport *self);
diff --git a/include/smolrtsp/rtp_transport.h b/include/smolrtsp/rtp_transport.h
@@ -83,3 +83,6 @@ int SmolRTSP_RtpTransport_send_packet(
  * usage.
  */
 declImplExtern99(SmolRTSP_Droppable, SmolRTSP_RtpTransport);
+
+bool
+SmolRTSP_RtpTransport_is_full(SmolRTSP_RtpTransport *self);
diff --git a/include/smolrtsp/transport.h b/include/smolrtsp/transport.h
@@ -27,7 +27,8 @@
      * @return -1 if an I/O error occurred and sets `errno` appropriately, 0   \
      * on success.                                                             \
      */                                                                        \
-    vfunc99(int, transmit, VSelf99, SmolRTSP_IoVecSlice bufs)
+    vfunc99(int, transmit, VSelf99, SmolRTSP_IoVecSlice bufs)                  \
+        vfunc99(bool, is_full, VSelf99)
 
 /**
  * The superinterfaces of #SmolRTSP_Transport_IFACE.
@@ -51,8 +52,9 @@ interface99(SmolRTSP_Transport);
  *
  * @pre `w.self && w.vptr`
  */
-SmolRTSP_Transport smolrtsp_transport_tcp(SmolRTSP_Writer w, uint8_t channel_id)
-    SMOLRTSP_PRIV_MUST_USE;
+SmolRTSP_Transport smolrtsp_transport_tcp(
+    SmolRTSP_Writer w, uint8_t channel_id,
+    size_t max_buffer) SMOLRTSP_PRIV_MUST_USE;
 
 /**
  * Creates a new UDP transport.

diff --git a/include/smolrtsp/writer.h b/include/smolrtsp/writer.h
@@ -30,6 +30,21 @@
      */                                                                        \
     vfunc99(ssize_t, write, VSelf99, CharSlice99 data)                         \
                                                                                \
+    /*                                                                         \
+     * Lock writer to prevent race conditions on TCP interleaved channels      \
+     */                                                                        \
+    vfunc99(void, lock, VSelf99)                                               \
+                                                                               \
+    /*                                                                         \
+     * Unlock writer locked by `lock`                                          \
+     */                                                                        \
+    vfunc99(void, unlock, VSelf99)                                             \
+                                                                               \
+    /*                                                                         \
+     * Get current size of output buffer                                       \
+     */                                                                        \
+    vfunc99(size_t, filled, VSelf99)                                           \
+                                                                               \
     /*                                                                         \
      * Writes a formatted string into itself.                                  \
      *                                                                         \

diff --git a/src/nal_transport.c b/src/nal_transport.c
@@ -58,6 +58,11 @@ static void SmolRTSP_NalTransport_drop(VSelf) {
 
 implExtern(SmolRTSP_Droppable, SmolRTSP_NalTransport);
 
+bool
+SmolRTSP_NalTransport_is_full(SmolRTSP_NalTransport *self) {
+	return SmolRTSP_RtpTransport_is_full(self->transport);
+}
+
 int SmolRTSP_NalTransport_send_packet(
     SmolRTSP_NalTransport *self, SmolRTSP_RtpTimestamp ts,
     SmolRTSP_NalUnit nalu) {

diff --git a/src/rtp_transport.c b/src/rtp_transport.c
@@ -104,3 +104,7 @@ compute_timestamp(SmolRTSP_RtpTimestamp ts, uint32_t clock_rate) {
 
     return 0;
 }
+
+bool SmolRTSP_RtpTransport_is_full(SmolRTSP_RtpTransport *self) {
+    return VCALL(self->transport, is_full);
+}
diff --git a/src/transport/tcp.c b/src/transport/tcp.c
@@ -14,19 +14,21 @@
 typedef struct {
     SmolRTSP_Writer w;
     int channel_id;
+    size_t max_buffer;
 } SmolRTSP_TcpTransport;
 
 declImpl(SmolRTSP_Transport, SmolRTSP_TcpTransport);
 
 SmolRTSP_Transport
-smolrtsp_transport_tcp(SmolRTSP_Writer w, uint8_t channel_id) {
+smolrtsp_transport_tcp(SmolRTSP_Writer w, uint8_t channel_id, size_t max_buffer) {
     assert(w.self && w.vptr);
 
     SmolRTSP_TcpTransport *self = malloc(sizeof *self);
     assert(self);
 
     self->w = w;
     self->channel_id = channel_id;
+    self->max_buffer = max_buffer;
 
     return DYN(SmolRTSP_TcpTransport, SmolRTSP_Transport, self);
 }
@@ -49,9 +51,11 @@ static int SmolRTSP_TcpTransport_transmit(VSelf, SmolRTSP_IoVecSlice bufs) {
     const uint32_t header =
         smolrtsp_interleaved_header(self->channel_id, htons(total_bytes));
 
+    VCALL(self->w, lock);
     ssize_t ret =
         VCALL(self->w, write, CharSlice99_new((char *)&header, sizeof header));
     if (ret != sizeof header) {
+        VCALL(self->w, unlock);
         return -1;
     }
 
@@ -60,11 +64,20 @@ static int SmolRTSP_TcpTransport_transmit(VSelf, SmolRTSP_IoVecSlice bufs) {
             CharSlice99_new(bufs.ptr[i].iov_base, bufs.ptr[i].iov_len);
         ret = VCALL(self->w, write, vec);
         if (ret != (ssize_t)vec.len) {
+            VCALL(self->w, unlock);
             return -1;
         }
     }
+    VCALL(self->w, unlock);
 
     return 0;
 }
 
+static bool SmolRTSP_TcpTransport_is_full(VSelf) {
+    VSELF(SmolRTSP_TcpTransport);
+    assert(self);
+
+    return VCALL(self->w, filled) > self->max_buffer;
+}
+
 impl(SmolRTSP_Transport, SmolRTSP_TcpTransport);
diff --git a/src/transport/udp.c b/src/transport/udp.c
@@ -58,6 +58,13 @@ static int SmolRTSP_UdpTransport_transmit(VSelf, SmolRTSP_IoVecSlice bufs) {
     return send_packet(self, msg);
 }
 
+static bool SmolRTSP_UdpTransport_is_full(VSelf) {
+    VSELF(SmolRTSP_UdpTransport);
+    (void)self;
+
+    return false;
+}
+
 impl(SmolRTSP_Transport, SmolRTSP_UdpTransport);
 
 static int send_packet(SmolRTSP_UdpTransport *self, struct msghdr message) {

diff --git a/src/writer/fd.c b/src/writer/fd.c
@@ -11,6 +11,22 @@ static ssize_t FdWriter_write(VSelf, CharSlice99 data) {
     return write(*self, data.ptr, data.len);
 }
 
+static void FdWriter_lock(VSelf) {
+    VSELF(FdWriter);
+    (void)self;
+}
+
+static void FdWriter_unlock(VSelf) {
+    VSELF(FdWriter);
+    (void)self;
+}
+
+static size_t FdWriter_filled(VSelf) {
+    VSELF(FdWriter);
+    (void)self;
+    return 0;
+}
+
 static int FdWriter_vwritef(VSelf, const char *restrict fmt, va_list ap) {
     VSELF(FdWriter);
 

diff --git a/src/writer/file.c b/src/writer/file.c
@@ -17,6 +17,22 @@ static ssize_t FileWriter_write(VSelf, CharSlice99 data) {
     return ret;
 }
 
+static void FileWriter_lock(VSelf) {
+    VSELF(FileWriter);
+    (void)self;
+}
+
+static void FileWriter_unlock(VSelf) {
+    VSELF(FileWriter);
+    (void)self;
+}
+
+static size_t FileWriter_filled(VSelf) {
+    VSELF(FileWriter);
+    (void)self;
+    return 0;
+}
+
 static int FileWriter_vwritef(VSelf, const char *restrict fmt, va_list ap) {
     VSELF(FileWriter);
 

diff --git a/src/writer/string.c b/src/writer/string.c
@@ -13,6 +13,22 @@ static ssize_t StringWriter_write(VSelf, CharSlice99 data) {
     return data.len;
 }
 
+static void StringWriter_lock(VSelf) {
+    VSELF(StringWriter);
+    (void)self;
+}
+
+static void StringWriter_unlock(VSelf) {
+    VSELF(StringWriter);
+    (void)self;
+}
+
+static size_t StringWriter_filled(VSelf) {
+    VSELF(StringWriter);
+    (void)self;
+    return 0;
+}
+
 static int StringWriter_vwritef(VSelf, const char *restrict fmt, va_list ap) {
     VSELF(StringWriter);