rtmp: Read and handle incoming packets while writing data

This makes sure all incoming packets are read and handled (and reacted
to) while sending an FLV stream over RTMP to a server. If there were
enough incoming data to fill the TCP buffers, this could potentially
make things block at unexpected places. For the upcoming RTMPT support,
we need to consume all incoming data before we can send the next
request.

Signed-off-by: Martin Storsjö <martin@martin.st>

3 files changed, 56 insertions, 3 deletions

diff --git a/libavformat/rtmppkt.c b/libavformat/rtmppkt.c
index ed8e6b203d..4ce238d5d0 100644
--- a/libavformat/rtmppkt.c
+++ b/libavformat/rtmppkt.c
@@ -74,15 +74,25 @@ void ff_amf_write_object_end(uint8_t **dst)
 int ff_rtmp_packet_read(URLContext *h, RTMPPacket *p,
                         int chunk_size, RTMPPacket *prev_pkt)
 {
-    uint8_t hdr, t, buf[16];
+    uint8_t hdr;
+
+    if (ffurl_read(h, &hdr, 1) != 1)
+        return AVERROR(EIO);
+
+    return ff_rtmp_packet_read_internal(h, p, chunk_size, prev_pkt, hdr);
+}
+
+int ff_rtmp_packet_read_internal(URLContext *h, RTMPPacket *p, int chunk_size,
+                                 RTMPPacket *prev_pkt, uint8_t hdr)
+{
+
+    uint8_t t, buf[16];
     int channel_id, timestamp, data_size, offset = 0;
     uint32_t extra = 0;
     enum RTMPPacketType type;
     int size = 0;
     int ret;
 
-    if (ffurl_read(h, &hdr, 1) != 1)
-        return AVERROR(EIO);
     size++;
     channel_id = hdr & 0x3F;
 
diff --git a/libavformat/rtmppkt.h b/libavformat/rtmppkt.h
index 8372484fbd..a83d0feb8f 100644
--- a/libavformat/rtmppkt.h
+++ b/libavformat/rtmppkt.h
@@ -115,6 +115,19 @@ void ff_rtmp_packet_destroy(RTMPPacket *pkt);
  */
 int ff_rtmp_packet_read(URLContext *h, RTMPPacket *p,
                         int chunk_size, RTMPPacket *prev_pkt);
+/**
+ * Read internal RTMP packet sent by the server.
+ *
+ * @param h          reader context
+ * @param p          packet
+ * @param chunk_size current chunk size
+ * @param prev_pkt   previously read packet headers for all channels
+ *                   (may be needed for restoring incomplete packet header)
+ * @param c          the first byte already read
+ * @return number of bytes read on success, negative value otherwise
+ */
+int ff_rtmp_packet_read_internal(URLContext *h, RTMPPacket *p, int chunk_size,
+                                 RTMPPacket *prev_pkt, uint8_t c);
 
 /**
  * Send RTMP packet to the server.
diff --git a/libavformat/rtmpproto.c b/libavformat/rtmpproto.c
index e64e2a322b..b3ae5a21e6 100644
--- a/libavformat/rtmpproto.c
+++ b/libavformat/rtmpproto.c
@@ -1287,6 +1287,7 @@ static int rtmp_write(URLContext *s, const uint8_t *buf, int size)
     int pktsize, pkttype;
     uint32_t ts;
     const uint8_t *buf_temp = buf;
+    uint8_t c;
     int ret;
 
     do {
@@ -1356,6 +1357,35 @@ static int rtmp_write(URLContext *s, const uint8_t *buf, int size)
             rt->flv_header_bytes = 0;
         }
     } while (buf_temp - buf < size);
+
+    /* set stream into nonblocking mode */
+    rt->stream->flags |= AVIO_FLAG_NONBLOCK;
+
+    /* try to read one byte from the stream */
+    ret = ffurl_read(rt->stream, &c, 1);
+
+    /* switch the stream back into blocking mode */
+    rt->stream->flags &= ~AVIO_FLAG_NONBLOCK;
+
+    if (ret == AVERROR(EAGAIN)) {
+        /* no incoming data to handle */
+        return size;
+    } else if (ret < 0) {
+        return ret;
+    } else if (ret == 1) {
+        RTMPPacket rpkt = { 0 };
+
+        if ((ret = ff_rtmp_packet_read_internal(rt->stream, &rpkt,
+                                                rt->chunk_size,
+                                                rt->prev_pkt[0], c)) <= 0)
+             return ret;
+
+        if ((ret = rtmp_parse_result(s, rt, &rpkt)) < 0)
+            return ret;
+
+        ff_rtmp_packet_destroy(&rpkt);
+    }
+
     return size;
 }