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 #ifndef TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_ 16 #define TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_ 17 18 #include <cstdint> 19 #include <cstdlib> 20 #include <string> 21 #include <vector> 22 23 #include "tensorflow/cc/framework/scope.h" 24 #include "tensorflow/core/framework/tensor.h" 25 #include "tensorflow/core/platform/status.h" 26 #include "tensorflow/core/public/session.h" 27 #include "tensorflow/core/public/session_options.h" 28 29 // Standard invoking function macro to dispatch to a fuzzer class. 30 #define STANDARD_TF_FUZZ_FUNCTION(FuzzerClass) \ 31 extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { \ 32 static FuzzerClass* fuzzer = new FuzzerClass(); \ 33 return fuzzer->Fuzz(data, size); \ 34 } 35 36 // Standard builder for hooking one placeholder to one op. 37 #define SINGLE_INPUT_OP_BUILDER(dtype, opName) \ 38 void BuildGraph(const Scope& scope) override { \ 39 auto op_node = \ 40 tensorflow::ops::Placeholder(scope.WithOpName("input"), dtype); \ 41 (void)tensorflow::ops::opName(scope.WithOpName("output"), op_node); \ 42 } 43 44 namespace tensorflow { 45 namespace fuzzing { 46 47 // Create a TensorFlow session using a specific GraphDef created 48 // by BuildGraph(), and make it available for fuzzing. 49 // Users must override BuildGraph and FuzzImpl to specify 50 // (1) which operations are being fuzzed; and 51 // (2) How to translate the uint8_t* buffer from the fuzzer 52 // to a Tensor or Tensors that are semantically appropriate 53 // for the op under test. 54 // For the simple cases of testing a single op that takes a single 55 // input Tensor, use the SINGLE_INPUT_OP_BUILDER(dtype, opName) macro in place 56 // of defining BuildGraphDef. 57 // 58 // Typical use: 59 // class FooFuzzer : public FuzzSession { 60 // SINGLE_INPUT_OP_BUILDER(DT_INT8, Identity); 61 // void FuzzImpl(const uint8_t* data, size_t size) { 62 // ... convert data and size to a Tensor, pass it to: 63 // RunInputs({{"input", input_tensor}}); 64 // 65 class FuzzSession { 66 public: FuzzSession()67 FuzzSession() : initialized_(false) {} ~FuzzSession()68 virtual ~FuzzSession() {} 69 70 // Constructs a Graph using the supplied Scope. 71 // By convention, the graph should have inputs named "input1", ... 72 // "inputN", and one output node, named "output". 73 // Users of FuzzSession should override this method to create their graph. 74 virtual void BuildGraph(const Scope& scope) = 0; 75 76 // Implements the logic that converts an opaque byte buffer 77 // from the fuzzer to Tensor inputs to the graph. Users must override. 78 virtual void FuzzImpl(const uint8_t* data, size_t size) = 0; 79 80 // Initializes the FuzzSession. Not safe for multithreading. 81 // Separate init function because the call to virtual BuildGraphDef 82 // can't be put into the constructor. InitIfNeeded()83 Status InitIfNeeded() { 84 if (initialized_) { 85 return OkStatus(); 86 } 87 initialized_ = true; 88 89 Scope root = Scope::DisabledShapeInferenceScope().ExitOnError(); 90 SessionOptions options; 91 session_ = std::unique_ptr<Session>(NewSession(options)); 92 93 BuildGraph(root); 94 95 GraphDef graph_def; 96 TF_CHECK_OK(root.ToGraphDef(&graph_def)); 97 98 Status status = session_->Create(graph_def); 99 if (!status.ok()) { 100 // This is FATAL, because this code is designed to fuzz an op 101 // within a session. Failure to create the session means we 102 // can't send any data to the op. 103 LOG(FATAL) << "Could not create session: " << status.error_message(); 104 } 105 return status; 106 } 107 108 // Runs the TF session by pulling on the "output" node, attaching 109 // the supplied input_tensor to the input node(s), and discarding 110 // any returned output. 111 // Note: We are ignoring Status from Run here since fuzzers don't need to 112 // check it (as that will slow them down and printing/logging is useless). RunInputs(const std::vector<std::pair<string,Tensor>> & inputs)113 void RunInputs(const std::vector<std::pair<string, Tensor> >& inputs) { 114 RunInputsWithStatus(inputs).IgnoreError(); 115 } 116 117 // Same as RunInputs but don't ignore status RunInputsWithStatus(const std::vector<std::pair<string,Tensor>> & inputs)118 Status RunInputsWithStatus( 119 const std::vector<std::pair<string, Tensor> >& inputs) { 120 return session_->Run(inputs, {}, {"output"}, nullptr); 121 } 122 123 // Dispatches to FuzzImpl; small amount of sugar to keep the code 124 // of the per-op fuzzers tiny. Fuzz(const uint8_t * data,size_t size)125 int Fuzz(const uint8_t* data, size_t size) { 126 Status status = InitIfNeeded(); 127 TF_CHECK_OK(status) << "Fuzzer graph initialization failed: " 128 << status.error_message(); 129 // No return value from fuzzing: Success is defined as "did not 130 // crash". The actual application results are irrelevant. 131 FuzzImpl(data, size); 132 return 0; 133 } 134 135 private: 136 bool initialized_; 137 std::unique_ptr<Session> session_; 138 }; 139 140 // A specialized fuzz implementation for ops that take 141 // a single string. Caller must still define the op 142 // to plumb by overriding BuildGraph or using 143 // a plumbing macro. 144 class FuzzStringInputOp : public FuzzSession { FuzzImpl(const uint8_t * data,size_t size)145 void FuzzImpl(const uint8_t* data, size_t size) final { 146 Tensor input_tensor(tensorflow::DT_STRING, TensorShape({})); 147 input_tensor.scalar<tstring>()() = 148 string(reinterpret_cast<const char*>(data), size); 149 RunInputs({{"input", input_tensor}}); 150 } 151 }; 152 153 } // end namespace fuzzing 154 } // end namespace tensorflow 155 156 #endif // TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_ 157