Limit size of entropy image to fix crash

3 files changed, 18 insertions, 13 deletions

diff --git a/examples/triage/entropy.cpp b/examples/triage/entropy.cpp
index e87615d5..87849dcc 100644
--- a/examples/triage/entropy.cpp
+++ b/examples/triage/entropy.cpp
@@ -5,10 +5,10 @@
 #include "theme.h"
 
 
-EntropyThread::EntropyThread(BinaryViewRef data, int width, size_t blockSize)
+EntropyThread::EntropyThread(BinaryViewRef data, size_t blockSize, QImage* image)
 {
 	m_data = data;
-	m_image = QImage(width, 1, QImage::Format_ARGB32);
+	m_image = image;
 	m_blockSize = blockSize;
 	m_updated = false;
 	m_running = true;
@@ -25,7 +25,7 @@ EntropyThread::~EntropyThread()
 
 void EntropyThread::Run()
 {
-	int width = m_image.width();
+	int width = m_image->width();
 	for (int i = 0; i < width; i++)
 	{
 		if (!m_running)
@@ -39,14 +39,14 @@ void EntropyThread::Run()
 		if (v >= 240)
 		{
 			QColor color = getThemeColor(YellowStandardHighlightColor);
-			m_image.setPixelColor(i, 0, color);
+			m_image->setPixelColor(i, 0, color);
 		}
 		else
 		{
 			QColor baseColor = getThemeColor(FeatureMapBaseColor);
 			QColor entropyColor = getThemeColor(BlueStandardHighlightColor);
 			QColor color = mixColor(baseColor, entropyColor, (uint8_t)v);
-			m_image.setPixelColor(i, 0, color);
+			m_image->setPixelColor(i, 0, color);
 		}
 		m_updated = true;
 	}
@@ -59,9 +59,13 @@ EntropyWidget::EntropyWidget(QWidget* parent, TriageView* view, BinaryViewRef da
 	m_data = data;
 	m_rawData = data->GetFile()->GetViewOfType("Raw");
 
-	m_blockSize = 1024;
+	m_blockSize = (size_t)((m_rawData->GetLength() / 4096) + 1);
+	if (m_blockSize < 1024)
+		m_blockSize = 1024;
 	m_width = (int)(m_rawData->GetLength() / (uint64_t)m_blockSize);
-	m_thread = new EntropyThread(m_rawData, m_width, m_blockSize);
+	m_image = QImage(m_width, 1, QImage::Format_ARGB32);
+	m_image.fill(QColor(0, 0, 0, 0));
+	m_thread = new EntropyThread(m_rawData, m_blockSize, &m_image);
 
 	QTimer* timer = new QTimer();
 	connect(timer, &QTimer::timeout, this, &EntropyWidget::timerExpired);
@@ -82,7 +86,7 @@ EntropyWidget::~EntropyWidget()
 void EntropyWidget::paintEvent(QPaintEvent*)
 {
 	QPainter p(this);
-	p.drawImage(rect(), m_thread->GetImage());
+	p.drawImage(rect(), m_image);
 	p.drawRect(rect());
 }
 
diff --git a/examples/triage/entropy.h b/examples/triage/entropy.h
index abe591b4..eb75cfee 100644
--- a/examples/triage/entropy.h
+++ b/examples/triage/entropy.h
@@ -9,20 +9,18 @@
 class EntropyThread
 {
 	BinaryViewRef m_data;
-	QImage m_image;
+	QImage* m_image;
 	size_t m_blockSize;
 	bool m_updated, m_running;
 	std::thread m_thread;
 
 public:
-	EntropyThread(BinaryViewRef data, int width, size_t blockSize);
+	EntropyThread(BinaryViewRef data, size_t blockSize, QImage* image);
 	~EntropyThread();
 
 	void Run();
 	bool IsUpdated() { return m_updated; }
 	void ResetUpdated() { m_updated = false; }
-
-	const QImage& GetImage() { return m_image; }
 };
 
 
@@ -34,6 +32,7 @@ class EntropyWidget: public QWidget
 	BinaryViewRef m_data, m_rawData;
 	size_t m_blockSize;
 	int m_width;
+	QImage m_image;
 	EntropyThread* m_thread;
 
 public:
diff --git a/python/examples/triage/entropy.py b/python/examples/triage/entropy.py
index 94afd650..57e9d2c4 100644
--- a/python/examples/triage/entropy.py
+++ b/python/examples/triage/entropy.py
@@ -37,7 +37,9 @@ class EntropyWidget(QWidget):
 		self.data = data
 		self.raw_data = data.file.raw
 
-		self.block_size = 1024
+		self.block_size = (len(self.raw_data) / 4096) + 1
+		if self.block_size < 1024:
+			self.block_size = 1024
 		self.width = int(len(self.raw_data) / self.block_size)
 		self.image = QImage(self.width, 1, QImage.Format_ARGB32)
 		self.image.fill(QColor(0, 0, 0, 0))