From aa59fc3ac649652bb92d5a245737548a0d0c2e92 Mon Sep 17 00:00:00 2001
From: Jason Ertel <jason.ertel@securityonionsolutions.com>
Date: Mon, 6 Jun 2022 16:32:12 -0400
Subject: [PATCH] Avoid recursive Sankey diagrams by omitting datasets that
 lead to recursion

---
 html/js/i18n.js             |  1 +
 html/js/routes/hunt.js      | 43 +++++++++++++++++++++++++++++++------
 html/js/routes/hunt.test.js | 19 ++++++++++++++--
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/html/js/i18n.js b/html/js/i18n.js
index e871bff6..96c913a7 100644
--- a/html/js/i18n.js
+++ b/html/js/i18n.js
@@ -131,6 +131,7 @@ const i18n = {
       category: 'Category',
 
       chartTitleBottom: 'Fewest Occurrences',
+      chartTitleIncomplete: '(partial)',
       chartTitleTimeline: 'Timeline',
       chartTitleTop: 'Most Occurrences',
       cheatsheet: 'Cheat Sheet',
diff --git a/html/js/routes/hunt.js b/html/js/routes/hunt.js
index 301cea84..6c506d6f 100644
--- a/html/js/routes/hunt.js
+++ b/html/js/routes/hunt.js
@@ -988,6 +988,7 @@ const huntComponent = {
           // chart_type:    ChartJS type, such as pie, bar, sankey, etc.
           // chart_options: ChartJS options. See setupBarChart, etc.
           // chart_data:    ChartJS labels and datasets. See setupBarChart and populateBarChart.
+          // is_incomplete: True if only partial data is rendered to avoid complete render failure.
           var group = {};
           group.title = fields.join(this.chartLabelFieldSeparator);
           group.fields = [...fields];
@@ -1078,34 +1079,62 @@ const huntComponent = {
       group.chart_type = "sankey";
       group.chart_options = {};
       group.chart_data = {};
-      this.setupSankeyChart(group.chart_options, group.chart_data, group.title);
-      this.applyLegendOption(group, groupIdx);
 
       // Sankey has a unique dataset format, build it out here instead of using populateChartData().
       // While building the new format, also calculate the max value across all nodes to be used
       // as a scale factor for choosing colors of the sankey flows.
       var flowMax = 0;
-      updateMaxMap = function(map, key, value) {
+      var updateMaxMap = function(map, key, value) {
         var max = map[key];
         if (!max) {
           max = 0;
         }
         max = max + value;
-        maxFlowMap[key] = max;
+        map[key] = max;
         flowMax = Math.max(flowMax, max);
-      }
+      };
+
+      var isRecursive = function(map, from, to, current, max) {
+        if (current > max || from == to) {
+          return true;
+        }
+
+        for (var i = 0; i < map.length; i++) {
+          var item = map[i];
+          if (item.from == to) {
+            if (isRecursive(map, item.from, item.to, current + 1, max)) {
+              return true;
+            }
+          }
+        }
+        return false;
+      };
+
       var data = [];
       var maxFlowMap = {};
       group.data.forEach(function(item, index) {
         for (var idx = 0; idx < group.fields.length - 1; idx++) {
           var from = item[group.fields[idx]];
           var to = item[group.fields[idx+1]];
-          updateMaxMap(maxFlowMap, from, item.count);
-          updateMaxMap(maxFlowMap, to, item.count);
           var flow = { from: from, to: to, flow: item.count };
           data.push(flow);
+
+          if (isRecursive(data, from, to, 0, group.fields.length)) {
+            group.is_incomplete = true;
+            data.pop();
+          } else {
+            updateMaxMap(maxFlowMap, from, item.count);
+            updateMaxMap(maxFlowMap, to, item.count);
+          }
         }
       });
+
+      if (group.is_incomplete) {
+        group.title += " " + this.i18n.chartTitleIncomplete;
+      }
+      this.setupSankeyChart(group.chart_options, group.chart_data, group.title);
+      this.applyLegendOption(group, groupIdx);
+
       group.chart_data.datasets[0].data = data;
       group.chart_data.flowMax = flowMax;
       Vue.set(this.groupBys, groupIdx, group);
diff --git a/html/js/routes/hunt.test.js b/html/js/routes/hunt.test.js
index ccc8dbd1..cdbed5d5 100644
--- a/html/js/routes/hunt.test.js
+++ b/html/js/routes/hunt.test.js
@@ -408,14 +408,19 @@ test('displayPieChart', () => {
 
 test('displaySankeyChart', () => {
   var group = {chart_type: ''};
-  group.data = [{ count: 1, foo: 'moo', bar: 'mar' }, { count: 12, foo: 'moo', bar: 'car' }]
+  group.data = [{ count: 10, foo: 'mog', bar: 'mop' }, { count: 1, foo: 'moo', bar: 'mar' }, { count: 12, foo: 'moo', bar: 'car' }, { count: 2, foo: 'moo', bar: 'mog' }, { count: 2, foo: 'mop', bar: 'moo' },{ count: 2, foo: 'moo', bar: 'moo' }, { count: 3, foo: 'mop', bar: 'baz' }]
   group.fields = ['foo', 'bar'];
   comp.groupBys = [group];
   comp.queryGroupByOptions = [[]];
   comp.displaySankeyChart(group, 0);
   expect(group.chart_type).toBe('sankey');
-  expect(group.chart_data.flowMax).toBe(13);
+  expect(group.chart_data.flowMax).toBe(15);
   expect(group.chart_data.datasets[0].data).toStrictEqual([
+    {
+      "flow": 10,
+      "from": "mog",
+      "to": "mop",
+    },
     {
       "flow": 1,
       "from": "moo",
@@ -426,6 +431,16 @@ test('displaySankeyChart', () => {
       "from": "moo",
       "to": "car",
     },
+    {
+      "flow": 2,
+      "from": "moo",
+      "to": "mog",
+    },
+    {
+      "flow": 3,
+      "from": "mop",
+      "to": "baz",
+    },
   ]);
 });