RFX: check for space in output buffer before writing tiles

5 files changed, 98 insertions, 77 deletions

diff --git a/src/rfxencode.c b/src/rfxencode.c
index a47a4a6..7822f14 100644
--- a/src/rfxencode.c
+++ b/src/rfxencode.c
@@ -314,6 +314,7 @@ rfxcodec_encode_ex(void *handle, char *cdata, int *cdata_bytes,
                    const char *quants, int num_quants, int flags)
 {
     struct rfxencode *enc;
+    int tiles_written;
     STREAM s;
 
     enc = (struct rfxencode *) handle;
@@ -327,18 +328,15 @@ rfxcodec_encode_ex(void *handle, char *cdata, int *cdata_bytes,
     {
         if (rfx_compose_message_header(enc, &s) != 0)
         {
-            return 1;
+            return -1;
         }
     }
-    if (rfx_compose_message_data(enc, &s, regions, num_regions,
-                                 buf, width, height, stride_bytes,
-                                 tiles, num_tiles, quants, num_quants,
-                                 flags) != 0)
-    {
-        return 1;
-    }
+    tiles_written = rfx_compose_message_data(enc, &s, regions, num_regions,
+                                            buf, width, height, stride_bytes,
+                                            tiles, num_tiles,
+                                            quants, num_quants, flags);
     *cdata_bytes = (int) (s.p - s.data);
-    return 0;
+    return tiles_written;
 }
 
 /******************************************************************************/
diff --git a/src/rfxencode_compose.c b/src/rfxencode_compose.c
index 3c0f87a..c9ecbd5 100644
--- a/src/rfxencode_compose.c
+++ b/src/rfxencode_compose.c
@@ -399,6 +399,8 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
     int size;
     int start_pos;
     int end_pos;
+    int tiles_end_checkpoint;
+    int tiles_written;
     int index;
     int numQuants;
     const char *quantVals;
@@ -449,6 +451,7 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
     memcpy(s->p, quantVals, numQuants * 5);
     s->p += numQuants * 5;
     end_pos = stream_get_pos(s);
+    tiles_written = 0;
     if (enc->format == RFX_FORMAT_YUV)
     {
         if (flags & RFX_FLAGS_ALPHAV1)
@@ -469,8 +472,10 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
                                                   quantIdxY, quantIdxCb, quantIdxCr,
                                                   x / 64, y / 64) != 0)
                 {
-                    return 1;
+                    break;
                 }
+                tiles_end_checkpoint = stream_get_pos(s);
+                tiles_written += 1;
             }
         }
         else
@@ -491,8 +496,10 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
                                                  quantIdxY, quantIdxCb, quantIdxCr,
                                                  x / 64, y / 64) != 0)
                 {
-                    return 1;
+                    break;
                 }
+                tiles_end_checkpoint = stream_get_pos(s);
+                tiles_written += 1;
             }
         }
     }
@@ -516,8 +523,10 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
                                                   quantIdxY, quantIdxCb, quantIdxCr,
                                                   x / 64, y / 64) != 0)
                 {
-                    return 1;
+                    break;
                 }
+                tiles_end_checkpoint = stream_get_pos(s);
+                tiles_written += 1;
             }
         }
         else
@@ -538,20 +547,24 @@ rfx_compose_message_tileset(struct rfxencode *enc, STREAM *s,
                                                  quantIdxY, quantIdxCb, quantIdxCr,
                                                  x / 64, y / 64) != 0)
                 {
-                    return 1;
+                    break;
                 }
+                tiles_end_checkpoint = stream_get_pos(s);
+                tiles_written += 1;
             }
         }
     }
-    tilesDataSize = stream_get_pos(s) - end_pos;
+    tilesDataSize = tiles_end_checkpoint - end_pos;
     size += tilesDataSize;
-    end_pos = stream_get_pos(s);
+    end_pos = tiles_end_checkpoint;
     stream_set_pos(s, start_pos + 2);
     stream_write_uint32(s, size); /* CodecChannelT.blockLen */
+    stream_set_pos(s, start_pos + 16);
+    stream_write_uint16(s, tiles_written);
     stream_set_pos(s, start_pos + 18);
     stream_write_uint32(s, tilesDataSize);
     stream_set_pos(s, end_pos);
