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: 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 47 virtual std::string GetStateFilePath() const { 48 return std::string(dir_.path()) + "/.guardraildata"; 49 } 50 51 virtual ~MockRateLimiter() override { 52 if (StateFileExists()) 53 remove(GetStateFilePath().c_str()); 54 } 55 56 bool LoadStateConcrete(PerfettoCmdState* state) { 57 return RateLimiter::LoadState(state); 58 } 59 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 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 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 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 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 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 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 130 TEST(RateLimiterTest, NotDropBox_FailedToTrace) { 131 StrictMock<MockRateLimiter> limiter; 132 133 ASSERT_FALSE(limiter.OnTraceDone({}, false, 0)); 134 ASSERT_FALSE(limiter.StateFileExists()); 135 } 136 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 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 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 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 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 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 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 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 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 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 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
