This is an automated email from the ASF dual-hosted git repository.

yjhjstz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/cloudberry.git

commit 546d447ad2cd0cc6de2ecba091a258399e7bddc3
Author: Hongxu Ma <[email protected]>
AuthorDate: Sun Oct 8 09:43:49 2023 +0800

    Fix ERROR: "Cannot add cell to table content: total cell count of XXX 
exceeded." (#16388)
    
    When user try to select a big number of rows via PSQL, they may encounter 
the following error:
    ERROR: "Cannot add cell to table content: total cell count of XXX exceeded.
    
    The root cause is the overflow of int type in printTableAddCell().
    Fix it in this commit.
    
    (And this issue also exists in upstream, I have reported it:
    
https://www.postgresql.org/message-id/flat/TYBP286MB0351B057B101C90D7C1239E6B4E2A%40TYBP286MB0351.JPNP286.PROD.OUTLOOK.COM)
---
 src/fe_utils/Makefile          |  3 +++
 src/fe_utils/print.c           | 23 +++++++++++------
 src/fe_utils/test/Makefile     | 14 ++++++++++
 src/fe_utils/test/print_test.c | 58 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 90 insertions(+), 8 deletions(-)

diff --git a/src/fe_utils/Makefile b/src/fe_utils/Makefile
index 604a3d87d4..3a33d43ced 100644
--- a/src/fe_utils/Makefile
+++ b/src/fe_utils/Makefile
@@ -64,3 +64,6 @@ clean distclean:
 # so do not clean it in the clean/distclean rules
 maintainer-clean: distclean
        rm -f psqlscan.c
+
+unittest-check:
+       $(MAKE) -C test check
diff --git a/src/fe_utils/print.c b/src/fe_utils/print.c
index 2d0f78b8a2..95168869cb 100644
--- a/src/fe_utils/print.c
+++ b/src/fe_utils/print.c
@@ -3058,6 +3058,8 @@ void
 printTableInit(printTableContent *const content, const printTableOpt *opt,
                           const char *title, const int ncolumns, const int 
nrows)
 {
+       long total_cells;
+
        content->opt = opt;
        content->title = title;
        content->ncolumns = ncolumns;
@@ -3065,7 +3067,8 @@ printTableInit(printTableContent *const content, const 
printTableOpt *opt,
 
        content->headers = pg_malloc0((ncolumns + 1) * 
sizeof(*content->headers));
 
-       content->cells = pg_malloc0((ncolumns * nrows + 1) * 
sizeof(*content->cells));
+       total_cells = (long)ncolumns * (long)nrows;
+       content->cells = pg_malloc0((total_cells + 1) * 
sizeof(*content->cells));
 
        content->cellmustfree = NULL;
        content->footers = NULL;
@@ -3135,15 +3138,18 @@ void
 printTableAddCell(printTableContent *const content, char *cell,
                                  const bool translate, const bool mustfree)
 {
+       long total_cells;
 #ifndef ENABLE_NLS
        (void) translate;                       /* unused parameter */
 #endif
 
-       if (content->cellsadded >= content->ncolumns * content->nrows)
+       /* product of cols and rows, using long type to prevent the int 
overflow */
+       total_cells = (long)content->ncolumns * (long)content->nrows;
+       if (content->cellsadded >= total_cells)
        {
                fprintf(stderr, _("Cannot add cell to table content: "
-                                                 "total cell count of %d 
exceeded.\n"),
-                               content->ncolumns * content->nrows);
+                                                 "total cell count of %ld 
exceeded, cells added: %ld.\n"),
+                               total_cells, content->cellsadded);
                exit(EXIT_FAILURE);
        }
 
@@ -3159,7 +3165,7 @@ printTableAddCell(printTableContent *const content, char 
*cell,
        {
                if (content->cellmustfree == NULL)
                        content->cellmustfree =
-                               pg_malloc0((content->ncolumns * content->nrows 
+ 1) * sizeof(bool));
+                               pg_malloc0((total_cells + 1) * sizeof(bool));
 
                content->cellmustfree[content->cellsadded] = true;
        }
@@ -3227,9 +3233,10 @@ printTableCleanup(printTableContent *const content)
 {
        if (content->cellmustfree)
        {
-               int                     i;
-
-               for (i = 0; i < content->nrows * content->ncolumns; i++)
+               long            i;
+               long            total_cells;
+               total_cells = (long)content->ncolumns * (long)content->nrows;
+               for (i = 0; i < total_cells; i++)
                {
                        if (content->cellmustfree[i])
                                free(unconstify(char *, content->cells[i]));
diff --git a/src/fe_utils/test/Makefile b/src/fe_utils/test/Makefile
new file mode 100644
index 0000000000..66d3098dc5
--- /dev/null
+++ b/src/fe_utils/test/Makefile
@@ -0,0 +1,14 @@
+subdir=src/fe_utils
+top_builddir = ../../..
+include $(top_builddir)/src/Makefile.global
+
+TARGETS=print
+PG_LIBS=-lpgcommon -lpgport -lpq -lm
+
+include $(top_srcdir)/src/backend/mock.mk
+
+# This test case doesn't follow all UT guides, because current the 
infrastructure of UT
+# is for backend test only. Using it leads to many weird link errors, let's 
keep it
+# simple here.
+print.t: print_test.o
+       $(CC) $(CFLAGS) $(LDFLAGS) -Wl,--wrap=exit $< ../mbprint.o 
$(CMOCKERY_OBJS) $(PG_LIBS) -o print.t
diff --git a/src/fe_utils/test/print_test.c b/src/fe_utils/test/print_test.c
new file mode 100644
index 0000000000..b3ee926059
--- /dev/null
+++ b/src/fe_utils/test/print_test.c
@@ -0,0 +1,58 @@
+#include <stdarg.h>
+#include <stddef.h>
+#include <setjmp.h>
+#include <stdbool.h>
+#include "cmockery.h"
+
+#include "../print.c" // the file which we want to test
+
+// wrap function for exit() of printTableAddCell
+void __wrap_exit(int code);
+void
+__wrap_exit(int code)
+{
+       check_expected(code);
+       mock();
+}
+
+static void
+test_printTableAddCell(void **state)
+{
+    struct printTableContent content;
+    char cell[] = "cell";
+
+    // normal case
+    content.ncolumns = 26;
+    content.nrows = 10;
+    content.cellsadded = 100;
+    printTableAddCell(&content, cell, false, false);
+    assert_true( content.cellsadded == 101 );
+
+    // bigger than int range (2147483647) case
+    const long LONG_VAL = 4000000000L;
+    content.ncolumns = 26;
+    content.nrows = 200000000;
+    content.cellsadded = LONG_VAL;
+    printTableAddCell(&content, cell, false, false);
+    assert_true( content.cellsadded == LONG_VAL+1 );
+
+    // error case (cellsadded > ncolumns*nrows)
+    content.ncolumns = 26;
+    content.nrows = 200000000;
+    content.cellsadded = 6000000000L;
+    expect_value(__wrap_exit, code, 1);
+    will_be_called(__wrap_exit);
+    printTableAddCell(&content, cell, false, false); 
+}
+
+int
+main(int argc, char *argv[])
+{
+       cmockery_parse_arguments(argc, argv);
+
+       const           UnitTest tests[] = {
+               unit_test(test_printTableAddCell)
+       };
+
+       return run_tests(tests);
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to