-    return 0;
+    return tiles_written;
 }
 
 /******************************************************************************/
@@ -578,23 +591,20 @@ rfx_compose_message_data(struct rfxencode *enc, STREAM *s,
                          const struct rfx_tile *tiles, int num_tiles,
                          const char *quants, int num_quants, int flags)
 {
+    int tiles_written;
     if (rfx_compose_message_frame_begin(enc, s) != 0)
     {
-        return 1;
+        return -1;
     }
     if (rfx_compose_message_region(enc, s, regions, num_regions) != 0)
     {
-        return 1;
-    }
-    if (rfx_compose_message_tileset(enc, s, buf, width, height, stride_bytes,
-                                    tiles, num_tiles, quants, num_quants,
-                                    flags) != 0)
-    {
-        return 1;
+        return -1;
     }
+    tiles_written = rfx_compose_message_tileset(enc, s, buf, width, height, stride_bytes,
+                   tiles, num_tiles, quants, num_quants, flags);
     if (rfx_compose_message_frame_end(enc, s) != 0)
     {
-        return 1;
+        return -1;
     }
-    return 0;
+    return tiles_written;
 }
diff --git a/src/rfxencode_tile.c b/src/rfxencode_tile.c
index b75e314..aa2b91f 100644
--- a/src/rfxencode_tile.c
+++ b/src/rfxencode_tile.c
@@ -98,6 +98,19 @@ rfx_encode_component_rlgr3(struct rfxencode *enc, const char *qtable,
 }
 
 /******************************************************************************/
