avcodec/elbg: fix integer overflows

Fixes #9977

1 files changed, 36 insertions, 20 deletions

diff --git a/libavcodec/elbg.c b/libavcodec/elbg.c
index d97a7bc3f9..7a6a84fb6b 100644
--- a/libavcodec/elbg.c
+++ b/libavcodec/elbg.c
@@ -44,13 +44,13 @@ typedef struct cell_s {
  * ELBG internal data
  */
 typedef struct ELBGContext {
-    int64_t error;
+    int error;
     int dim;
     int num_cb;
     int *codebook;
     cell **cells;
-    int64_t *utility;
-    int64_t *utility_inc;
+    int *utility;
+    int *utility_inc;
     int *nearest_cb;
     int *points;
     int *temp_points;
@@ -75,9 +75,12 @@ static inline int distance_limited(int *a, int *b, int dim, int limit)
 {
     int i, dist=0;
     for (i=0; i<dim; i++) {
-        dist += (a[i] - b[i])*(a[i] - b[i]);
-        if (dist > limit)
-            return INT_MAX;
+        int64_t distance = a[i] - b[i];
+
+        distance *= distance;
+        if (dist >= limit - distance)
+            return limit;
+        dist += distance;
     }
 
     return dist;
@@ -97,8 +100,12 @@ static inline void vect_division(int *res, int *vect, int div, int dim)
 static int eval_error_cell(ELBGContext *elbg, int *centroid, cell *cells)
 {
     int error=0;
-    for (; cells; cells=cells->next)
-        error += distance_limited(centroid, elbg->points + cells->index*elbg->dim, elbg->dim, INT_MAX);
+    for (; cells; cells=cells->next) {
+        int distance = distance_limited(centroid, elbg->points + cells->index*elbg->dim, elbg->dim, INT_MAX);
+        if (error >= INT_MAX - distance)
+            return INT_MAX;
+        error += distance;
+    }
 
     return error;
 }
@@ -178,10 +185,13 @@ static int simple_lbg(ELBGContext *elbg,
         int dist[2] = {distance_limited(centroid[0], points + tempcell->index*dim, dim, INT_MAX),
                        distance_limited(centroid[1], points + tempcell->index*dim, dim, INT_MAX)};
         int idx = dist[0] > dist[1];
-        newutility[idx] += dist[idx];
+        if (newutility[idx] >= INT_MAX - dist[idx])
+            newutility[idx] = INT_MAX;
+        else
+            newutility[idx] += dist[idx];
     }
 
-    return newutility[0] + newutility[1];
+    return (newutility[0] >= INT_MAX - newutility[1]) ? INT_MAX : newutility[0] + newutility[1];
 }
 
 static void get_new_centroids(ELBGContext *elbg, int huc, int *newcentroid_i,
@@ -253,9 +263,9 @@ static void evaluate_utility_inc(ELBGContext *elbg)
     int64_t inc=0;
 
     for (int i = 0; i < elbg->num_cb; i++) {
-        if (elbg->num_cb * elbg->utility[i] > elbg->error)
+        if (elbg->num_cb * (int64_t)elbg->utility[i] > elbg->error)
             inc += elbg->utility[i];
-        elbg->utility_inc[i] = inc;
+        elbg->utility_inc[i] = FFMIN(inc, INT_MAX);
     }
 }
 
@@ -278,7 +288,7 @@ static void update_utility_and_n_cb(ELBGContext *elbg, int idx, int newutility)
  */
 static void try_shift_candidate(ELBGContext *elbg, int idx[3])
 {
-    int j, k, cont=0;
+    int j, k, cont=0, tmp;
     int64_t olderror=0, newerror;
     int newutility[3];
     int *newcentroid[3] = {
@@ -305,12 +315,17 @@ static void try_shift_candidate(ELBGContext *elbg, int idx[3])
     get_new_centroids(elbg, idx[1], newcentroid[0], newcentroid[1]);
 
     newutility[2]  = eval_error_cell(elbg, newcentroid[2], elbg->cells[idx[0]]);
-    newutility[2] += eval_error_cell(elbg, newcentroid[2], elbg->cells[idx[2]]);
+    tmp            = eval_error_cell(elbg, newcentroid[2], elbg->cells[idx[2]]);
+    newutility[2]  = (tmp >= INT_MAX - newutility[2]) ? INT_MAX : newutility[2] + tmp;
 
     newerror = newutility[2];
 
-    newerror += simple_lbg(elbg, elbg->dim, newcentroid, newutility, elbg->points,
+    tmp = simple_lbg(elbg, elbg->dim, newcentroid, newutility, elbg->points,
                            elbg->cells[idx[1]]);
+    if (tmp >= INT_MAX - newerror)
+        newerror = INT_MAX;
+    else
+        newerror += tmp;
 
     if (olderror > newerror) {
         shift_codebook(elbg, idx, newcentroid);
@@ -334,7 +349,7 @@ static void do_shiftings(ELBGContext *elbg)
     evaluate_utility_inc(elbg);
 
     for (idx[0]=0; idx[0] < elbg->num_cb; idx[0]++)
-        if (elbg->num_cb * elbg->utility[idx[0]] < elbg->error) {
+        if (elbg->num_cb * (int64_t)elbg->utility[idx[0]] < elbg->error) {
             if (elbg->utility_inc[elbg->num_cb - 1] == 0)
                 return;
 
@@ -352,9 +367,9 @@ static void do_elbg(ELBGContext *av_restrict elbg, int *points, int numpoints,
     int *const size_part = elbg->size_part;
     int i, j, steps = 0;
     int best_idx = 0;
-    int64_t last_error;
+    int last_error;
 
-    elbg->error = INT64_MAX;
+    elbg->error = INT_MAX;
     elbg->points = points;
 
     do {
@@ -382,8 +397,9 @@ static void do_elbg(ELBGContext *av_restrict elbg, int *points, int numpoints,
                 }
             }
             elbg->nearest_cb[i] = best_idx;
-            elbg->error += best_dist;
-            elbg->utility[elbg->nearest_cb[i]] += best_dist;
+            elbg->error = (elbg->error >= INT_MAX - best_dist) ? INT_MAX : elbg->error + best_dist;
+            elbg->utility[elbg->nearest_cb[i]] = (elbg->utility[elbg->nearest_cb[i]] >= INT_MAX - best_dist) ?
+                                                  INT_MAX : elbg->utility[elbg->nearest_cb[i]] + best_dist;
             free_cells->index = i;
             free_cells->next = elbg->cells[elbg->nearest_cb[i]];
             elbg->cells[elbg->nearest_cb[i]] = free_cells;