2 files changed, 52 insertions, 6 deletions

diff --git a/src/profiling/test/BufferTests.cpp b/src/profiling/test/BufferTests.cpp
index 0225d750a7..804335138d 100644
--- a/src/profiling/test/BufferTests.cpp
+++ b/src/profiling/test/BufferTests.cpp
@@ -1,5 +1,5 @@
 //
-// Copyright © 2019 Arm Ltd. All rights reserved.
+// Copyright © 2019 Arm Ltd and Contributors. All rights reserved.
 // SPDX-License-Identifier: MIT
 //
 
@@ -363,4 +363,48 @@ BOOST_AUTO_TEST_CASE(BufferMarkReadTest)
     BOOST_TEST(packetBuffer3.get());
 }
 
+BOOST_AUTO_TEST_CASE(ReadSwTraceMessageExceptionTest0)
+{
+    IPacketBufferPtr packetBuffer = std::make_unique<PacketBuffer>(512);
+
+    BOOST_TEST(packetBuffer->GetSize() == 0);
+
+    // Write zero data to the buffer
+    WriteUint32(packetBuffer, 0, 0);
+    WriteUint32(packetBuffer, 4, 0);
+    WriteUint32(packetBuffer, 8, 0);
+    WriteUint32(packetBuffer, 12, 0);
+
+    // Commit
+    packetBuffer->Commit(16);
+
+    unsigned int uint32_t_size = sizeof(uint32_t);
+    unsigned int offset = uint32_t_size;
+    BOOST_CHECK_THROW(ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize()),
+                      armnn::RuntimeException);
+
+}
+
+BOOST_AUTO_TEST_CASE(ReadSwTraceMessageExceptionTest1)
+{
+    IPacketBufferPtr packetBuffer = std::make_unique<PacketBuffer>(512);
+
+    BOOST_TEST(packetBuffer->GetSize() == 0);
+
+    // Write data to the buffer
+    WriteUint32(packetBuffer, 0, 10);
+    WriteUint32(packetBuffer, 4, 20);
+    WriteUint32(packetBuffer, 8, 30);
+    WriteUint32(packetBuffer, 12, 40);
+
+    // Commit
+    packetBuffer->Commit(16);
+
+    unsigned int uint32_t_size = sizeof(uint32_t);
+    unsigned int offset = uint32_t_size;
+    BOOST_CHECK_THROW(ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize()),
+                      armnn::RuntimeException);
+
+}
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/profiling/test/SendTimelinePacketTests.cpp b/src/profiling/test/SendTimelinePacketTests.cpp
index da30cef90c..5e9f6bddea 100644
--- a/src/profiling/test/SendTimelinePacketTests.cpp
+++ b/src/profiling/test/SendTimelinePacketTests.cpp
@@ -72,7 +72,9 @@ BOOST_AUTO_TEST_CASE(SendTimelineMessageDirectoryPackageTest)
     BOOST_CHECK(DeclCount == 5);
 
     offset += uint32_t_size;
-    SwTraceMessage swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset);
+    SwTraceMessage swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(),
+                                                       offset,
+                                                       packetBuffer->GetSize());
 
     BOOST_CHECK(swTraceMessage.m_Id == 0);
     BOOST_CHECK(swTraceMessage.m_Name == "declareLabel");
@@ -84,7 +86,7 @@ BOOST_AUTO_TEST_CASE(SendTimelineMessageDirectoryPackageTest)
     BOOST_CHECK(swTraceMessage.m_ArgNames[0] == "guid");
     BOOST_CHECK(swTraceMessage.m_ArgNames[1] == "value");
 
-    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset);
+    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize());
 
     BOOST_CHECK(swTraceMessage.m_Id == 1);
     BOOST_CHECK(swTraceMessage.m_Name == "declareEntity");
@@ -94,7 +96,7 @@ BOOST_AUTO_TEST_CASE(SendTimelineMessageDirectoryPackageTest)
     BOOST_CHECK(swTraceMessage.m_ArgNames.size() == 1);
     BOOST_CHECK(swTraceMessage.m_ArgNames[0] == "guid");
 
-    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset);
+    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize());
 
     BOOST_CHECK(swTraceMessage.m_Id == 2);
     BOOST_CHECK(swTraceMessage.m_Name == "declareEventClass");
@@ -106,7 +108,7 @@ BOOST_AUTO_TEST_CASE(SendTimelineMessageDirectoryPackageTest)
     BOOST_CHECK(swTraceMessage.m_ArgNames[0] == "guid");
     BOOST_CHECK(swTraceMessage.m_ArgNames[1] == "nameGuid");
 
-    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset);
+    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize());
 
     BOOST_CHECK(swTraceMessage.m_Id == 3);
     BOOST_CHECK(swTraceMessage.m_Name == "declareRelationship");
@@ -124,7 +126,7 @@ BOOST_AUTO_TEST_CASE(SendTimelineMessageDirectoryPackageTest)
     BOOST_CHECK(swTraceMessage.m_ArgNames[3] == "tailGuid");
     BOOST_CHECK(swTraceMessage.m_ArgNames[4] == "attributeGuid");
 
-    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset);
+    swTraceMessage = ReadSwTraceMessage(packetBuffer->GetReadableData(), offset, packetBuffer->GetSize());
 
     BOOST_CHECK(swTraceMessage.m_Id == 4);
     BOOST_CHECK(swTraceMessage.m_Name == "declareEvent");