+static int
+check_and_rfx_encode(struct rfxencode *enc, const char *qtable,
+                     const uint8 *data,
+                     uint8 *buffer, int buffer_size, int *size)
+{
+    if (buffer_size < 4096 * 2)
+    {
+        return 1;
+    }
+    return enc->rfx_encode(enc, qtable, data, buffer, buffer_size, size);
+}
+
+/******************************************************************************/
 int
 rfx_encode_rgb(struct rfxencode *enc, const char *rgb_data,
                int width, int height, int stride_bytes,
@@ -118,28 +131,28 @@ rfx_encode_rgb(struct rfxencode *enc, const char *rgb_data,
     y_r_buffer = enc->y_r_buffer;
     u_g_buffer = enc->u_g_buffer;
     v_b_buffer = enc->v_b_buffer;
-    if (enc->rfx_encode(enc, y_quants, y_r_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        y_size) != 0)
+    if (check_and_rfx_encode(enc, y_quants, y_r_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             y_size) != 0)
     {
         return 1;
     }
     LLOGLN(10, ("rfx_encode_rgb: y_size %d", *y_size));
     stream_seek(data_out, *y_size);
-    if (enc->rfx_encode(enc, u_quants, u_g_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        u_size) != 0)
+    if (check_and_rfx_encode(enc, u_quants, u_g_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             u_size) != 0)
     {
         return 1;
     }
     LLOGLN(10, ("rfx_encode_rgb: u_size %d", *u_size));
     stream_seek(data_out, *u_size);
-    if (enc->rfx_encode(enc, v_quants, v_b_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        v_size) != 0)
+    if (check_and_rfx_encode(enc, v_quants, v_b_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             v_size) != 0)
     {
         return 1;
     }
@@ -172,28 +185,28 @@ rfx_encode_argb(struct rfxencode *enc, const char *argb_data,
     y_r_buffer = enc->y_r_buffer;
     u_g_buffer = enc->u_g_buffer;
     v_b_buffer = enc->v_b_buffer;
-    if (enc->rfx_encode(enc, y_quants, y_r_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        y_size) != 0)
+    if (check_and_rfx_encode(enc, y_quants, y_r_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             y_size) != 0)
     {
         return 1;
     }
     LLOGLN(10, ("rfx_encode_rgb: y_size %d", *y_size));
     stream_seek(data_out, *y_size);
-    if (enc->rfx_encode(enc, u_quants, u_g_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        u_size) != 0)
+    if (check_and_rfx_encode(enc, u_quants, u_g_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             u_size) != 0)
     {
         return 1;
     }
     LLOGLN(10, ("rfx_encode_rgb: u_size %d", *u_size));
     stream_seek(data_out, *u_size);
-    if (enc->rfx_encode(enc, v_quants, v_b_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        v_size) != 0)
+    if (check_and_rfx_encode(enc, v_quants, v_b_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             v_size) != 0)
     {
         return 1;
     }
@@ -218,26 +231,26 @@ rfx_encode_yuv(struct rfxencode *enc, const char *yuv_data,
     y_buffer = (const uint8 *) yuv_data;
     u_buffer = (const uint8 *) (yuv_data + RFX_YUV_BTES);
     v_buffer = (const uint8 *) (yuv_data + RFX_YUV_BTES * 2);
-    if (enc->rfx_encode(enc, y_quants, y_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        y_size) != 0)
+    if (check_and_rfx_encode(enc, y_quants, y_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             y_size) != 0)
     {
         return 1;
     }
     stream_seek(data_out, *y_size);
-    if (enc->rfx_encode(enc, u_quants, u_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        u_size) != 0)
+    if (check_and_rfx_encode(enc, u_quants, u_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             u_size) != 0)
     {
         return 1;
     }
     stream_seek(data_out, *u_size);
-    if (enc->rfx_encode(enc, v_quants, v_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        v_size) != 0)
+    if (check_and_rfx_encode(enc, v_quants, v_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             v_size) != 0)
     {
         return 1;
     }
@@ -263,26 +276,26 @@ rfx_encode_yuva(struct rfxencode *enc, const char *yuva_data,
     u_buffer = (const uint8 *) (yuva_data + RFX_YUV_BTES);
     v_buffer = (const uint8 *) (yuva_data + RFX_YUV_BTES * 2);
     a_buffer = (const uint8 *) (yuva_data + RFX_YUV_BTES * 3);
-    if (enc->rfx_encode(enc, y_quants, y_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        y_size) != 0)
+    if (check_and_rfx_encode(enc, y_quants, y_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             y_size) != 0)
     {
         return 1;
     }
     stream_seek(data_out, *y_size);
-    if (enc->rfx_encode(enc, u_quants, u_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        u_size) != 0)
+    if (check_and_rfx_encode(enc, u_quants, u_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             u_size) != 0)
     {
         return 1;
     }
     stream_seek(data_out, *u_size);
-    if (enc->rfx_encode(enc, v_quants, v_buffer,
-                        stream_get_tail(data_out),
-                        stream_get_left(data_out),
-                        v_size) != 0)
+    if (check_and_rfx_encode(enc, v_quants, v_buffer,
+                             stream_get_tail(data_out),
+                             stream_get_left(data_out),
+                             v_size) != 0)
     {
         return 1;
     }
diff --git a/tests/rfxcodectest.c b/tests/rfxcodectest.c
index 3b22dc7..9a5782e 100644
--- a/tests/rfxcodectest.c
+++ b/tests/rfxcodectest.c
@@ -125,7 +125,7 @@ speed_random(int count, const char *quants)
         error = rfxcodec_encode_ex(han, cdata, &cdata_bytes, buf, 64, 64, 64 * 4,
                                    regions, num_regions, tiles, num_tiles,
                                    quants, num_quants, flags);
-        if (error != 0)
+        if (error < num_tiles)
         {
             break;
         }
@@ -295,7 +295,7 @@ encode_file(char *data, int width, int height, char *cdata, int *cdata_bytes,
     error = rfxcodec_encode_ex(han, cdata, cdata_bytes, data, width, height, width * 4,
                                regions, num_regions, tiles, num_tiles,
                                quants, num_quants, 0);
-    if (error != 0)
+    if (error < num_tiles)
     {
         printf("encode_file: rfxcodec_encode failed error %d\n", error);
         return 1;
diff --git a/tests/rfxencode.c b/tests/rfxencode.c
index d8d1dfc..150b4ed 100644
--- a/tests/rfxencode.c
+++ b/tests/rfxencode.c
@@ -248,7 +248,7 @@ int process(void)
             error = rfxcodec_encode(han, out_data, &out_bytes, bmp_data,
                                     width, height, width * 4,
                                     &region, 1, tiles, num_tiles, NULL, 0);
-            if (error != 0)
+            if (error < num_tiles)
             {
                 break;
             }