Limit size of entropy image to fix crash

1 files changed, 12 insertions, 8 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());
 }