1 /* 2 * Copyright (C) 2021 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef SRC_PROTOZERO_FILTERING_MESSAGE_FILTER_H_ 18 #define SRC_PROTOZERO_FILTERING_MESSAGE_FILTER_H_ 19 20 #include <stdint.h> 21 22 #include <memory> 23 #include <string> 24 #include <unordered_map> 25 26 #include "src/protozero/filtering/filter_bytecode_parser.h" 27 #include "src/protozero/filtering/message_tokenizer.h" 28 29 namespace protozero { 30 31 // A class to filter binary-encoded proto messages using an allow-list of field 32 // ids, also known as "filter bytecode". The filter determines which fields are 33 // allowed to be passed through in output and strips all the other fields. 34 // See go/trace-filtering for full design. 35 // This class takes in input: 36 // 1) The filter bytecode, loaded once via the LoadFilterBytecode() method. 37 // 2) A proto-encoded binary message. The message doesn't have to be contiguous, 38 // it can be passed as an array of arbitrarily chunked fragments. 39 // The FilterMessage*() method returns in output a proto message, stripping out 40 // all unknown fields. If the input is malformed (e.g., unknown proto field wire 41 // types, lengths out of bound) the whole filtering failed and the |error| flag 42 // of the FilteredMessage object is set to true. 43 // The filtering operation is based on rewriting a copy of the message into a 44 // self-allocated buffer, which is then returned in the output. The input buffer 45 // is NOT altered. 46 // Note also that the process of rewriting the protos gets rid of most redundant 47 // varint encoding (if present). So even if all fields are allow-listed, the 48 // output might NOT be bitwise identical to the input (but it will be 49 // semantically equivalent). 50 // Furthermore the enable_field_usage_tracking() method allows to keep track of 51 // a histogram of allowed / denied fields. It slows down filtering and is 52 // intended only on host tools. 53 class MessageFilter { 54 public: 55 MessageFilter(); 56 ~MessageFilter(); 57 58 struct InputSlice { 59 const void* data; 60 size_t len; 61 }; 62 63 struct FilteredMessage { FilteredMessageFilteredMessage64 FilteredMessage(std::unique_ptr<uint8_t[]> d, size_t s) 65 : data(std::move(d)), size(s) {} 66 std::unique_ptr<uint8_t[]> data; 67 size_t size; // The used bytes in |data|. This is <= sizeof(data). 68 bool error = false; 69 }; 70 71 // Loads the filter bytecode that will be used to filter any subsequent 72 // message. Must be called before the first call to FilterMessage*(). 73 // |filter_data| must point to a byte buffer for a proto-encoded ProtoFilter 74 // message (see proto_filter.proto). 75 bool LoadFilterBytecode(const void* filter_data, size_t len); 76 77 // This affects the filter starting point of the subsequent FilterMessage*() 78 // calls. By default the filtering process starts from the message @ index 0, 79 // the root message passed to proto_filter when generating the bytecode 80 // (in typical tracing use-cases, this is perfetto.protos.Trace). However, the 81 // caller (TracingServiceImpl) might want to filter packets from the 2nd level 82 // (perfetto.protos.TracePacket) because the root level is pre-pended after 83 // the fact. This call allows to change the root message for the filter. 84 // The argument |field_ids| is an array of proto field ids and determines the 85 // path to the new root. For instance, in the case of [1,2,3] SetFilterRoot 86 // will identify the sub-message for the field "root.1.2.3" and use that. 87 // In order for this to succeed all the fields in the path must be allowed 88 // in the filter and must be a nested message type. 89 bool SetFilterRoot(const uint32_t* field_ids, size_t num_fields); 90 91 // Takes an input message, fragmented in arbitrary slices, and returns a 92 // filtered message in output. 93 FilteredMessage FilterMessageFragments(const InputSlice*, size_t num_slices); 94 95 // Helper for tests, where the input is a contiguous buffer. FilterMessage(const void * data,size_t len)96 FilteredMessage FilterMessage(const void* data, size_t len) { 97 InputSlice slice{data, len}; 98 return FilterMessageFragments(&slice, 1); 99 } 100 101 // When enabled returns a map of "field path" to "usage counter". 102 // The key (std::string) is a binary buffer (i.e. NOT an ASCII/UTF-8 string) 103 // which contains a varint for each field. Consider the following: 104 // message Root { Sub1 f1 = 1; }; 105 // message Sub1 { Sub2 f2 = 7;} 106 // message Sub2 { string f3 = 5; } 107 // The field .f1.f2.f3 will be encoded as \x01\0x07\x05. 108 // The value is the number of times that field has been encountered. If the 109 // field is not allow-listed in the bytecode (the field is stripped in output) 110 // the count will be negative. enable_field_usage_tracking(bool x)111 void enable_field_usage_tracking(bool x) { track_field_usage_ = x; } field_usage()112 const std::unordered_map<std::string, int32_t>& field_usage() const { 113 return field_usage_; 114 } 115 116 // Exposed only for DCHECKS in TracingServiceImpl. root_msg_index()117 uint32_t root_msg_index() { return root_msg_index_; } 118 119 private: 120 // This is called by FilterMessageFragments(). 121 // Inlining allows the compiler turn the per-byte call/return into a for loop, 122 // while, at the same time, keeping the code easy to read and reason about. 123 // It gives a 20-25% speedup (265ms vs 215ms for a 25MB trace). 124 void FilterOneByte(uint8_t octet) PERFETTO_ALWAYS_INLINE; 125 126 // No-inline because this is a slowpath (only when usage tracking is enabled). 127 void IncrementCurrentFieldUsage(uint32_t field_id, 128 bool allowed) PERFETTO_NO_INLINE; 129 130 // Gets into an error state which swallows all the input and emits no output. 131 void SetUnrecoverableErrorState(); 132 133 // We keep track of the the nest of messages in a stack. Each StackState 134 // object corresponds to a level of nesting in the proto message structure. 135 // Every time a new field of type len-delimited that has a corresponding 136 // sub-message in the bytecode is encountered, a new StackState is pushed in 137 // |stack_|. stack_[0] is a sentinel to prevent over-popping without adding 138 // extra branches in the fastpath. 139 // |stack_|. stack_[1] is the state of the root message. 140 struct StackState { 141 uint32_t in_bytes = 0; // Number of input bytes processed. 142 143 // When |in_bytes| reaches this value, the current state should be popped. 144 // This is set when recursing into nested submessages. This is 0 only for 145 // stack_[0] (we don't know the size of the root message upfront). 146 uint32_t in_bytes_limit = 0; 147 148 // This is set when a len-delimited message is encountered, either a string 149 // or a nested submessage that is NOT allow-listed in the bytecode. 150 // This causes input bytes to be consumed without being parsed from the 151 // input stream. If |passthrough_eaten_bytes| == true, they will be copied 152 // as-is in output (e.g. in the case of an allowed string/bytes field). 153 uint32_t eat_next_bytes = 0; 154 155 // Keeps tracks of the stream_writer output counter (out_.written()) then 156 // the StackState is pushed. This is used to work out, when popping, how 157 // many bytes have been written for the current submessage. 158 uint32_t out_bytes_written_at_start = 0; 159 160 uint32_t field_id = 0; // The proto field id for the current message. 161 uint32_t msg_index = 0; // The index of the message filter in the bytecode. 162 163 // This is a pointer to the proto preamble for the current submessage 164 // (it's nullptr for stack_[0] and non-null elsewhere). This will be filled 165 // with the actual size of the message (out_.written() - 166 // |out_bytes_written_at_start|) when finishing (popping) the message. 167 // This must be filled using WriteRedundantVarint(). Note that the 168 // |size_field_len| is variable and depends on the actual length of the 169 // input message. If the output message has roughly the same size of the 170 // input message, the length will not be redundant. 171 // In other words: the length of the field is reserved when the submessage 172 // starts. At that point we know the upper-bound for the output message 173 // (a filtered submessage can be <= the original one, but not >). So we 174 // reserve as many bytes it takes to write the input length in varint. 175 // Then, when the message is finalized and we know the actual output size 176 // we backfill the field. 177 // Consider the example of a submessage where the input size = 130 (>127, 178 // 2 varint bytes) and the output is 120 bytes. The length will be 2 bytes 179 // wide even though could have been encoded with just one byte. 180 uint8_t* size_field = nullptr; 181 uint32_t size_field_len = 0; 182 183 // When true the next |eat_next_bytes| are copied as-is in output. 184 // It seems that keeping this field at the end rather than next to 185 // |eat_next_bytes| makes the filter a little (but measurably) faster. 186 // (likely something related with struct layout vs cache sizes). 187 bool passthrough_eaten_bytes = false; 188 }; 189 out_written()190 uint32_t out_written() { return static_cast<uint32_t>(out_ - &out_buf_[0]); } 191 192 std::unique_ptr<uint8_t[]> out_buf_; 193 uint8_t* out_ = nullptr; 194 uint8_t* out_end_ = nullptr; 195 uint32_t root_msg_index_ = 0; 196 197 FilterBytecodeParser filter_; 198 MessageTokenizer tokenizer_; 199 std::vector<StackState> stack_; 200 201 bool error_ = false; 202 bool track_field_usage_ = false; 203 std::unordered_map<std::string, int32_t> field_usage_; 204 }; 205 206 } // namespace protozero 207 208 #endif // SRC_PROTOZERO_FILTERING_MESSAGE_FILTER_H_ 209