From a995cbf9ecafc44a73da17c193da982c2aa2492a Mon Sep 17 00:00:00 2001
From: alexander-g <3867427+alexander-g@users.noreply.github.com>
Date: Sat, 30 Jan 2021 19:41:47 +0100
Subject: [PATCH] disallowing zero sized tensors

---
 src/Tensor.cpp              |  6 +++++-
 test/TestOpTensorCreate.cpp | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/Tensor.cpp b/src/Tensor.cpp
index 214ac2e..63e6670 100644
--- a/src/Tensor.cpp
+++ b/src/Tensor.cpp
@@ -268,10 +268,14 @@ Tensor::createBuffer()
         throw std::runtime_error("Kompute Tensor device is null");
     }
 
-    this->mFreeBuffer = true;
 
     vk::BufferUsageFlags usageFlags = this->getBufferUsageFlags();
     vk::DeviceSize bufferSize = this->memorySize();
+    if(bufferSize<1){
+        throw std::runtime_error("Kompute Tensor attempted to create a zero-sized buffer");
+    }
+    
+    this->mFreeBuffer = true;
 
     SPDLOG_DEBUG("Kompute Tensor creating buffer with memory size: {}, and "
                  "usage flags: {}",
diff --git a/test/TestOpTensorCreate.cpp b/test/TestOpTensorCreate.cpp
index 6f8a39c..e25c135 100644
--- a/test/TestOpTensorCreate.cpp
+++ b/test/TestOpTensorCreate.cpp
@@ -113,3 +113,21 @@ TEST(TestOpTensorCreate, NoErrorIfTensorFreedBefore)
     EXPECT_FALSE(tensorA->isInit());
     EXPECT_FALSE(tensorB->isInit());
 }
+
+
+TEST(TestOpTensorCreate, ExceptionOnZeroSizeTensor)
+{
+    std::vector<float> testVecA;
+
+    std::shared_ptr<kp::Tensor> tensorA{ new kp::Tensor(testVecA) };
+
+    kp::Manager mgr;
+
+    try{
+        mgr.evalOpDefault<kp::OpTensorCreate>({ tensorA });
+    } catch( const std::runtime_error& err ) {
+         // check exception
+        ASSERT_STREQ("Kompute Tensor attempted to create a zero-sized buffer", err.what() );
+    }
+
+}