• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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