Verify contents of tensors
PiperOrigin-RevId: 184003263
diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD
index efb6449..ac6f013 100644
--- a/tensorflow/contrib/BUILD
+++ b/tensorflow/contrib/BUILD
@@ -24,6 +24,7 @@
"//tensorflow/contrib/bayesflow:bayesflow_py",
"//tensorflow/contrib/boosted_trees:init_py",
"//tensorflow/contrib/cloud:cloud_py",
+ "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip",
"//tensorflow/contrib/cluster_resolver:cluster_resolver_py",
"//tensorflow/contrib/coder:coder_ops_py",
"//tensorflow/contrib/compiler:compiler_py",
diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD
index 1bffcfb..4d3b553 100644
--- a/tensorflow/contrib/lite/tools/BUILD
+++ b/tensorflow/contrib/lite/tools/BUILD
@@ -99,8 +99,11 @@
srcs = ["verifier.cc"],
hdrs = ["verifier.h"],
deps = [
+ "//tensorflow/contrib/lite:framework",
"//tensorflow/contrib/lite:schema_fbs_version",
+ "//tensorflow/contrib/lite:string_util",
"//tensorflow/contrib/lite/schema:schema_fbs",
+ "@com_google_absl//absl/base:core_headers",
],
)
@@ -112,6 +115,7 @@
":verifier",
"//tensorflow/contrib/lite:framework",
"//tensorflow/contrib/lite:schema_fbs_version",
+ "//tensorflow/contrib/lite:string_util",
"//tensorflow/contrib/lite/schema:schema_fbs",
"//tensorflow/contrib/lite/testing:util",
"@com_google_googletest//:gtest",
diff --git a/tensorflow/contrib/lite/tools/verifier.cc b/tensorflow/contrib/lite/tools/verifier.cc
index 95a0895..726e2aa 100644
--- a/tensorflow/contrib/lite/tools/verifier.cc
+++ b/tensorflow/contrib/lite/tools/verifier.cc
@@ -14,13 +14,32 @@
==============================================================================*/
#include "tensorflow/contrib/lite/tools/verifier.h"
+#include <climits>
#include "tensorflow/contrib/lite/schema/schema_generated.h"
+#include "tensorflow/contrib/lite/string_util.h"
#include "tensorflow/contrib/lite/version.h"
namespace tflite {
namespace {
+// Reports error message when the reporter is set.
+void ReportError(ErrorReporter* error_reporter, const char* format, ...) {
+ if (error_reporter) {
+ va_list args;
+ va_start(args, format);
+ error_reporter->Report(format, args);
+ va_end(args);
+ }
+}
+
+// Returns the int32_t value pointed by ptr.
+const uint32_t* GetIntPtr(const char* ptr) {
+ return reinterpret_cast<const uint32_t*>(ptr);
+}
+
+// Verifies flatbuffer format of the model contents and returns the in-memory
+// model.
const Model* VerifyFlatbufferAndGetModel(const void* buf, size_t len) {
::flatbuffers::Verifier verifier(static_cast<const uint8_t*>(buf), len);
if (VerifyModelBuffer(verifier)) {
@@ -30,14 +49,159 @@
}
}
-} // namespace
+const uint32_t kMaxNumString = UINT_MAX / sizeof(int32_t) - 2;
-bool Verify(const void* buf, size_t len) {
- const Model* model = VerifyFlatbufferAndGetModel(buf, len);
- if (model == nullptr) {
+// Verifies string tensor has legit buffer contents that follow the schema
+// defined in lite/string_util.h
+bool VerifyStringTensorBuffer(const Buffer& buffer,
+ ErrorReporter* error_reporter) {
+ uint32_t buffer_size = buffer.data()->size();
+ const char* buffer_ptr = reinterpret_cast<const char*>(buffer.data()->data());
+
+ uint32_t num_strings = *GetIntPtr(buffer_ptr);
+ if (num_strings > kMaxNumString) {
+ ReportError(error_reporter,
+ "String tensor has invalid num of string set: %d", num_strings);
+ return false;
+ }
+ uint32_t header_offsets =
+ static_cast<uint32_t>(num_strings + 2) * sizeof(int32_t);
+
+ if (buffer_size < header_offsets) {
+ ReportError(error_reporter,
+ "String tensor buffer requires at least %d bytes, but is "
+ "allocated with %d bytes",
+ header_offsets, buffer_size);
return false;
}
- return model->version() == TFLITE_SCHEMA_VERSION;
+ uint32_t prev_ptr = header_offsets;
+ uint32_t offset = sizeof(int32_t);
+
+ if (*GetIntPtr(buffer_ptr + offset) != header_offsets) {
+ ReportError(error_reporter,
+ "String tensor buffer initial offset must be: %d",
+ header_offsets);
+ return false;
+ }
+ offset += sizeof(int32_t);
+ for (int i = 1; i <= num_strings; i++, offset += sizeof(int32_t)) {
+ int string_offset = *GetIntPtr(buffer_ptr + offset);
+ if (string_offset < prev_ptr || string_offset > buffer_size) {
+ ReportError(error_reporter, "String tensor buffer is invalid: index %d",
+ i);
+ return false;
+ }
+ }
+ if (*GetIntPtr(buffer_ptr + offset - sizeof(int32_t)) != buffer_size) {
+ ReportError(error_reporter, "String tensor buffer last offset must be %d",
+ buffer_size);
+ return false;
+ }
+ return true;
+}
+
+// Verifies numeric tensor has legit buffer.
+bool VerifyNumericTensorBuffer(const Tensor& tensor, const Buffer& buffer,
+ ErrorReporter* error_reporter) {
+ uint64_t bytes_required = 1;
+ for (int dim : *tensor.shape()) {
+ bytes_required *= dim;
+ if (bytes_required > UINT_MAX) {
+ ReportError(error_reporter, "Tensor dimension overflow");
+ return false;
+ }
+ }
+ switch (tensor.type()) {
+ case TensorType_FLOAT32:
+ bytes_required *= sizeof(float);
+ break;
+ case TensorType_INT32:
+ bytes_required *= sizeof(int32_t);
+ break;
+ case TensorType_UINT8:
+ bytes_required *= sizeof(uint8_t);
+ break;
+ case TensorType_INT64:
+ bytes_required *= sizeof(int64_t);
+ break;
+ case TensorType_FLOAT16:
+ // FALLTHROUGH_INTENDED;
+ default:
+ ReportError(error_reporter, "Invalid tensor type: %d", tensor.type());
+ return false;
+ }
+ if (bytes_required > UINT_MAX) {
+ ReportError(error_reporter, "Tensor dimension overflow");
+ return false;
+ }
+
+ if (bytes_required != buffer.data()->size()) {
+ ReportError(
+ error_reporter,
+ "Tensor requires %d bytes, but is allocated with %d bytes buffer",
+ bytes_required, buffer.data()->size());
+ return false;
+ }
+ return true;
+
+ // TODO(yichengfan): verify quantized tensors.
+}
+
+// Verifies tensors have valid properties and legit buffer if set.
+bool VerifyTensors(const Model& model, ErrorReporter* error_reporter) {
+ if (!model.subgraphs()) {
+ return true;
+ }
+ for (const auto& subgraph : *model.subgraphs()) {
+ if (!subgraph->tensors()) {
+ return true;
+ }
+ for (const auto& tensor : *subgraph->tensors()) {
+ if (!tensor->buffer()) {
+ return true;
+ }
+ if (tensor->buffer() >= model.buffers()->size()) {
+ ReportError(error_reporter, "Invalid tensor buffer index: %d",
+ tensor->buffer());
+ return false;
+ }
+ auto* buffer = model.buffers()->Get(tensor->buffer());
+ if (!buffer || !buffer->data()) {
+ ReportError(error_reporter, "Tensor buffer %d not set",
+ tensor->buffer());
+ return false;
+ }
+
+ if (tensor->type() == TensorType_STRING) {
+ if (!VerifyStringTensorBuffer(*buffer, error_reporter)) {
+ return false;
+ }
+ } else {
+ if (!VerifyNumericTensorBuffer(*tensor, *buffer, error_reporter)) {
+ return false;
+ }
+ }
+ }
+ }
+ return true;
+}
+
+} // namespace
+
+bool Verify(const void* buf, size_t len, ErrorReporter* error_reporter) {
+ const Model* model = VerifyFlatbufferAndGetModel(buf, len);
+ if (model == nullptr) {
+ ReportError(error_reporter, "Invalid flatbuffer format");
+ return false;
+ }
+ if (model->version() != TFLITE_SCHEMA_VERSION) {
+ ReportError(error_reporter, "Invalid model version %d", model->version());
+ return false;
+ }
+ if (!VerifyTensors(*model, error_reporter)) {
+ return false;
+ }
+ return true;
}
} // namespace tflite
diff --git a/tensorflow/contrib/lite/tools/verifier.h b/tensorflow/contrib/lite/tools/verifier.h
index 03e1f22..d2bf3c9 100644
--- a/tensorflow/contrib/lite/tools/verifier.h
+++ b/tensorflow/contrib/lite/tools/verifier.h
@@ -18,13 +18,15 @@
#include <stdio.h>
+#include "tensorflow/contrib/lite/error_reporter.h"
+
namespace tflite {
// Verifies the integrity of a Tensorflow Lite flatbuffer model file.
// Currently, it verifies:
// * The file is following a legit flatbuffer schema.
// * The model is in supported version.
-bool Verify(const void* buf, size_t len);
+bool Verify(const void* buf, size_t len, ErrorReporter* error_reporter);
} // namespace tflite
diff --git a/tensorflow/contrib/lite/tools/verifier_test.cc b/tensorflow/contrib/lite/tools/verifier_test.cc
index 0481a55..244d4f0 100644
--- a/tensorflow/contrib/lite/tools/verifier_test.cc
+++ b/tensorflow/contrib/lite/tools/verifier_test.cc
@@ -28,31 +28,62 @@
using flatbuffers::Offset;
using flatbuffers::Vector;
-// Class that abstracts the list of buffers at the end of the TF Lite structure
-class DeferredBufferWriter {
+// Build single subgraph model.
+class TfLiteFlatbufferModelBuilder {
public:
- DeferredBufferWriter() {
- data_.push_back({}); // sentinel empty buffer.
+ TfLiteFlatbufferModelBuilder() {
+ buffers_.push_back(
+ CreateBuffer(builder_, builder_.CreateVector(std::vector<uint8_t>{})));
}
- Offset<Vector<Offset<Buffer>>> BuildBuffers(FlatBufferBuilder *builder) {
- std::vector<Offset<Buffer>> buffer_vector;
- for (const auto &vec : data_) {
- auto data_buffer = builder->CreateVector(vec.data(), vec.size());
- buffer_vector.push_back(tflite::CreateBuffer(*builder, data_buffer));
+ void AddTensor(const std::vector<int>& shape, tflite::TensorType type,
+ const std::vector<uint8_t>& buffer, const char* name) {
+ int buffer_index = 0;
+ if (!buffer.empty()) {
+ buffer_index = buffers_.size();
+ buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector(buffer)));
}
- return builder->CreateVector(buffer_vector);
+ tensors_.push_back(CreateTensorDirect(builder_, &shape, type, buffer_index,
+ name, /*quantization=*/0));
}
- // Registers a buffer index and takes ownership of the data to write to it.
- int Record(std::vector<uint8_t> data) {
- int buffer_index = data_.size();
- data_.emplace_back(std::move(data));
- return buffer_index;
+ void AddOperator(const std::vector<int32_t>& inputs,
+ const std::vector<int32_t>& outputs,
+ tflite::BuiltinOperator builtin_op, const char* custom_op) {
+ operator_codes_.push_back(
+ CreateOperatorCodeDirect(builder_, builtin_op, custom_op));
+ operators_.push_back(CreateOperator(
+ builder_, operator_codes_.size() - 1, builder_.CreateVector(inputs),
+ builder_.CreateVector(outputs), BuiltinOptions_NONE,
+ /*builtin_options=*/0,
+ /*custom_options=*/0, tflite::CustomOptionsFormat_FLEXBUFFERS));
+ }
+
+ void FinishModel(const std::vector<int32_t>& inputs,
+ const std::vector<int32_t>& outputs) {
+ auto subgraph = std::vector<Offset<SubGraph>>({CreateSubGraph(
+ builder_, builder_.CreateVector(tensors_),
+ builder_.CreateVector(inputs), builder_.CreateVector(outputs),
+ builder_.CreateVector(operators_),
+ builder_.CreateString("test_subgraph"))});
+ auto result = CreateModel(
+ builder_, TFLITE_SCHEMA_VERSION, builder_.CreateVector(operator_codes_),
+ builder_.CreateVector(subgraph), builder_.CreateString("test_model"),
+ builder_.CreateVector(buffers_));
+ tflite::FinishModelBuffer(builder_, result);
+ }
+
+ bool Verify() {
+ return tflite::Verify(builder_.GetBufferPointer(), builder_.GetSize(),
+ DefaultErrorReporter());
}
private:
- std::vector<std::vector<unsigned char>> data_;
+ FlatBufferBuilder builder_;
+ std::vector<Offset<Operator>> operators_;
+ std::vector<Offset<OperatorCode>> operator_codes_;
+ std::vector<Offset<Tensor>> tensors_;
+ std::vector<Offset<Buffer>> buffers_;
};
TEST(VerifyModel, TestEmptyModel) {
@@ -62,43 +93,26 @@
/*description=*/0, /*buffers=*/0);
::tflite::FinishModelBuffer(builder, model);
- ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize()));
+ ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(),
+ DefaultErrorReporter()));
}
TEST(VerifyModel, TestSimpleModel) {
- FlatBufferBuilder builder;
- auto inputs = builder.CreateVector<int32_t>({0});
- auto outputs = builder.CreateVector<int32_t>({1});
- auto operator_codes = builder.CreateVector(std::vector<Offset<OperatorCode>>{
- CreateOperatorCodeDirect(builder, BuiltinOperator_CUSTOM, "test")});
- auto operators =
- builder.CreateVector(std::vector<Offset<Operator>>{CreateOperator(
- builder, /*opcode_index=*/0,
- /*inputs=*/builder.CreateVector<int32_t>({0}),
- /*outputs=*/builder.CreateVector<int32_t>({1}), BuiltinOptions_NONE,
- /*builtin_options=*/0,
- /*custom_options=*/0, ::tflite::CustomOptionsFormat_FLEXBUFFERS)});
- std::vector<int> shape;
- auto tensors = builder.CreateVector(std::vector<Offset<Tensor>>{
- CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/0,
- "input", /*quantization=*/0),
- CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/0,
- "output", /*quantization=*/0)});
- auto subgraph = std::vector<Offset<SubGraph>>(
- {CreateSubGraph(builder, tensors, inputs, outputs, operators,
- builder.CreateString("Main"))});
-
- auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, operator_codes,
- builder.CreateVector(subgraph),
- builder.CreateString("SmartReply"), /*buffers=*/0);
-
- ::tflite::FinishModelBuffer(builder, model);
- ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize()));
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "test");
+ builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4, 5, 6}, "input");
+ builder.AddTensor(
+ {2}, TensorType_STRING,
+ {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 19, 0, 0, 0, 'A', 'B', 'C'},
+ "data");
+ builder.AddTensor({2, 3}, TensorType_INT32, {}, "output");
+ builder.FinishModel({0, 1}, {2});
+ ASSERT_TRUE(builder.Verify());
}
TEST(VerifyModel, TestCorruptedData) {
string model = "123";
- ASSERT_FALSE(Verify(model.data(), model.size()));
+ ASSERT_FALSE(Verify(model.data(), model.size(), /*error_reporter=*/nullptr));
}
TEST(VerifyModel, TestUnsupportedVersion) {
@@ -106,7 +120,8 @@
auto model = CreateModel(builder, /*version=*/1, /*operator_codes=*/0,
/*subgraphs=*/0, /*description=*/0, /*buffers=*/0);
::tflite::FinishModelBuffer(builder, model);
- ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize()));
+ ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(),
+ DefaultErrorReporter()));
}
TEST(VerifyModel, TestRandomModificationIsNotAllowed) {
@@ -116,20 +131,105 @@
/*subgraphs=*/0, /*description=*/0, /*buffers=*/0);
::tflite::FinishModelBuffer(builder, model);
- string model_content(reinterpret_cast<char *>(builder.GetBufferPointer()),
+ string model_content(reinterpret_cast<char*>(builder.GetBufferPointer()),
builder.GetSize());
for (int i = 0; i < model_content.size(); i++) {
model_content[i] = (model_content[i] + 137) % 255;
- EXPECT_FALSE(Verify(model_content.data(), model_content.size()))
+ EXPECT_FALSE(Verify(model_content.data(), model_content.size(),
+ DefaultErrorReporter()))
<< "Fail at position: " << i;
}
}
+TEST(VerifyModel, TestIntTensorShapeIsGreaterThanBuffer) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, TestIntTensorShapeIsSmallerThanBuffer) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor({2, 1}, TensorType_UINT8, {1, 2, 3, 4}, "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, TestIntTensorShapeOverflow) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor({1024, 2048, 4096}, TensorType_UINT8, {1, 2, 3, 4},
+ "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, TensorBufferIsNotValid) {
+ FlatBufferBuilder builder;
+ std::vector<int> shape = {2, 3};
+ auto tensors = builder.CreateVector(std::vector<Offset<Tensor>>{
+ CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/2,
+ "input", /*quantization=*/0)});
+ auto subgraph = std::vector<Offset<SubGraph>>(
+ {CreateSubGraph(builder, tensors, /*inputs=*/0, /*outputs=*/0,
+ /*operators=*/0, builder.CreateString("Main"))});
+
+ auto buffers = builder.CreateVector(std::vector<Offset<Buffer>>{
+ CreateBuffer(builder,
+ builder.CreateVector(std::vector<uint8>{1, 2, 3, 4, 5, 6})),
+ });
+
+ auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, /*operator_codes=*/0,
+ builder.CreateVector(subgraph),
+ builder.CreateString("SmartReply"), buffers);
+
+ ::tflite::FinishModelBuffer(builder, model);
+ ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(),
+ DefaultErrorReporter()));
+}
+
+TEST(VerifyModel, StringTensorHasInvalidNumString) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor(
+ {2}, TensorType_STRING,
+ {0x00, 0x00, 0x00, 0x20, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'},
+ "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, StringTensorOffsetTooSmall) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor(
+ {2}, TensorType_STRING,
+ {2, 0, 0, 0, 12, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, StringTensorOffsetOutOfRange) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor(
+ {2}, TensorType_STRING,
+ {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 22, 0, 0, 0, 'A', 'B'}, "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
+TEST(VerifyModel, StringTensorIsLargerThanRequired) {
+ TfLiteFlatbufferModelBuilder builder;
+ builder.AddTensor(
+ {2}, TensorType_STRING,
+ {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B', 'C'},
+ "input");
+ builder.FinishModel({}, {});
+ ASSERT_FALSE(builder.Verify());
+}
+
// TODO(yichengfan): make up malicious files to test with.
} // namespace tflite
-int main(int argc, char **argv) {
+int main(int argc, char** argv) {
::tflite::LogToStderr();
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();