From 7495d1b6e16f419c567ac6ace4e725349f495dc9 Mon Sep 17 00:00:00 2001
From: Marek Niepieklo <marek.niepieklo@mudita.com>
Date: Fri, 9 Jul 2021 19:09:33 +0200
Subject: [PATCH] [CP-327] Fix UB in QueryResult

QueryResult iterator nextRow() may get into infinite loop.
Additionally, fixed handling of the case when each row of result
could have a different count of fields.
---
 module-db/Database/Database.cpp    |  8 +++++---
 module-db/Database/Database.hpp    |  4 ++--
 module-db/Database/QueryResult.cpp | 18 ++++++------------
 module-db/Database/QueryResult.hpp |  8 +++-----
 4 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/module-db/Database/Database.cpp b/module-db/Database/Database.cpp
index 8aebfedeb3fd5d91fad74f6b62184dd33b97ba39..26c27f006d3cdb1c7e5e1025d2d7af77598651cb 100644
--- a/module-db/Database/Database.cpp
+++ b/module-db/Database/Database.cpp
@@ -195,11 +195,13 @@ uint32_t Database::getLastInsertRowId()
     return sqlite3_last_insert_rowid(dbConnection);
 }
 
-void Database::pragmaQuery(const std::string &pragmaStatemnt)
+void Database::pragmaQuery(const std::string &pragmaStatement)
 {
-    if (auto results = query(pragmaStatemnt.c_str()); results) {
-        const auto fieldsCount = results->getFieldCount();
+    auto results = query(pragmaStatement.c_str());
+
+    if (results && results->getRowCount()) {
         do {
+            const auto fieldsCount = results->getFieldCount();
             for (uint32_t i = 0; i < fieldsCount; i++) {
                 Field field = (*results)[i];
             }
diff --git a/module-db/Database/Database.hpp b/module-db/Database/Database.hpp
index 5c6005f95dbfa4522333fa1824e10cc8cf10acc9..48bd66bb6c304a35de5d181a80145d850f1c13fa 100644
--- a/module-db/Database/Database.hpp
+++ b/module-db/Database/Database.hpp
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2020, Mudita Sp. z.o.o. All rights reserved.
+// Copyright (c) 2017-2021, Mudita Sp. z.o.o. All rights reserved.
 // For licensing, see https://github.com/mudita/MuditaOS/LICENSE.md
 
 #pragma once
@@ -37,7 +37,7 @@ class Database
     bool storeIntoFile(const std::filesystem::path &backupPath);
 
     uint32_t getLastInsertRowId();
-    void pragmaQuery(const std::string &pragmaStatemnt);
+    void pragmaQuery(const std::string &pragmaStatement);
 
     [[nodiscard]] bool isInitialized() const noexcept
     {
diff --git a/module-db/Database/QueryResult.cpp b/module-db/Database/QueryResult.cpp
index d73bdb04ea094dcb5ad6ec11c67a68aea4b51a42..51d7f2911786594ffd6bc224b734ee6762922eff 100644
--- a/module-db/Database/QueryResult.cpp
+++ b/module-db/Database/QueryResult.cpp
@@ -1,25 +1,19 @@
-// Copyright (c) 2017-2020, Mudita Sp. z.o.o. All rights reserved.
+// Copyright (c) 2017-2021, Mudita Sp. z.o.o. All rights reserved.
 // For licensing, see https://github.com/mudita/MuditaOS/LICENSE.md
 
 #include "QueryResult.hpp"
 
-QueryResult::QueryResult() : currentRow(0), fieldCount(0), rowCount(0)
+QueryResult::QueryResult() : currentRow(0)
 {}
 
 void QueryResult::addRow(const std::vector<Field> &row)
 {
     rows.push_back(row);
-    rowCount++;
-    fieldCount = row.size();
 }
 
 bool QueryResult::nextRow()
 {
-    if (currentRow < (rowCount - 1)) {
-        currentRow++;
-        return true;
-    }
-    else {
-        return false;
-    }
-}
\ No newline at end of file
+    ++currentRow;
+
+    return (currentRow < rows.size());
+}
diff --git a/module-db/Database/QueryResult.hpp b/module-db/Database/QueryResult.hpp
index 18bcca38e0c379cdcfd36c37967188907be4a0eb..4ee01a200a0246e3a5f3831e612b185e3b91f2bd 100644
--- a/module-db/Database/QueryResult.hpp
+++ b/module-db/Database/QueryResult.hpp
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2020, Mudita Sp. z.o.o. All rights reserved.
+// Copyright (c) 2017-2021, Mudita Sp. z.o.o. All rights reserved.
 // For licensing, see https://github.com/mudita/MuditaOS/LICENSE.md
 
 #pragma once
@@ -27,17 +27,15 @@ class QueryResult
 
     uint32_t getFieldCount() const
     {
-        return fieldCount;
+        return rows[currentRow].size();
     }
 
     uint32_t getRowCount() const
     {
-        return rowCount;
+        return rows.size();
     }
 
   private:
     uint32_t currentRow;
     std::vector<std::vector<Field>> rows;
-    uint32_t fieldCount;
-    uint32_t rowCount;
 };