1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15#include <string> 16#include <vector> 17 18#include "flatbuffers/flatbuffers.h" 19#include "flatbuffers/util.h" 20#include <gtest/gtest.h> 21#include "tensorflow/contrib/lite/allocation.h" 22#include "tensorflow/contrib/lite/error_reporter.h" 23#include "tensorflow/contrib/lite/schema/schema_generated.h" 24#include "tensorflow/contrib/lite/testing/util.h" 25#include "tensorflow/contrib/lite/tools/mutable_op_resolver.h" 26#include "tensorflow/contrib/lite/tools/verifier.h" 27#include "tensorflow/contrib/lite/version.h" 28#include "tensorflow/core/framework/numeric_types.h" 29 30namespace tflite { 31 32using flatbuffers::FlatBufferBuilder; 33using flatbuffers::Offset; 34using flatbuffers::Vector; 35 36// Build single subgraph model. 37class TfLiteFlatbufferModelBuilder { 38 public: 39 TfLiteFlatbufferModelBuilder() { 40 buffers_.push_back( 41 CreateBuffer(builder_, builder_.CreateVector(std::vector<uint8_t>{}))); 42 } 43 44 TfLiteFlatbufferModelBuilder(const std::vector<BuiltinOperator>& builtin_ops, 45 const std::vector<string>& custom_ops) { 46 buffers_.push_back( 47 CreateBuffer(builder_, builder_.CreateVector(std::vector<uint8_t>{}))); 48 49 for (const auto& iter : builtin_ops) { 50 resolver_.AddBuiltin(iter, &fake_op_); 51 } 52 for (const auto& iter : custom_ops) { 53 resolver_.AddCustom(iter.data(), &fake_op_); 54 } 55 } 56 57 void AddTensor(const std::vector<int>& shape, tflite::TensorType type, 58 const std::vector<uint8_t>& buffer, const char* name) { 59 int buffer_index = 0; 60 if (!buffer.empty()) { 61 buffer_index = buffers_.size(); 62 buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector(buffer))); 63 } 64 tensors_.push_back(CreateTensorDirect(builder_, &shape, type, buffer_index, 65 name, /*quantization=*/0)); 66 } 67 68 void AddOperator(const std::vector<int32_t>& inputs, 69 const std::vector<int32_t>& outputs, 70 tflite::BuiltinOperator builtin_op, const char* custom_op) { 71 operator_codes_.push_back( 72 CreateOperatorCodeDirect(builder_, builtin_op, custom_op)); 73 operators_.push_back(CreateOperator( 74 builder_, operator_codes_.size() - 1, builder_.CreateVector(inputs), 75 builder_.CreateVector(outputs), BuiltinOptions_NONE, 76 /*builtin_options=*/0, 77 /*custom_options=*/0, tflite::CustomOptionsFormat_FLEXBUFFERS)); 78 } 79 80 void FinishModel(const std::vector<int32_t>& inputs, 81 const std::vector<int32_t>& outputs) { 82 auto subgraph = std::vector<Offset<SubGraph>>({CreateSubGraph( 83 builder_, builder_.CreateVector(tensors_), 84 builder_.CreateVector(inputs), builder_.CreateVector(outputs), 85 builder_.CreateVector(operators_), 86 builder_.CreateString("test_subgraph"))}); 87 auto result = CreateModel( 88 builder_, TFLITE_SCHEMA_VERSION, builder_.CreateVector(operator_codes_), 89 builder_.CreateVector(subgraph), builder_.CreateString("test_model"), 90 builder_.CreateVector(buffers_)); 91 tflite::FinishModelBuffer(builder_, result); 92 } 93 94 bool Verify() { 95 return tflite::Verify(builder_.GetBufferPointer(), builder_.GetSize(), 96 resolver_, DefaultErrorReporter()); 97 } 98 99 private: 100 FlatBufferBuilder builder_; 101 MutableOpResolver resolver_; 102 TfLiteRegistration fake_op_; 103 std::vector<Offset<Operator>> operators_; 104 std::vector<Offset<OperatorCode>> operator_codes_; 105 std::vector<Offset<Tensor>> tensors_; 106 std::vector<Offset<Buffer>> buffers_; 107}; 108 109TEST(VerifyModel, TestEmptyModel) { 110 FlatBufferBuilder builder; 111 auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, 112 /*operator_codes=*/0, /*subgraphs=*/0, 113 /*description=*/0, /*buffers=*/0); 114 ::tflite::FinishModelBuffer(builder, model); 115 116 ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(), 117 MutableOpResolver{}, DefaultErrorReporter())); 118} 119 120TEST(VerifyModel, TestSimpleModel) { 121 TfLiteFlatbufferModelBuilder builder({}, {"test"}); 122 builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "test"); 123 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4, 5, 6}, "input"); 124 builder.AddTensor( 125 {2}, TensorType_STRING, 126 {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 19, 0, 0, 0, 'A', 'B', 'C'}, 127 "data"); 128 builder.AddTensor({2, 3}, TensorType_INT32, {}, "output"); 129 builder.FinishModel({0, 1}, {2}); 130 ASSERT_TRUE(builder.Verify()); 131} 132 133TEST(VerifyModel, TestCorruptedData) { 134 std::string model = "123"; 135 ASSERT_FALSE(Verify(model.data(), model.size(), MutableOpResolver{}, 136 /*error_reporter=*/nullptr)); 137} 138 139TEST(VerifyModel, TestUnsupportedVersion) { 140 FlatBufferBuilder builder; 141 auto model = CreateModel(builder, /*version=*/1, /*operator_codes=*/0, 142 /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); 143 ::tflite::FinishModelBuffer(builder, model); 144 ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), 145 MutableOpResolver{}, DefaultErrorReporter())); 146} 147 148TEST(VerifyModel, TestRandomModificationIsNotAllowed) { 149 FlatBufferBuilder builder; 150 auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, 151 /*operator_codes=*/0, 152 /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); 153 ::tflite::FinishModelBuffer(builder, model); 154 155 std::string model_content(reinterpret_cast<char*>(builder.GetBufferPointer()), 156 builder.GetSize()); 157 for (int i = 0; i < model_content.size(); i++) { 158 model_content[i] = (model_content[i] + 137) % 255; 159 EXPECT_FALSE(Verify(model_content.data(), model_content.size(), 160 MutableOpResolver{}, DefaultErrorReporter())) 161 << "Fail at position: " << i; 162 } 163} 164 165TEST(VerifyModel, TestIntTensorShapeIsGreaterThanBuffer) { 166 TfLiteFlatbufferModelBuilder builder; 167 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input"); 168 builder.FinishModel({}, {}); 169 ASSERT_FALSE(builder.Verify()); 170} 171 172TEST(VerifyModel, TestIntTensorShapeIsSmallerThanBuffer) { 173 TfLiteFlatbufferModelBuilder builder; 174 builder.AddTensor({2, 1}, TensorType_UINT8, {1, 2, 3, 4}, "input"); 175 builder.FinishModel({}, {}); 176 ASSERT_FALSE(builder.Verify()); 177} 178 179TEST(VerifyModel, TestIntTensorShapeOverflow) { 180 TfLiteFlatbufferModelBuilder builder; 181 builder.AddTensor({1024, 2048, 4096}, TensorType_UINT8, {1, 2, 3, 4}, 182 "input"); 183 builder.FinishModel({}, {}); 184 ASSERT_FALSE(builder.Verify()); 185} 186 187TEST(VerifyModel, TensorBufferIsNotValid) { 188 FlatBufferBuilder builder; 189 std::vector<int> shape = {2, 3}; 190 auto tensors = builder.CreateVector(std::vector<Offset<Tensor>>{ 191 CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/2, 192 "input", /*quantization=*/0)}); 193 auto subgraph = std::vector<Offset<SubGraph>>( 194 {CreateSubGraph(builder, tensors, /*inputs=*/0, /*outputs=*/0, 195 /*operators=*/0, builder.CreateString("Main"))}); 196 197 auto buffers = builder.CreateVector(std::vector<Offset<Buffer>>{ 198 CreateBuffer(builder, 199 builder.CreateVector(std::vector<uint8>{1, 2, 3, 4, 5, 6})), 200 }); 201 202 auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, /*operator_codes=*/0, 203 builder.CreateVector(subgraph), 204 builder.CreateString("SmartReply"), buffers); 205 206 ::tflite::FinishModelBuffer(builder, model); 207 ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), 208 MutableOpResolver{}, DefaultErrorReporter())); 209} 210 211TEST(VerifyModel, StringTensorHasInvalidNumString) { 212 TfLiteFlatbufferModelBuilder builder; 213 builder.AddTensor( 214 {2}, TensorType_STRING, 215 {0x00, 0x00, 0x00, 0x20, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, 216 "input"); 217 builder.FinishModel({}, {}); 218 ASSERT_FALSE(builder.Verify()); 219} 220 221TEST(VerifyModel, StringTensorOffsetTooSmall) { 222 TfLiteFlatbufferModelBuilder builder; 223 builder.AddTensor( 224 {2}, TensorType_STRING, 225 {2, 0, 0, 0, 12, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, "input"); 226 builder.FinishModel({}, {}); 227 ASSERT_FALSE(builder.Verify()); 228} 229 230TEST(VerifyModel, StringTensorOffsetOutOfRange) { 231 TfLiteFlatbufferModelBuilder builder; 232 builder.AddTensor( 233 {2}, TensorType_STRING, 234 {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 22, 0, 0, 0, 'A', 'B'}, "input"); 235 builder.FinishModel({}, {}); 236 ASSERT_FALSE(builder.Verify()); 237} 238 239TEST(VerifyModel, StringTensorIsLargerThanRequired) { 240 TfLiteFlatbufferModelBuilder builder; 241 builder.AddTensor( 242 {2}, TensorType_STRING, 243 {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B', 'C'}, 244 "input"); 245 builder.FinishModel({}, {}); 246 ASSERT_FALSE(builder.Verify()); 247} 248 249TEST(VerifyModel, AllOpsAreSupported) { 250 TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"CustomOp"}); 251 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); 252 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); 253 builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); 254 builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); 255 builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "CustomOp"); 256 builder.FinishModel({}, {}); 257 ASSERT_FALSE(builder.Verify()); 258} 259 260TEST(VerifyModel, UseUnsupportedBuiltinOps) { 261 TfLiteFlatbufferModelBuilder builder({BuiltinOperator_SUB}, {"CustomOp"}); 262 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); 263 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); 264 builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); 265 builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); 266 builder.FinishModel({}, {}); 267 ASSERT_FALSE(builder.Verify()); 268} 269 270TEST(VerifyModel, UseUnsupportedCustomOps) { 271 TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"NewOp"}); 272 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); 273 builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); 274 builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); 275 builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "Not supported"); 276 builder.FinishModel({}, {}); 277 ASSERT_FALSE(builder.Verify()); 278} 279 280// TODO(yichengfan): make up malicious files to test with. 281 282} // namespace tflite 283 284int main(int argc, char** argv) { 285 ::tflite::LogToStderr(); 286 ::testing::InitGoogleTest(&argc, argv); 287 return RUN_ALL_TESTS(); 288} 289