1 /*
2 * Copyright (C) 2018 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/perfetto_cmd/rate_limiter.h"
18
19 #include <stdio.h>
20
21 #include "perfetto/ext/base/file_utils.h"
22 #include "perfetto/ext/base/scoped_file.h"
23 #include "perfetto/ext/base/temp_file.h"
24 #include "perfetto/ext/base/utils.h"
25
26 #include "test/gtest_and_gmock.h"
27
28 using testing::_;
29 using testing::Contains;
30 using testing::Invoke;
31 using testing::NiceMock;
32 using testing::Return;
33 using testing::StrictMock;
34
35 namespace perfetto {
36
37 namespace {
38
39 class MockRateLimiter : public RateLimiter {
40 public:
MockRateLimiter()41 MockRateLimiter() : dir_(base::TempDir::Create()) {
42 ON_CALL(*this, LoadState(_))
43 .WillByDefault(Invoke(this, &MockRateLimiter::LoadStateConcrete));
44 ON_CALL(*this, SaveState(_))
45 .WillByDefault(Invoke(this, &MockRateLimiter::SaveStateConcrete));
46 }
47
GetStateFilePath() const48 virtual std::string GetStateFilePath() const {
49 return std::string(dir_.path()) + "/.guardraildata";
50 }
51
~MockRateLimiter()52 virtual ~MockRateLimiter() override {
53 if (StateFileExists())
54 remove(GetStateFilePath().c_str());
55 }
56
LoadStateConcrete(gen::PerfettoCmdState * state)57 bool LoadStateConcrete(gen::PerfettoCmdState* state) {
58 return RateLimiter::LoadState(state);
59 }
60
SaveStateConcrete(const gen::PerfettoCmdState & state)61 bool SaveStateConcrete(const gen::PerfettoCmdState& state) {
62 return RateLimiter::SaveState(state);
63 }
64
65 MOCK_METHOD1(LoadState, bool(gen::PerfettoCmdState*));
66 MOCK_METHOD1(SaveState, bool(const gen::PerfettoCmdState&));
67
68 private:
69 base::TempDir dir_;
70 };
71
WriteGarbageToFile(const std::string & path)72 void WriteGarbageToFile(const std::string& path) {
73 base::ScopedFile fd(base::OpenFile(path, O_WRONLY | O_CREAT, 0600));
74 constexpr char data[] = "Some random bytes.";
75 if (base::WriteAll(fd.get(), data, sizeof(data)) != sizeof(data))
76 ADD_FAILURE() << "Could not write garbage";
77 }
78
TEST(RateLimiterTest,RoundTripState)79 TEST(RateLimiterTest, RoundTripState) {
80 NiceMock<MockRateLimiter> limiter;
81
82 gen::PerfettoCmdState input{};
83 gen::PerfettoCmdState output{};
84
85 input.set_total_bytes_uploaded(42);
86 ASSERT_TRUE(limiter.SaveState(input));
87 ASSERT_TRUE(limiter.LoadState(&output));
88 ASSERT_EQ(output.total_bytes_uploaded(), 42u);
89 ASSERT_EQ(output.session_state_size(), 0);
90 }
91
TEST(RateLimiterTest,FileIsSensiblyTruncated)92 TEST(RateLimiterTest, FileIsSensiblyTruncated) {
93 NiceMock<MockRateLimiter> limiter;
94
95 gen::PerfettoCmdState input{};
96 gen::PerfettoCmdState output{};
97
98 input.set_total_bytes_uploaded(42);
99 input.set_first_trace_timestamp(1);
100 input.set_last_trace_timestamp(2);
101
102 for (size_t i = 0; i < 100; ++i) {
103 auto* session = input.add_session_state();
104 session->set_session_name("session_" + std::to_string(i));
105 session->set_total_bytes_uploaded(i * 100);
106 session->set_last_trace_timestamp(i);
107 }
108
109 ASSERT_TRUE(limiter.SaveState(input));
110 ASSERT_TRUE(limiter.LoadState(&output));
111
112 ASSERT_EQ(output.total_bytes_uploaded(), 42u);
113 ASSERT_EQ(output.first_trace_timestamp(), 1u);
114 ASSERT_EQ(output.last_trace_timestamp(), 2u);
115 ASSERT_LE(output.session_state_size(), 50);
116 ASSERT_GE(output.session_state_size(), 5);
117
118 {
119 gen::PerfettoCmdState::PerSessionState session;
120 session.set_session_name("session_99");
121 session.set_total_bytes_uploaded(99 * 100);
122 session.set_last_trace_timestamp(99);
123 ASSERT_THAT(output.session_state(), Contains(session));
124 }
125 }
126
TEST(RateLimiterTest,LoadFromEmpty)127 TEST(RateLimiterTest, LoadFromEmpty) {
128 NiceMock<MockRateLimiter> limiter;
129
130 gen::PerfettoCmdState input{};
131 input.set_total_bytes_uploaded(0);
132 input.set_last_trace_timestamp(0);
133 input.set_first_trace_timestamp(0);
134 gen::PerfettoCmdState output{};
135
136 ASSERT_TRUE(limiter.SaveState(input));
137 ASSERT_TRUE(limiter.LoadState(&output));
138 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
139 }
140
TEST(RateLimiterTest,LoadFromNoFileFails)141 TEST(RateLimiterTest, LoadFromNoFileFails) {
142 NiceMock<MockRateLimiter> limiter;
143 gen::PerfettoCmdState output{};
144 ASSERT_FALSE(limiter.LoadState(&output));
145 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
146 }
147
TEST(RateLimiterTest,LoadFromGarbageFails)148 TEST(RateLimiterTest, LoadFromGarbageFails) {
149 NiceMock<MockRateLimiter> limiter;
150
151 WriteGarbageToFile(limiter.GetStateFilePath().c_str());
152
153 gen::PerfettoCmdState output{};
154 ASSERT_FALSE(limiter.LoadState(&output));
155 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
156 }
157
TEST(RateLimiterTest,NotDropBox)158 TEST(RateLimiterTest, NotDropBox) {
159 StrictMock<MockRateLimiter> limiter;
160
161 ASSERT_EQ(limiter.ShouldTrace({}), RateLimiter::kOkToTrace);
162 ASSERT_TRUE(limiter.OnTraceDone({}, true, 10000));
163 ASSERT_FALSE(limiter.StateFileExists());
164 }
165
TEST(RateLimiterTest,NotDropBox_FailedToTrace)166 TEST(RateLimiterTest, NotDropBox_FailedToTrace) {
167 StrictMock<MockRateLimiter> limiter;
168
169 ASSERT_FALSE(limiter.OnTraceDone({}, false, 0));
170 ASSERT_FALSE(limiter.StateFileExists());
171 }
172
TEST(RateLimiterTest,DropBox_IgnoreGuardrails)173 TEST(RateLimiterTest, DropBox_IgnoreGuardrails) {
174 StrictMock<MockRateLimiter> limiter;
175 RateLimiter::Args args;
176
177 args.allow_user_build_tracing = true;
178 args.is_uploading = true;
179 args.ignore_guardrails = true;
180 args.current_time = base::TimeSeconds(41);
181
182 EXPECT_CALL(limiter, SaveState(_));
183 EXPECT_CALL(limiter, LoadState(_));
184 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
185
186 EXPECT_CALL(limiter, SaveState(_));
187 ASSERT_TRUE(limiter.OnTraceDone(args, true, 42u));
188
189 gen::PerfettoCmdState output{};
190 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
191 ASSERT_EQ(output.first_trace_timestamp(), 41u);
192 ASSERT_EQ(output.last_trace_timestamp(), 41u);
193 ASSERT_EQ(output.total_bytes_uploaded(), 42u);
194 }
195
TEST(RateLimiterTest,DropBox_EmptyState)196 TEST(RateLimiterTest, DropBox_EmptyState) {
197 StrictMock<MockRateLimiter> limiter;
198 RateLimiter::Args args;
199
200 args.allow_user_build_tracing = true;
201 args.is_uploading = true;
202 args.current_time = base::TimeSeconds(10000);
203
204 EXPECT_CALL(limiter, SaveState(_));
205 EXPECT_CALL(limiter, LoadState(_));
206 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
207
208 EXPECT_CALL(limiter, SaveState(_));
209 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
210
211 gen::PerfettoCmdState output{};
212 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
213 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
214 EXPECT_EQ(output.first_trace_timestamp(), 10000u);
215 EXPECT_EQ(output.last_trace_timestamp(), 10000u);
216 }
217
TEST(RateLimiterTest,DropBox_NormalUpload)218 TEST(RateLimiterTest, DropBox_NormalUpload) {
219 StrictMock<MockRateLimiter> limiter;
220 RateLimiter::Args args;
221
222 gen::PerfettoCmdState input{};
223 input.set_first_trace_timestamp(10000);
224 input.set_last_trace_timestamp(10000 + 60 * 10);
225 input.set_total_bytes_uploaded(1024 * 1024 * 2);
226 ASSERT_TRUE(limiter.SaveStateConcrete(input));
227
228 args.allow_user_build_tracing = true;
229 args.is_uploading = true;
230 args.current_time = base::TimeSeconds(input.last_trace_timestamp() + 60 * 10);
231
232 EXPECT_CALL(limiter, LoadState(_));
233 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
234
235 EXPECT_CALL(limiter, SaveState(_));
236 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
237
238 gen::PerfettoCmdState output{};
239 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
240 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u * 3);
241 EXPECT_EQ(output.first_trace_timestamp(), input.first_trace_timestamp());
242 EXPECT_EQ(output.last_trace_timestamp(),
243 static_cast<uint64_t>(args.current_time.count()));
244 }
245
TEST(RateLimiterTest,DropBox_NormalUploadWithSessionName)246 TEST(RateLimiterTest, DropBox_NormalUploadWithSessionName) {
247 StrictMock<MockRateLimiter> limiter;
248 RateLimiter::Args args;
249
250 gen::PerfettoCmdState input{};
251 input.set_first_trace_timestamp(10000);
252 input.set_last_trace_timestamp(10000 + 60 * 10);
253 input.set_total_bytes_uploaded(1024 * 1024 * 2);
254 ASSERT_TRUE(limiter.SaveStateConcrete(input));
255
256 args.allow_user_build_tracing = true;
257 args.is_uploading = true;
258 args.unique_session_name = "foo";
259 args.current_time = base::TimeSeconds(input.last_trace_timestamp() + 60 * 10);
260
261 EXPECT_CALL(limiter, LoadState(_));
262 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
263
264 EXPECT_CALL(limiter, SaveState(_));
265 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
266
267 gen::PerfettoCmdState output{};
268 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
269 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u * 2);
270 EXPECT_EQ(output.first_trace_timestamp(), input.first_trace_timestamp());
271 EXPECT_EQ(output.last_trace_timestamp(),
272 static_cast<uint64_t>(args.current_time.count()));
273 ASSERT_GE(output.session_state_size(), 1);
274
275 {
276 gen::PerfettoCmdState::PerSessionState session;
277 session.set_session_name("foo");
278 session.set_total_bytes_uploaded(1024 * 1024);
279 session.set_last_trace_timestamp(
280 static_cast<uint64_t>(args.current_time.count()));
281 ASSERT_THAT(output.session_state(), Contains(session));
282 }
283 }
284
TEST(RateLimiterTest,DropBox_FailedToLoadState)285 TEST(RateLimiterTest, DropBox_FailedToLoadState) {
286 StrictMock<MockRateLimiter> limiter;
287 RateLimiter::Args args;
288
289 args.allow_user_build_tracing = true;
290 args.is_uploading = true;
291
292 WriteGarbageToFile(limiter.GetStateFilePath().c_str());
293
294 EXPECT_CALL(limiter, LoadState(_));
295 EXPECT_CALL(limiter, SaveState(_));
296 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kInvalidState);
297
298 gen::PerfettoCmdState output{};
299 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
300 EXPECT_EQ(output.total_bytes_uploaded(), 0u);
301 EXPECT_EQ(output.first_trace_timestamp(), 0u);
302 EXPECT_EQ(output.last_trace_timestamp(), 0u);
303 }
304
TEST(RateLimiterTest,DropBox_NoTimeTravel)305 TEST(RateLimiterTest, DropBox_NoTimeTravel) {
306 StrictMock<MockRateLimiter> limiter;
307 RateLimiter::Args args;
308
309 gen::PerfettoCmdState input{};
310 input.set_first_trace_timestamp(100);
311 input.set_last_trace_timestamp(100);
312 ASSERT_TRUE(limiter.SaveStateConcrete(input));
313
314 args.allow_user_build_tracing = true;
315 args.is_uploading = true;
316 args.current_time = base::TimeSeconds(99);
317
318 EXPECT_CALL(limiter, LoadState(_));
319 EXPECT_CALL(limiter, SaveState(_));
320 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kInvalidState);
321
322 gen::PerfettoCmdState output{};
323 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
324 EXPECT_EQ(output.total_bytes_uploaded(), 0u);
325 EXPECT_EQ(output.first_trace_timestamp(), 0u);
326 EXPECT_EQ(output.last_trace_timestamp(), 0u);
327 }
328
TEST(RateLimiterTest,DropBox_TooMuch_OtherSession)329 TEST(RateLimiterTest, DropBox_TooMuch_OtherSession) {
330 StrictMock<MockRateLimiter> limiter;
331 RateLimiter::Args args;
332
333 gen::PerfettoCmdState input{};
334 auto* session = input.add_session_state();
335 session->set_session_name("foo");
336 session->set_total_bytes_uploaded(100 * 1024 * 1024);
337
338 ASSERT_TRUE(limiter.SaveStateConcrete(input));
339
340 args.is_user_build = true;
341 args.allow_user_build_tracing = true;
342 args.is_uploading = true;
343 args.unique_session_name = "bar";
344 args.current_time = base::TimeSeconds(60 * 60);
345
346 EXPECT_CALL(limiter, LoadState(_));
347 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
348 }
349
TEST(RateLimiterTest,DropBox_TooMuch_Session)350 TEST(RateLimiterTest, DropBox_TooMuch_Session) {
351 StrictMock<MockRateLimiter> limiter;
352 RateLimiter::Args args;
353
354 gen::PerfettoCmdState input{};
355 auto* session = input.add_session_state();
356 session->set_session_name("foo");
357 session->set_total_bytes_uploaded(100 * 1024 * 1024);
358
359 ASSERT_TRUE(limiter.SaveStateConcrete(input));
360
361 args.is_user_build = true;
362 args.allow_user_build_tracing = true;
363 args.is_uploading = true;
364 args.unique_session_name = "foo";
365 args.current_time = base::TimeSeconds(60 * 60);
366
367 EXPECT_CALL(limiter, LoadState(_));
368 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kHitUploadLimit);
369 }
370
TEST(RateLimiterTest,DropBox_TooMuch_User)371 TEST(RateLimiterTest, DropBox_TooMuch_User) {
372 StrictMock<MockRateLimiter> limiter;
373 RateLimiter::Args args;
374
375 gen::PerfettoCmdState input{};
376 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
377 ASSERT_TRUE(limiter.SaveStateConcrete(input));
378
379 args.is_user_build = true;
380 args.allow_user_build_tracing = true;
381 args.is_uploading = true;
382 args.current_time = base::TimeSeconds(60 * 60);
383
384 EXPECT_CALL(limiter, LoadState(_));
385 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kHitUploadLimit);
386 }
387
TEST(RateLimiterTest,DropBox_TooMuch_Override)388 TEST(RateLimiterTest, DropBox_TooMuch_Override) {
389 StrictMock<MockRateLimiter> limiter;
390 RateLimiter::Args args;
391
392 gen::PerfettoCmdState input{};
393 auto* session = input.add_session_state();
394 session->set_session_name("foo");
395 session->set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
396 ASSERT_TRUE(limiter.SaveStateConcrete(input));
397
398 args.allow_user_build_tracing = true;
399 args.is_uploading = true;
400 args.current_time = base::TimeSeconds(60 * 60);
401 args.max_upload_bytes_override = 10 * 1024 * 1024 + 2;
402 args.unique_session_name = "foo";
403
404 EXPECT_CALL(limiter, LoadState(_));
405 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
406 }
407
408 // Override doesn't apply to traces without session name.
TEST(RateLimiterTest,DropBox_OverrideOnEmptySesssionName)409 TEST(RateLimiterTest, DropBox_OverrideOnEmptySesssionName) {
410 StrictMock<MockRateLimiter> limiter;
411 RateLimiter::Args args;
412
413 gen::PerfettoCmdState input{};
414 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
415 ASSERT_TRUE(limiter.SaveStateConcrete(input));
416
417 args.allow_user_build_tracing = true;
418 args.is_uploading = true;
419 args.current_time = base::TimeSeconds(60 * 60);
420 args.max_upload_bytes_override = 10 * 1024 * 1024 + 2;
421
422 EXPECT_CALL(limiter, LoadState(_));
423 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kHitUploadLimit);
424 }
425
TEST(RateLimiterTest,DropBox_TooMuchWasUploaded)426 TEST(RateLimiterTest, DropBox_TooMuchWasUploaded) {
427 StrictMock<MockRateLimiter> limiter;
428 RateLimiter::Args args;
429
430 gen::PerfettoCmdState input{};
431 input.set_first_trace_timestamp(1);
432 input.set_last_trace_timestamp(1);
433 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
434 ASSERT_TRUE(limiter.SaveStateConcrete(input));
435
436 args.is_uploading = true;
437 args.current_time = base::TimeSeconds(60 * 60 * 24 + 2);
438
439 EXPECT_CALL(limiter, LoadState(_));
440 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
441
442 EXPECT_CALL(limiter, SaveState(_));
443 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
444
445 gen::PerfettoCmdState output{};
446 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
447 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
448 EXPECT_EQ(output.first_trace_timestamp(),
449 static_cast<uint64_t>(args.current_time.count()));
450 EXPECT_EQ(output.last_trace_timestamp(),
451 static_cast<uint64_t>(args.current_time.count()));
452 }
453
TEST(RateLimiterTest,DropBox_FailedToUpload)454 TEST(RateLimiterTest, DropBox_FailedToUpload) {
455 StrictMock<MockRateLimiter> limiter;
456 RateLimiter::Args args;
457
458 args.is_uploading = true;
459 args.current_time = base::TimeSeconds(10000);
460
461 EXPECT_CALL(limiter, SaveState(_));
462 EXPECT_CALL(limiter, LoadState(_));
463 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
464 ASSERT_FALSE(limiter.OnTraceDone(args, false, 1024 * 1024));
465 }
466
TEST(RateLimiterTest,DropBox_FailedToSave)467 TEST(RateLimiterTest, DropBox_FailedToSave) {
468 StrictMock<MockRateLimiter> limiter;
469 RateLimiter::Args args;
470
471 args.is_uploading = true;
472 args.current_time = base::TimeSeconds(10000);
473
474 EXPECT_CALL(limiter, SaveState(_));
475 EXPECT_CALL(limiter, LoadState(_));
476 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
477
478 EXPECT_CALL(limiter, SaveState(_)).WillOnce(Return(false));
479 ASSERT_FALSE(limiter.OnTraceDone(args, true, 1024 * 1024));
480 }
481
TEST(RateLimiterTest,DropBox_CantTraceOnUser)482 TEST(RateLimiterTest, DropBox_CantTraceOnUser) {
483 StrictMock<MockRateLimiter> limiter;
484 RateLimiter::Args args;
485
486 args.is_user_build = true;
487 args.allow_user_build_tracing = false;
488 args.is_uploading = true;
489 args.current_time = base::TimeSeconds(10000);
490
491 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kNotAllowedOnUserBuild);
492 }
493
TEST(RateLimiterTest,DropBox_CanTraceOnUser)494 TEST(RateLimiterTest, DropBox_CanTraceOnUser) {
495 StrictMock<MockRateLimiter> limiter;
496 RateLimiter::Args args;
497
498 args.is_user_build = false;
499 args.allow_user_build_tracing = false;
500 args.is_uploading = true;
501 args.current_time = base::TimeSeconds(10000);
502
503 EXPECT_CALL(limiter, SaveState(_));
504 EXPECT_CALL(limiter, LoadState(_));
505 ASSERT_EQ(limiter.ShouldTrace(args), RateLimiter::kOkToTrace);
506 }
507
508 } // namespace
509
510 } // namespace perfetto
511