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