avformat/webpenc: Write correct size for single images when unseekable

The earlier code writes the file and then tries to patch up
the size later. This is avoidable for the common case of
a single image because one can know the complete size
in advance and write it.

Fixes ticket #4609.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@outlook.com>

1 files changed, 41 insertions, 27 deletions

diff --git a/libavformat/webpenc.c b/libavformat/webpenc.c
index d4acea7ba6..ea7a321975 100644
--- a/libavformat/webpenc.c
+++ b/libavformat/webpenc.c
@@ -21,6 +21,7 @@
 
 #include "libavutil/intreadwrite.h"
 #include "libavutil/opt.h"
+#include "libavcodec/bytestream.h"
 #include "avformat.h"
 #include "internal.h"
 #include "mux.h"
@@ -76,11 +77,16 @@ static int is_animated_webp_packet(AVPacket *pkt)
     return 0;
 }
 
+/**
+ * Returns 1 if it has written a RIFF header with a correct length field
+ */
 static int flush(AVFormatContext *s, int trailer, int64_t pts)
 {
     WebpContext *w = s->priv_data;
     AVStream *st = s->streams[0];
-    int skip = 0;
+    uint8_t buf[12 /* RIFF+WEBP */ + 18 /* VP8X */ +
+                14 /* ANIM */ + 24 /* ANMF */], *bufp = buf;
+    int writing_webp_header = 0, skip = 0;
     unsigned flags = 0;
     int vp8x = 0;
 
@@ -97,7 +103,10 @@ static int flush(AVFormatContext *s, int trailer, int64_t pts)
     }
 
     if (!w->wrote_webp_header) {
-        avio_write(s->pb, "RIFF\0\0\0\0WEBP", 12);
+        bytestream_put_le32(&bufp, MKTAG('R', 'I', 'F', 'F'));
+        bytestream_put_le32(&bufp, 0); /* Size to be patched later */
+        bytestream_put_le32(&bufp, MKTAG('W', 'E', 'B', 'P'));
+        writing_webp_header  = 1;
         w->wrote_webp_header = 1;
         if (w->frame_count > 1)  // first non-empty packet
             w->frame_count = 1;  // so we don't count previous empty packets.
@@ -110,38 +119,41 @@ static int flush(AVFormatContext *s, int trailer, int64_t pts)
         }
 
         if (vp8x) {
-            avio_write(s->pb, "VP8X", 4);
-            avio_wl32(s->pb, 10);
-            avio_w8(s->pb, flags);
-            avio_wl24(s->pb, 0);
-            avio_wl24(s->pb, st->codecpar->width - 1);
-            avio_wl24(s->pb, st->codecpar->height - 1);
+            bytestream_put_le32(&bufp, MKTAG('V', 'P', '8', 'X'));
+            bytestream_put_le32(&bufp, 10);
+            bytestream_put_byte(&bufp, flags);
+            bytestream_put_le24(&bufp, 0);
+            bytestream_put_le24(&bufp, st->codecpar->width  - 1);
+            bytestream_put_le24(&bufp, st->codecpar->height - 1);
         }
         if (!trailer) {
-            avio_write(s->pb, "ANIM", 4);
-            avio_wl32(s->pb, 6);
-            avio_wl32(s->pb, 0xFFFFFFFF);
-            avio_wl16(s->pb, w->loop);
+            bytestream_put_le32(&bufp, MKTAG('A', 'N', 'I', 'M'));
+            bytestream_put_le32(&bufp, 6);
+            bytestream_put_le32(&bufp, 0xFFFFFFFF);
+            bytestream_put_le16(&bufp, w->loop);
         }
     }
 
     if (w->frame_count > trailer) {
-        avio_write(s->pb, "ANMF", 4);
-        avio_wl32(s->pb, 16 + w->last_pkt->size - skip);
-        avio_wl24(s->pb, 0);
-        avio_wl24(s->pb, 0);
-        avio_wl24(s->pb, st->codecpar->width - 1);
-        avio_wl24(s->pb, st->codecpar->height - 1);
+        bytestream_put_le32(&bufp, MKTAG('A', 'N', 'M', 'F'));
+        bytestream_put_le32(&bufp, 16 + w->last_pkt->size - skip);
+        bytestream_put_le24(&bufp, 0);
+        bytestream_put_le24(&bufp, 0);
+        bytestream_put_le24(&bufp, st->codecpar->width  - 1);
+        bytestream_put_le24(&bufp, st->codecpar->height - 1);
         if (w->last_pkt->pts != AV_NOPTS_VALUE && pts != AV_NOPTS_VALUE) {
-            avio_wl24(s->pb, pts - w->last_pkt->pts);
+            bytestream_put_le24(&bufp, pts - w->last_pkt->pts);
         } else
-            avio_wl24(s->pb, w->last_pkt->duration);
-        avio_w8(s->pb, 0);
+            bytestream_put_le24(&bufp, w->last_pkt->duration);
+        bytestream_put_byte(&bufp, 0);
     }
+    if (trailer && writing_webp_header)
+        AV_WL32(buf + 4, bufp - (buf + 8) + w->last_pkt->size - skip);
+    avio_write(s->pb, buf, bufp - buf);
     avio_write(s->pb, w->last_pkt->data + skip, w->last_pkt->size - skip);
     av_packet_unref(w->last_pkt);
 
-    return 0;
+    return trailer && writing_webp_header;
 }
 
 static int webp_write_packet(AVFormatContext *s, AVPacket *pkt)
@@ -185,11 +197,13 @@ static int webp_write_trailer(AVFormatContext *s)
         if ((ret = flush(s, 1, AV_NOPTS_VALUE)) < 0)
             return ret;
 
-        filesize = avio_tell(s->pb);
-        if (avio_seek(s->pb, 4, SEEK_SET) == 4) {
-            avio_wl32(s->pb, filesize - 8);
-            // Note: without the following, avio only writes 8 bytes to the file.
-            avio_seek(s->pb, filesize, SEEK_SET);
+        if (!ret) {
+            filesize = avio_tell(s->pb);
+            if (avio_seek(s->pb, 4, SEEK_SET) == 4) {
+                avio_wl32(s->pb, filesize - 8);
+                // Note: without the following, avio only writes 8 bytes to the file.
+                avio_seek(s->pb, filesize, SEEK_SET);
+            }
         }
     }