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/base/scoped_file.h"
22 #include "perfetto/base/temp_file.h"
23 #include "perfetto/base/utils.h"
24
25 #include "gmock/gmock.h"
26 #include "gtest/gtest.h"
27
28 using testing::_;
29 using testing::NiceMock;
30 using testing::StrictMock;
31 using testing::Invoke;
32 using testing::Return;
33
34 namespace perfetto {
35
36 namespace {
37
38 class MockRateLimiter : public RateLimiter {
39 public:
MockRateLimiter()40 MockRateLimiter() : dir_(base::TempDir::Create()) {
41 ON_CALL(*this, LoadState(_))
42 .WillByDefault(Invoke(this, &MockRateLimiter::LoadStateConcrete));
43 ON_CALL(*this, SaveState(_))
44 .WillByDefault(Invoke(this, &MockRateLimiter::SaveStateConcrete));
45 }
46
GetStateFilePath() const47 virtual std::string GetStateFilePath() const {
48 return std::string(dir_.path()) + "/.guardraildata";
49 }
50
~MockRateLimiter()51 virtual ~MockRateLimiter() override {
52 if (StateFileExists())
53 remove(GetStateFilePath().c_str());
54 }
55
LoadStateConcrete(PerfettoCmdState * state)56 bool LoadStateConcrete(PerfettoCmdState* state) {
57 return RateLimiter::LoadState(state);
58 }
59
SaveStateConcrete(const PerfettoCmdState & state)60 bool SaveStateConcrete(const PerfettoCmdState& state) {
61 return RateLimiter::SaveState(state);
62 }
63
64 MOCK_METHOD1(LoadState, bool(PerfettoCmdState*));
65 MOCK_METHOD1(SaveState, bool(const PerfettoCmdState&));
66
67 private:
68 base::TempDir dir_;
69 };
70
WriteGarbageToFile(const std::string & path)71 void WriteGarbageToFile(const std::string& path) {
72 base::ScopedFile fd;
73 fd.reset(open(path.c_str(), O_WRONLY | O_CREAT, 0600));
74 constexpr char data[] = "Some random bytes.";
75 if (write(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 PerfettoCmdState input{};
83 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 }
90
TEST(RateLimiterTest,LoadFromEmpty)91 TEST(RateLimiterTest, LoadFromEmpty) {
92 NiceMock<MockRateLimiter> limiter;
93
94 PerfettoCmdState input{};
95 input.set_total_bytes_uploaded(0);
96 input.set_last_trace_timestamp(0);
97 input.set_first_trace_timestamp(0);
98 PerfettoCmdState output{};
99
100 ASSERT_TRUE(limiter.SaveState(input));
101 ASSERT_TRUE(limiter.LoadState(&output));
102 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
103 }
104
TEST(RateLimiterTest,LoadFromNoFileFails)105 TEST(RateLimiterTest, LoadFromNoFileFails) {
106 NiceMock<MockRateLimiter> limiter;
107 PerfettoCmdState output{};
108 ASSERT_FALSE(limiter.LoadState(&output));
109 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
110 }
111
TEST(RateLimiterTest,LoadFromGarbageFails)112 TEST(RateLimiterTest, LoadFromGarbageFails) {
113 NiceMock<MockRateLimiter> limiter;
114
115 WriteGarbageToFile(limiter.GetStateFilePath().c_str());
116
117 PerfettoCmdState output{};
118 ASSERT_FALSE(limiter.LoadState(&output));
119 ASSERT_EQ(output.total_bytes_uploaded(), 0u);
120 }
121
TEST(RateLimiterTest,NotDropBox)122 TEST(RateLimiterTest, NotDropBox) {
123 StrictMock<MockRateLimiter> limiter;
124
125 ASSERT_TRUE(limiter.ShouldTrace({}));
126 ASSERT_TRUE(limiter.OnTraceDone({}, true, 10000));
127 ASSERT_FALSE(limiter.StateFileExists());
128 }
129
TEST(RateLimiterTest,NotDropBox_FailedToTrace)130 TEST(RateLimiterTest, NotDropBox_FailedToTrace) {
131 StrictMock<MockRateLimiter> limiter;
132
133 ASSERT_FALSE(limiter.OnTraceDone({}, false, 0));
134 ASSERT_FALSE(limiter.StateFileExists());
135 }
136
TEST(RateLimiterTest,DropBox_IgnoreGuardrails)137 TEST(RateLimiterTest, DropBox_IgnoreGuardrails) {
138 StrictMock<MockRateLimiter> limiter;
139 RateLimiter::Args args;
140
141 args.is_dropbox = true;
142 args.ignore_guardrails = true;
143 args.current_time = base::TimeSeconds(41);
144
145 EXPECT_CALL(limiter, SaveState(_));
146 EXPECT_CALL(limiter, LoadState(_));
147 ASSERT_TRUE(limiter.ShouldTrace(args));
148
149 EXPECT_CALL(limiter, SaveState(_));
150 ASSERT_TRUE(limiter.OnTraceDone(args, true, 42u));
151
152 PerfettoCmdState output{};
153 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
154 ASSERT_EQ(output.first_trace_timestamp(), 41u);
155 ASSERT_EQ(output.last_trace_timestamp(), 41u);
156 ASSERT_EQ(output.total_bytes_uploaded(), 42u);
157 }
158
TEST(RateLimiterTest,DropBox_EmptyState)159 TEST(RateLimiterTest, DropBox_EmptyState) {
160 StrictMock<MockRateLimiter> limiter;
161 RateLimiter::Args args;
162
163 args.is_dropbox = true;
164 args.current_time = base::TimeSeconds(10000);
165
166 EXPECT_CALL(limiter, SaveState(_));
167 EXPECT_CALL(limiter, LoadState(_));
168 ASSERT_TRUE(limiter.ShouldTrace(args));
169
170 EXPECT_CALL(limiter, SaveState(_));
171 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
172
173 PerfettoCmdState output{};
174 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
175 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
176 EXPECT_EQ(output.first_trace_timestamp(), 10000u);
177 EXPECT_EQ(output.last_trace_timestamp(), 10000u);
178 }
179
TEST(RateLimiterTest,DropBox_NormalUpload)180 TEST(RateLimiterTest, DropBox_NormalUpload) {
181 StrictMock<MockRateLimiter> limiter;
182 RateLimiter::Args args;
183
184 PerfettoCmdState input{};
185 input.set_first_trace_timestamp(10000);
186 input.set_last_trace_timestamp(10000 + 60 * 10);
187 input.set_total_bytes_uploaded(1024 * 1024 * 2);
188 ASSERT_TRUE(limiter.SaveStateConcrete(input));
189
190 args.is_dropbox = true;
191 args.current_time = base::TimeSeconds(input.last_trace_timestamp() + 60 * 10);
192
193 EXPECT_CALL(limiter, LoadState(_));
194 ASSERT_TRUE(limiter.ShouldTrace(args));
195
196 EXPECT_CALL(limiter, SaveState(_));
197 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
198
199 PerfettoCmdState output{};
200 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
201 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u * 3);
202 EXPECT_EQ(output.first_trace_timestamp(), input.first_trace_timestamp());
203 EXPECT_EQ(output.last_trace_timestamp(),
204 static_cast<uint64_t>(args.current_time.count()));
205 }
206
TEST(RateLimiterTest,DropBox_FailedToLoadState)207 TEST(RateLimiterTest, DropBox_FailedToLoadState) {
208 StrictMock<MockRateLimiter> limiter;
209 RateLimiter::Args args;
210
211 args.is_dropbox = true;
212
213 WriteGarbageToFile(limiter.GetStateFilePath().c_str());
214
215 EXPECT_CALL(limiter, LoadState(_));
216 EXPECT_CALL(limiter, SaveState(_));
217 ASSERT_FALSE(limiter.ShouldTrace(args));
218
219 PerfettoCmdState output{};
220 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
221 EXPECT_EQ(output.total_bytes_uploaded(), 0u);
222 EXPECT_EQ(output.first_trace_timestamp(), 0u);
223 EXPECT_EQ(output.last_trace_timestamp(), 0u);
224 }
225
TEST(RateLimiterTest,DropBox_NoTimeTravel)226 TEST(RateLimiterTest, DropBox_NoTimeTravel) {
227 StrictMock<MockRateLimiter> limiter;
228 RateLimiter::Args args;
229
230 PerfettoCmdState input{};
231 input.set_first_trace_timestamp(100);
232 input.set_last_trace_timestamp(100);
233 ASSERT_TRUE(limiter.SaveStateConcrete(input));
234
235 args.is_dropbox = true;
236 args.current_time = base::TimeSeconds(99);
237
238 EXPECT_CALL(limiter, LoadState(_));
239 EXPECT_CALL(limiter, SaveState(_));
240 ASSERT_FALSE(limiter.ShouldTrace(args));
241
242 PerfettoCmdState output{};
243 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
244 EXPECT_EQ(output.total_bytes_uploaded(), 0u);
245 EXPECT_EQ(output.first_trace_timestamp(), 0u);
246 EXPECT_EQ(output.last_trace_timestamp(), 0u);
247 }
248
TEST(RateLimiterTest,DropBox_TooSoon)249 TEST(RateLimiterTest, DropBox_TooSoon) {
250 StrictMock<MockRateLimiter> limiter;
251 RateLimiter::Args args;
252
253 PerfettoCmdState input{};
254 input.set_first_trace_timestamp(10000);
255 input.set_last_trace_timestamp(10000);
256 ASSERT_TRUE(limiter.SaveStateConcrete(input));
257
258 args.is_dropbox = true;
259 args.current_time = base::TimeSeconds(10000 + 60 * 4);
260
261 EXPECT_CALL(limiter, LoadState(_));
262 ASSERT_FALSE(limiter.ShouldTrace(args));
263 }
264
TEST(RateLimiterTest,DropBox_TooMuch)265 TEST(RateLimiterTest, DropBox_TooMuch) {
266 StrictMock<MockRateLimiter> limiter;
267 RateLimiter::Args args;
268
269 PerfettoCmdState input{};
270 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
271 ASSERT_TRUE(limiter.SaveStateConcrete(input));
272
273 args.is_dropbox = true;
274 args.current_time = base::TimeSeconds(60 * 60);
275
276 EXPECT_CALL(limiter, LoadState(_));
277 ASSERT_FALSE(limiter.ShouldTrace(args));
278 }
279
TEST(RateLimiterTest,DropBox_TooMuch_Override)280 TEST(RateLimiterTest, DropBox_TooMuch_Override) {
281 StrictMock<MockRateLimiter> limiter;
282 RateLimiter::Args args;
283
284 PerfettoCmdState input{};
285 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
286 ASSERT_TRUE(limiter.SaveStateConcrete(input));
287
288 args.is_dropbox = true;
289 args.current_time = base::TimeSeconds(60 * 60);
290 args.max_upload_bytes_override = 10 * 1024 * 1024 + 2;
291
292 EXPECT_CALL(limiter, LoadState(_));
293 ASSERT_TRUE(limiter.ShouldTrace(args));
294 }
295
TEST(RateLimiterTest,DropBox_TooMuchWasUploaded)296 TEST(RateLimiterTest, DropBox_TooMuchWasUploaded) {
297 StrictMock<MockRateLimiter> limiter;
298 RateLimiter::Args args;
299
300 PerfettoCmdState input{};
301 input.set_first_trace_timestamp(1);
302 input.set_last_trace_timestamp(1);
303 input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
304 ASSERT_TRUE(limiter.SaveStateConcrete(input));
305
306 args.is_dropbox = true;
307 args.current_time = base::TimeSeconds(60 * 60 * 24 + 2);
308
309 EXPECT_CALL(limiter, LoadState(_));
310 ASSERT_TRUE(limiter.ShouldTrace(args));
311
312 EXPECT_CALL(limiter, SaveState(_));
313 ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
314
315 PerfettoCmdState output{};
316 ASSERT_TRUE(limiter.LoadStateConcrete(&output));
317 EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
318 EXPECT_EQ(output.first_trace_timestamp(),
319 static_cast<uint64_t>(args.current_time.count()));
320 EXPECT_EQ(output.last_trace_timestamp(),
321 static_cast<uint64_t>(args.current_time.count()));
322 }
323
TEST(RateLimiterTest,DropBox_FailedToUpload)324 TEST(RateLimiterTest, DropBox_FailedToUpload) {
325 StrictMock<MockRateLimiter> limiter;
326 RateLimiter::Args args;
327
328 args.is_dropbox = true;
329 args.current_time = base::TimeSeconds(10000);
330
331 EXPECT_CALL(limiter, SaveState(_));
332 EXPECT_CALL(limiter, LoadState(_));
333 ASSERT_TRUE(limiter.ShouldTrace(args));
334 ASSERT_FALSE(limiter.OnTraceDone(args, false, 1024 * 1024));
335 }
336
TEST(RateLimiterTest,DropBox_FailedToSave)337 TEST(RateLimiterTest, DropBox_FailedToSave) {
338 StrictMock<MockRateLimiter> limiter;
339 RateLimiter::Args args;
340
341 args.is_dropbox = true;
342 args.current_time = base::TimeSeconds(10000);
343
344 EXPECT_CALL(limiter, SaveState(_));
345 EXPECT_CALL(limiter, LoadState(_));
346 ASSERT_TRUE(limiter.ShouldTrace(args));
347
348 EXPECT_CALL(limiter, SaveState(_)).WillOnce(Return(false));
349 ASSERT_FALSE(limiter.OnTraceDone(args, true, 1024 * 1024));
350 }
351
352 } // namespace
353
354 } // namespace perfetto
355