• 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 #include "src/protozero/filtering/message_filter.h"
18 
19 #include "perfetto/base/logging.h"
20 #include "perfetto/protozero/proto_utils.h"
21 
22 namespace protozero {
23 
24 namespace {
25 
26 // Inline helpers to append proto fields in output. They are the equivalent of
27 // the protozero::Message::AppendXXX() fields but don't require building and
28 // maintaining a full protozero::Message object or dealing with scattered
29 // output slices.
30 // All these functions assume there is enough space in the output buffer, which
31 // should be always the case assuming that we don't end up generating more
32 // output than input.
33 
AppendVarInt(uint32_t field_id,uint64_t value,uint8_t ** out)34 inline void AppendVarInt(uint32_t field_id, uint64_t value, uint8_t** out) {
35   *out = proto_utils::WriteVarInt(proto_utils::MakeTagVarInt(field_id), *out);
36   *out = proto_utils::WriteVarInt(value, *out);
37 }
38 
39 // For fixed32 / fixed64.
40 template <typename INT_T /* uint32_t | uint64_t*/>
AppendFixed(uint32_t field_id,INT_T value,uint8_t ** out)41 inline void AppendFixed(uint32_t field_id, INT_T value, uint8_t** out) {
42   *out = proto_utils::WriteVarInt(proto_utils::MakeTagFixed<INT_T>(field_id),
43                                   *out);
44   memcpy(*out, &value, sizeof(value));
45   *out += sizeof(value);
46 }
47 
48 // For length-delimited (string, bytes) fields. Note: this function appends only
49 // the proto preamble and the varint field that states the length of the payload
50 // not the payload itself.
51 // In the case of submessages, the caller needs to re-write the length at the
52 // end in the in the returned memory area.
53 // The problem here is that, because of filtering, the length of a submessage
54 // might be < original length (the original length is still an upper-bound).
55 // Returns a pair with: (1) the pointer where the final length should be written
56 // into, (2) the length of the size field.
57 // The caller must write a redundant varint to match the original size (i.e.
58 // needs to use WriteRedundantVarInt()).
AppendLenDelim(uint32_t field_id,uint32_t len,uint8_t ** out)59 inline std::pair<uint8_t*, uint32_t> AppendLenDelim(uint32_t field_id,
60                                                     uint32_t len,
61                                                     uint8_t** out) {
62   *out = proto_utils::WriteVarInt(proto_utils::MakeTagLengthDelimited(field_id),
63                                   *out);
64   uint8_t* size_field_start = *out;
65   *out = proto_utils::WriteVarInt(len, *out);
66   const size_t size_field_len = static_cast<size_t>(*out - size_field_start);
67   return std::make_pair(size_field_start, size_field_len);
68 }
69 }  // namespace
70 
MessageFilter()71 MessageFilter::MessageFilter() {
72   // Push a state on the stack for the implicit root message.
73   stack_.emplace_back();
74 }
75 
76 MessageFilter::~MessageFilter() = default;
77 
LoadFilterBytecode(const void * filter_data,size_t len)78 bool MessageFilter::LoadFilterBytecode(const void* filter_data, size_t len) {
79   return filter_.Load(filter_data, len);
80 }
81 
SetFilterRoot(const uint32_t * field_ids,size_t num_fields)82 bool MessageFilter::SetFilterRoot(const uint32_t* field_ids,
83                                   size_t num_fields) {
84   uint32_t root_msg_idx = 0;
85   for (const uint32_t* it = field_ids; it < field_ids + num_fields; ++it) {
86     uint32_t field_id = *it;
87     auto res = filter_.Query(root_msg_idx, field_id);
88     if (!res.allowed || res.simple_field())
89       return false;
90     root_msg_idx = res.nested_msg_index;
91   }
92   root_msg_index_ = root_msg_idx;
93   return true;
94 }
95 
FilterMessageFragments(const InputSlice * slices,size_t num_slices)96 MessageFilter::FilteredMessage MessageFilter::FilterMessageFragments(
97     const InputSlice* slices,
98     size_t num_slices) {
99   // First compute the upper bound for the output. The filtered message cannot
100   // be > the original message.
101   uint32_t total_len = 0;
102   for (size_t i = 0; i < num_slices; ++i)
103     total_len += slices[i].len;
104   out_buf_.reset(new uint8_t[total_len]);
105   out_ = out_buf_.get();
106   out_end_ = out_ + total_len;
107 
108   // Reset the parser state.
109   tokenizer_ = MessageTokenizer();
110   error_ = false;
111   stack_.clear();
112   stack_.resize(2);
113   // stack_[0] is a sentinel and should never be hit in nominal cases. If we
114   // end up there we will just keep consuming the input stream and detecting
115   // at the end, without hurting the fastpath.
116   stack_[0].in_bytes_limit = UINT32_MAX;
117   stack_[0].eat_next_bytes = UINT32_MAX;
118   // stack_[1] is the actual root message.
119   stack_[1].in_bytes_limit = total_len;
120   stack_[1].msg_index = root_msg_index_;
121 
122   // Process the input data and write the output.
123   for (size_t slice_idx = 0; slice_idx < num_slices; ++slice_idx) {
124     const InputSlice& slice = slices[slice_idx];
125     const uint8_t* data = static_cast<const uint8_t*>(slice.data);
126     for (size_t i = 0; i < slice.len; ++i)
127       FilterOneByte(data[i]);
128   }
129 
130   // Construct the output object.
131   PERFETTO_CHECK(out_ >= out_buf_.get() && out_ <= out_end_);
132   auto used_size = static_cast<size_t>(out_ - out_buf_.get());
133   FilteredMessage res{std::move(out_buf_), used_size};
134   res.error = error_;
135   if (stack_.size() != 1 || !tokenizer_.idle() ||
136       stack_[0].in_bytes != total_len) {
137     res.error = true;
138   }
139   return res;
140 }
141 
FilterOneByte(uint8_t octet)142 void MessageFilter::FilterOneByte(uint8_t octet) {
143   PERFETTO_DCHECK(!stack_.empty());
144 
145   auto* state = &stack_.back();
146   StackState next_state{};
147   bool push_next_state = false;
148 
149   if (state->eat_next_bytes > 0) {
150     // This is the case where the previous tokenizer_.Push() call returned a
151     // length delimited message which is NOT a submessage (a string or a bytes
152     // field). We just want to consume it, and pass it through in output
153     // if the field was allowed.
154     --state->eat_next_bytes;
155     if (state->passthrough_eaten_bytes)
156       *(out_++) = octet;
157   } else {
158     MessageTokenizer::Token token = tokenizer_.Push(octet);
159     // |token| will not be valid() in most cases and this is WAI. When pushing
160     // a varint field, only the last byte yields a token, all the other bytes
161     // return an invalid token, they just update the internal tokenizer state.
162     if (token.valid()) {
163       auto filter = filter_.Query(state->msg_index, token.field_id);
164       switch (token.type) {
165         case proto_utils::ProtoWireType::kVarInt:
166           if (filter.allowed && filter.simple_field())
167             AppendVarInt(token.field_id, token.value, &out_);
168           break;
169         case proto_utils::ProtoWireType::kFixed32:
170           if (filter.allowed && filter.simple_field())
171             AppendFixed(token.field_id, static_cast<uint32_t>(token.value),
172                         &out_);
173           break;
174         case proto_utils::ProtoWireType::kFixed64:
175           if (filter.allowed && filter.simple_field())
176             AppendFixed(token.field_id, static_cast<uint64_t>(token.value),
177                         &out_);
178           break;
179         case proto_utils::ProtoWireType::kLengthDelimited:
180           // Here we have two cases:
181           // A. A simple string/bytes field: we just want to consume the next
182           //    bytes (the string payload), optionally passing them through in
183           //    output if the field is allowed.
184           // B. This is a nested submessage. In this case we want to recurse and
185           //    push a new state on the stack.
186           // Note that we can't tell the difference between a
187           // "non-allowed string" and a "non-allowed submessage". But it doesn't
188           // matter because in both cases we just want to skip the next N bytes.
189           const auto submessage_len = static_cast<uint32_t>(token.value);
190           auto in_bytes_left = state->in_bytes_limit - state->in_bytes - 1;
191           if (PERFETTO_UNLIKELY(submessage_len > in_bytes_left)) {
192             // This is a malicious / malformed string/bytes/submessage that
193             // claims to be larger than the outer message that contains it.
194             return SetUnrecoverableErrorState();
195           }
196 
197           if (filter.allowed && !filter.simple_field() && submessage_len > 0) {
198             // submessage_len == 0 is the edge case of a message with a 0-len
199             // (but present) submessage. In this case, if allowed, we don't want
200             // to push any further state (doing so would desync the FSM) but we
201             // still want to emit it.
202             // At this point |submessage_len| is only an upper bound. The
203             // final message written in output can be <= the one in input,
204             // only some of its fields might be allowed (also remember that
205             // this class implicitly removes redundancy varint encoding of
206             // len-delimited field lengths). The final length varint (the
207             // return value of AppendLenDelim()) will be filled when popping
208             // from |stack_|.
209             auto size_field =
210                 AppendLenDelim(token.field_id, submessage_len, &out_);
211             push_next_state = true;
212             next_state.field_id = token.field_id;
213             next_state.msg_index = filter.nested_msg_index;
214             next_state.in_bytes_limit = submessage_len;
215             next_state.size_field = size_field.first;
216             next_state.size_field_len = size_field.second;
217             next_state.out_bytes_written_at_start = out_written();
218           } else {
219             // A string or bytes field, or a 0 length submessage.
220             state->eat_next_bytes = submessage_len;
221             state->passthrough_eaten_bytes = filter.allowed;
222             if (filter.allowed)
223               AppendLenDelim(token.field_id, submessage_len, &out_);
224           }
225           break;
226       }  // switch(type)
227 
228       if (PERFETTO_UNLIKELY(track_field_usage_)) {
229         IncrementCurrentFieldUsage(token.field_id, filter.allowed);
230       }
231     }  // if (token.valid)
232   }    // if (eat_next_bytes == 0)
233 
234   ++state->in_bytes;
235   while (state->in_bytes >= state->in_bytes_limit) {
236     PERFETTO_DCHECK(state->in_bytes == state->in_bytes_limit);
237     push_next_state = false;
238 
239     // We can't possibly write more than we read.
240     const uint32_t msg_bytes_written = static_cast<uint32_t>(
241         out_written() - state->out_bytes_written_at_start);
242     PERFETTO_DCHECK(msg_bytes_written <= state->in_bytes_limit);
243 
244     // Backfill the length field of the
245     proto_utils::WriteRedundantVarInt(msg_bytes_written, state->size_field,
246                                       state->size_field_len);
247 
248     const uint32_t in_bytes_processes_for_last_msg = state->in_bytes;
249     stack_.pop_back();
250     PERFETTO_CHECK(!stack_.empty());
251     state = &stack_.back();
252     state->in_bytes += in_bytes_processes_for_last_msg;
253     if (PERFETTO_UNLIKELY(!tokenizer_.idle())) {
254       // If we hit this case, it means that we got to the end of a submessage
255       // while decoding a field. We can't recover from this and we don't want to
256       // propagate a broken sub-message.
257       return SetUnrecoverableErrorState();
258     }
259   }
260 
261   if (push_next_state) {
262     PERFETTO_DCHECK(tokenizer_.idle());
263     stack_.emplace_back(std::move(next_state));
264     state = &stack_.back();
265   }
266 }
267 
SetUnrecoverableErrorState()268 void MessageFilter::SetUnrecoverableErrorState() {
269   error_ = true;
270   stack_.clear();
271   stack_.resize(1);
272   auto& state = stack_[0];
273   state.eat_next_bytes = UINT32_MAX;
274   state.in_bytes_limit = UINT32_MAX;
275   state.passthrough_eaten_bytes = false;
276   out_ = out_buf_.get();  // Reset the write pointer.
277 }
278 
IncrementCurrentFieldUsage(uint32_t field_id,bool allowed)279 void MessageFilter::IncrementCurrentFieldUsage(uint32_t field_id,
280                                                bool allowed) {
281   // Slowpath. Used mainly in offline tools and tests to workout used fields in
282   // a proto.
283   PERFETTO_DCHECK(track_field_usage_);
284 
285   // Field path contains a concatenation of varints, one for each nesting level.
286   // e.g. y in message Root { Sub x = 2; }; message Sub { SubSub y = 7; }
287   // is encoded as [varint(2) + varint(7)].
288   // We use varint to take the most out of SSO (small string opt). In most cases
289   // the path will fit in the on-stack 22 bytes, requiring no heap.
290   std::string field_path;
291 
292   auto append_field_id = [&field_path](uint32_t id) {
293     uint8_t buf[10];
294     uint8_t* end = proto_utils::WriteVarInt(id, buf);
295     field_path.append(reinterpret_cast<char*>(buf),
296                       static_cast<size_t>(end - buf));
297   };
298 
299   // Append all the ancestors IDs from the state stack.
300   // The first entry of the stack has always ID 0 and we skip it (we don't know
301   // the ID of the root message itself).
302   PERFETTO_DCHECK(stack_.size() >= 2 && stack_[1].field_id == 0);
303   for (size_t i = 2; i < stack_.size(); ++i)
304     append_field_id(stack_[i].field_id);
305   // Append the id of the field in the current message.
306   append_field_id(field_id);
307   field_usage_[field_path] += allowed ? 1 : -1;
308 }
309 
310 }  // namespace protozero
311