touch/focaltech: check the heatmap data size to avoid invalid access.

Change list:
1. Remove to switch the setting between uncompressed and compressed heatmap in fts_test_get_strength().
2. Check if the size of compressed heatmap is valid.

Bug: 236026908
Test: Stress test pass, there is no memory corruption when draw circles
      on the panel while running the dump touch state script.

Signed-off-by: Steve Pfetsch <spfetsch@google.com>
Change-Id: I4ecc1b3470e02f02cfd12cd3752ece540e6705c0

4 files changed, 50 insertions, 41 deletions

diff --git a/ft3658/focaltech_common.h b/ft3658/focaltech_common.h
index 34c3d42..4688d96 100644
--- a/ft3658/focaltech_common.h
+++ b/ft3658/focaltech_common.h
@@ -146,6 +146,8 @@
                                                     // bit 7   : LPWG
 #define FTS_CAP_DATA_LEN                    91
 #define FTS_SELF_DATA_LEN                   68
+#define FTS_FULL_HEATMAP_RAW_SIZE(tx_num, rx_num) \
+    (FTS_CAP_DATA_LEN + ((tx_num) * (rx_num) + FTS_SELF_DATA_LEN * 2) * sizeof(u16))
 #define FTS_PRESSURE_SCALE                  85      // 255 / 3
 #define FTS_CUSTOMER_STATUS_LEN             4
 #define FTS_CUSTOMER_STATUS1_MASK           0x0F
diff --git a/ft3658/focaltech_core.c b/ft3658/focaltech_core.c
index 65f4501..f982bd7 100644
--- a/ft3658/focaltech_core.c
+++ b/ft3658/focaltech_core.c
@@ -821,7 +821,7 @@ static int fts_read_touchdata(struct fts_ts_data *data)
     ret = fts_get_heatmap(data);
     if (ret < 0)
         return ret;
