[FFmpeg-devel] [PATCH 3/3] avformat/webpenc: Write correct size for single images when unseekable

[Andreas Rheinhardt]

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 at outlook.com>
---
 libavformat/webpenc.c | 68 ++++++++++++++++++++++++++-----------------
 1 file 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);
+            }
         }
     }
 
-- 
2.34.1