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;
 }