1 /* Copyright 2016 Google Inc. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_KERNELS_FUZZING_FUZZ_SESSION_H_ 17 #define TENSORFLOW_CORE_KERNELS_FUZZING_FUZZ_SESSION_H_ 18 19 #include "tensorflow/cc/framework/scope.h" 20 #include "tensorflow/core/graph/graph.h" 21 #include "tensorflow/core/public/session.h" 22 23 // Standard invoking function macro to dispatch to a fuzzer class. 24 #ifndef PLATFORM_WINDOWS 25 #define STANDARD_TF_FUZZ_FUNCTION(FuzzerClass) \ 26 extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { \ 27 static FuzzerClass* fuzzer = new FuzzerClass(); \ 28 return fuzzer->Fuzz(data, size); \ 29 } 30 #else 31 // We don't compile this for Windows, MSVC doesn't like it as pywrap in Windows 32 // links all the code into one big object file and there are conflicting 33 // function names. 34 #define STANDARD_TF_FUZZ_FUNCTION(FuzzerClass) 35 #endif 36 37 // Standard builder for hooking one placeholder to one op. 38 #define SINGLE_INPUT_OP_BUILDER(dtype, opName) \ 39 void BuildGraph(const Scope& scope) override { \ 40 auto op_node = \ 41 tensorflow::ops::Placeholder(scope.WithOpName("input1"), dtype); \ 42 std::ignore = \ 43 tensorflow::ops::opName(scope.WithOpName("output"), op_node); \ 44 } 45 46 namespace tensorflow { 47 namespace fuzzing { 48 49 // Create a TensorFlow session using a specific GraphDef created 50 // by BuildGraph(), and make it available for fuzzing. 51 // Users must override BuildGraph and FuzzImpl to specify 52 // (1) which operations are being fuzzed; and 53 // (2) How to translate the uint8_t* buffer from the fuzzer 54 // to a Tensor or Tensors that are semantically appropriate 55 // for the op under test. 56 // For the simple cases of testing a single op that takes a single 57 // input Tensor, use the SINGLE_INPUT_OP_BUILDER(dtype, opName) macro in place 58 // of defining BuildGraphDef. 59 // 60 // Typical use: 61 // class FooFuzzer : public FuzzSession { 62 // SINGLE_INPUT_OP_BUILDER(DT_INT8, Identity); 63 // void FuzzImpl(const uint8_t* data, size_t size) { 64 // ... convert data and size to a Tensor, pass it to: 65 // RunOneInput(input_tensor); 66 // 67 class FuzzSession { 68 public: 69 FuzzSession() : initialized_(false) {} 70 virtual ~FuzzSession() {} 71 72 // Constructs a Graph using the supplied Scope. 73 // By convention, the graph should have inputs named "input1", ... 74 // "inputN", and one output node, named "output". 75 // Users of FuzzSession should override this method to create their graph. 76 virtual void BuildGraph(const Scope& scope) {} 77 78 // Implements the logic that converts an opaque byte buffer 79 // from the fuzzer to Tensor inputs to the graph. Users must override. 80 virtual void FuzzImpl(const uint8_t* data, size_t size) {} 81 82 // Initializes the FuzzSession. Not safe for multithreading. 83 // Separate init function because the call to virtual BuildGraphDef 84 // can't be put into the constructor. 85 Status InitIfNeeded() { 86 if (initialized_) { 87 return Status::OK(); 88 } 89 initialized_ = true; 90 91 Scope root = Scope::DisabledShapeInferenceScope().ExitOnError(); 92 SessionOptions options; 93 session_ = std::unique_ptr<Session>(NewSession(options)); 94 95 BuildGraph(root); 96 97 GraphDef graph_def; 98 TF_CHECK_OK(root.ToGraphDef(&graph_def)); 99 100 Status status = session_->Create(graph_def); 101 if (!status.ok()) { 102 // This is FATAL, because this code is designed to fuzz an op 103 // within a session. Failure to create the session means we 104 // can't send any data to the op. 105 LOG(FATAL) << "Could not create session: " << status.error_message(); 106 } 107 return status; 108 } 109 110 // Runs the TF session by pulling on the "output" node, attaching 111 // the supplied input_tensor to the "input1" node, and discarding 112 // any returned output. 113 Status RunOneInput(const Tensor& input_tensor) { 114 return session_->Run({{"input1", input_tensor}}, {}, {"output"}, nullptr); 115 } 116 117 Status RunTwoInputs(const Tensor& input1, const Tensor& input2) { 118 return session_->Run({{"input1", input1}, {"input2", input2}}, {}, 119 {"output"}, nullptr); 120 } 121 122 // Dispatches to FuzzImpl; small amount of sugar to keep the code 123 // of the per-op fuzzers tiny. 124 int Fuzz(const uint8_t* data, size_t size) { 125 Status status = InitIfNeeded(); 126 TF_CHECK_OK(status) << "Fuzzer graph initialization failed: " 127 << status.error_message(); 128 // No return value from fuzzing: Success is defined as "did not 129 // crash". The actual application results are irrelevant. 130 FuzzImpl(data, size); 131 return 0; 132 } 133 134 private: 135 bool initialized_; 136 std::unique_ptr<Session> session_; 137 }; 138 139 // A specialized fuzz implementation for ops that take 140 // a single string. Caller must still define the op 141 // to plumb by overriding BuildGraph or using 142 // a plumbing macro. 143 class FuzzStringInputOp : public FuzzSession { 144 void FuzzImpl(const uint8_t* data, size_t size) final { 145 Tensor input_tensor(tensorflow::DT_STRING, TensorShape({})); 146 input_tensor.scalar<string>()() = 147 string(reinterpret_cast<const char*>(data), size); 148 // TODO(b/32704451): Don't just ignore the ::tensorflow::Status object! 149 RunOneInput(input_tensor).IgnoreError(); 150 } 151 }; 152 153 } // end namespace fuzzing 154 } // end namespace tensorflow 155 156 #endif // TENSORFLOW_CORE_KERNELS_FUZZING_FUZZ_SESSION_H_ 157