-    memcpy(buf + 1,data->heatmap_raw, data->pnt_buf_size - 1);
+    memcpy(buf + 1, data->heatmap_raw, data->pnt_buf_size - 1);
 #else
     ret = fts_read(cmd, 1, buf + 1, data->pnt_buf_size - 1);
     if (ret < 0) {
@@ -1488,12 +1488,20 @@ static int fts_get_heatmap(struct fts_ts_data *ts_data) {
         /* Total touch data: (cap header(91) + heatmap(N-MS + W-SS + N-SS)). */
         total_data_size = FTS_CAP_DATA_LEN + self_data_size * 2 +
                           mutual_data_size;
+
+        if (total_data_size > ts_data->heatmap_raw_size) {
+            FTS_DEBUG("Warning : The total touch data size is %d!!",
+                total_data_size);
+            total_data_size = ts_data->heatmap_raw_size;
+        }
+
         ret = fts_read(cmd, 1, ts_data->heatmap_raw, total_data_size);
         if (ret < 0) {
             FTS_ERROR("Failed to get heatmap raw data, ret=%d.", ret);
             ret = -EIO;
             goto exit;
         }
+
         /* Get the self-sensing type. */
         ts_data->self_sensing_type =
             ts_data->heatmap_raw[FTS_CAP_DATA_LEN - 1] & 0x80;
@@ -1511,40 +1519,52 @@ static int fts_get_heatmap(struct fts_ts_data *ts_data) {
          * |-     91     -|-   68*2   -|-   68*2    -|- (B2[1]<<8+B2[2])*2    -|
          */
 
+        if (ts_data->compress_heatmap_wlen < 0 ||
+            (ts_data->compress_heatmap_wlen * sizeof(u16)) > mutual_data_size) {
+            FTS_DEBUG("Warning : The compressed heatmap size is %d!!",
+                ts_data->compress_heatmap_wlen);
+            ts_data->compress_heatmap_wlen = 0;
+            memset(ts_data->trans_raw, 0, ts_data->trans_raw_size);
+        }
+
         /* Total touch data:(cap header + W-SS + N-SS + compressed heatmap(N-MS)
          */
         total_data_size = FTS_CAP_DATA_LEN +
                           self_data_size * 2 +
                           ts_data->compress_heatmap_wlen * sizeof(u16);
 
+        if (total_data_size > ts_data->heatmap_raw_size) {
+            FTS_DEBUG("Warning : The total touch data size is %d!!",
+                total_data_size);
+            total_data_size = ts_data->heatmap_raw_size;
+        }
+
         ret = fts_read(cmd, 1, ts_data->heatmap_raw, total_data_size);
         if (ret < 0) {
             FTS_ERROR("Failed to get compressed heatmap raw data,ret=%d.", ret);
             ret = -EIO;
             goto exit;
         }
-        if (ts_data->compress_heatmap_wlen == 0) {
-            FTS_DEBUG("Warning : The compressed heatmap length is 0!!");
-            goto exit;
-        }
 
         /* Get the self-sensing type. */
         ts_data->self_sensing_type =
             ts_data->heatmap_raw[FTS_CAP_DATA_LEN - 1] & 0x80;
 
-        /* decode the compressed data from heatmap_raw to heatmap_buff. */
-        fts_ptflib_decoder(ts_data,
-            (u16*)(&ts_data->heatmap_raw[idx_ms_raw]),
-            ts_data->compress_heatmap_wlen,
-            ts_data->heatmap_buff,
-            mutual_data_size / sizeof(u16));
-
-        /* MS: Transform the order from RX->TX. */
-        /* After decoding, the data become to little-endian, but the output of
-         * transpose_raw is big-endian.
-         */
-        transpose_raw(&((u8*)ts_data->heatmap_buff)[0], ts_data->trans_raw,
-            tx, rx, false);
+        if (ts_data->compress_heatmap_wlen > 0) {
+            /* decode the compressed data from heatmap_raw to heatmap_buff. */
+            fts_ptflib_decoder(ts_data,
+                (u16*)(&ts_data->heatmap_raw[idx_ms_raw]),
+                ts_data->compress_heatmap_wlen,
+                ts_data->heatmap_buff,
+                mutual_data_size / sizeof(u16));
+
+            /* MS: Transform the order from RX->TX. */
+            /* After decoding, the data become to little-endian, but the output of
+             * transpose_raw is big-endian.
+             */
+            transpose_raw(&((u8*)ts_data->heatmap_buff)[0], ts_data->trans_raw,
+                tx, rx, false);
+        }
     }
 #if IS_ENABLED(GOOGLE_HEATMAP_DEBUG)
     FTS_DEBUG("Copy matual data,idx_buff=%d,idx_ms_raw=%d.",
diff --git a/ft3658/focaltech_test/focaltech_test.c b/ft3658/focaltech_test/focaltech_test.c
index 9fc6d79..c2f3c15 100644
--- a/ft3658/focaltech_test/focaltech_test.c
+++ b/ft3658/focaltech_test/focaltech_test.c
@@ -2618,7 +2618,7 @@ static const struct file_operations proc_test_int_fops = {
 
 extern int fts_test_get_raw(int *raw, u8 tx, u8 rx);
 extern int fts_test_get_baseline(int *raw,int *base_raw, u8 tx, u8 rx);
-extern int fts_test_get_strength(u8 *base_raw, u8 tx, u8 rx);
+extern int fts_test_get_strength(u8 *base_raw, u16 base_raw_size);
 extern int fts_test_get_uniformity_data(int *rawdata_linearity, u8 tx, u8 rx);
 extern int fts_test_get_scap_raw(int *scap_raw, u8 tx, u8 rx, int *fwcheck);
 extern int fts_test_get_scap_cb(int *scap_cb, u8 tx, u8 rx, int *fwcheck);
@@ -2853,7 +2853,7 @@ static int proc_test_strength_show(struct seq_file *s, void *v)
 
     u8 *base_raw = NULL;
     u8 *trans_raw = NULL;
-    int base_raw_len = 0;
+    int base_raw_size = 0;
     int base = 0;
     int fast_events_x = 0;
     int fast_events_y = 0;
@@ -2868,10 +2868,9 @@ static int proc_test_strength_show(struct seq_file *s, void *v)
     node_num = tx * rx;
     self_node = tx + rx;
 
-    base_raw_len =
-        FTS_CAP_DATA_LEN + (FTS_SELF_DATA_LEN * 2 + node_num) * sizeof(u16);
-    FTS_DEBUG("heapmap base_raw length = %d", base_raw_len);
-    base_raw = fts_malloc(base_raw_len);
+    base_raw_size = FTS_FULL_HEATMAP_RAW_SIZE(tx, rx);
+    FTS_DEBUG("heapmap base_raw size = %d", base_raw_size);
+    base_raw = fts_malloc(base_raw_size);
     if (!base_raw) {
         FTS_ERROR("malloc memory for raw fails");
         ret = -ENOMEM;
@@ -2886,7 +2885,7 @@ static int proc_test_strength_show(struct seq_file *s, void *v)
     }
 
     /* get strength data. */
-    ret = fts_test_get_strength(base_raw, tx, rx);
+    ret = fts_test_get_strength(base_raw, base_raw_size);
     if (ret < 0) {
         FTS_ERROR("get strength fails");
         goto exit;
diff --git a/ft3658/focaltech_test/supported_ic/focaltech_test_ft3658u.c b/ft3658/focaltech_test/supported_ic/focaltech_test_ft3658u.c
index 0d9617a..719dc29 100644
--- a/ft3658/focaltech_test/supported_ic/focaltech_test_ft3658u.c
+++ b/ft3658/focaltech_test/supported_ic/focaltech_test_ft3658u.c
@@ -2037,31 +2037,19 @@ test_err:
     return ret;
 }
 
-int fts_test_get_strength(u8 *base_raw, u8 tx, u8 rx)
+int fts_test_get_strength(u8 *base_raw, u16 base_raw_size)
 {
     int ret = 0;
-    struct fts_ts_data *ts_data = fts_data;
-    int self_cap_offset = 91;
-    int self_cap_len = 68;
-    int self_data_len = self_cap_len * 2; /* cap and cap_off */
-    int fast_num_len = self_cap_offset +
-                       ((tx * rx) + self_data_len) * sizeof(u16);
-    u8 id_cmd[4] = {0};
+    u8 id_cmd[1] = {0};
 
     FTS_TEST_INFO("====== Test Item: strength test start\n");
-    id_cmd[0] = 0x01;
-    /* Enable uncompressed heatmap. */
-    fts_set_heatmap_mode(ts_data, FW_HEATMAP_MODE_UNCOMPRESSED);
+    id_cmd[0] = FTS_CMD_READ_TOUCH_DATA;
     sys_delay(500);
-    FTS_TEST_DBG("Allocate heatmap length = %d.\n", fast_num_len);
-    ret = fts_read(id_cmd, 1, base_raw, fast_num_len);
-
+    ret = fts_read(id_cmd, 1, base_raw, base_raw_size);
     if (ret < 0) {
         FTS_TEST_ERROR("get strength fail,ret=%d\n", ret);
     }
 
-    /* Enable compressed heatmap. */
-    fts_set_heatmap_mode(ts_data, FW_HEATMAP_MODE_COMPRESSED);
     FTS_TEST_INFO("====== Test Item: strength test end\n");
     return ret;
 }