From 2c01713dbeba8446b135ef19fc443b6abe098210 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonny=20Sv=C3=A4rd?= <jonny.svaerd@arm.com>
Date: Mon, 12 Apr 2021 15:56:44 +0200
Subject: Add nullptr checks

Change-Id: Ie76db24dec03a6b45c64070d89a83805397d2f1f
---
 driver_library/src/ethosu.cpp | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/driver_library/src/ethosu.cpp b/driver_library/src/ethosu.cpp
index a888928..1768271 100644
--- a/driver_library/src/ethosu.cpp
+++ b/driver_library/src/ethosu.cpp
@@ -50,6 +50,10 @@ int eioctl(int fd, unsigned long cmd, void *data = nullptr) {
 size_t getShapeSize(const flatbuffers::Vector<int32_t> *shape) {
     size_t size = 1;
 
+    if (shape == nullptr) {
+        throw EthosU::Exception("getShapeSize(): nullptr arg");
+    }
+
     for (auto it = shape->begin(); it != shape->end(); ++it) {
         size *= *it;
     }
@@ -74,6 +78,10 @@ size_t getTensorTypeSize(const enum tflite::TensorType type) {
 vector<size_t> getSubGraphDims(const tflite::SubGraph *subgraph, const flatbuffers::Vector<int32_t> *tensorMap) {
     vector<size_t> dims;
 
+    if (subgraph == nullptr || tensorMap == nullptr) {
+        throw EthosU::Exception("getSubGraphDims(): nullptr arg(s)");
+    }
+
     for (auto index = tensorMap->begin(); index != tensorMap->end(); ++index) {
         auto tensor = subgraph->tensors()->Get(*index);
         size_t size = getShapeSize(tensor->shape());
@@ -190,6 +198,10 @@ Network::Network(Device &device, shared_ptr<Buffer> &buffer) : fd(-1), buffer(bu
     // Create model handle
     const tflite::Model *model = tflite::GetModel(reinterpret_cast<void *>(buffer->data()));
 
+    if (model->subgraphs() == nullptr) {
+        throw EthosU::Exception("Failed to get subgraphs: nullptr");
+    }
+
     // Get input dimensions for first subgraph
     auto *subgraph = *model->subgraphs()->begin();
     ifmDims        = getSubGraphDims(subgraph, subgraph->inputs());
-- 
cgit v1.2.1