keychain/ckks/tests/RateLimiterTests.m

   1 /*
   2  * Copyright (c) 2017 Apple Inc. All Rights Reserved.
   3  *
   4  * @APPLE_LICENSE_HEADER_START@
   5  *
   6  * This file contains Original Code and/or Modifications of Original Code
   7  * as defined in and that are subject to the Apple Public Source License
   8  * Version 2.0 (the 'License'). You may not use this file except in
   9  * compliance with the License. Please obtain a copy of the License at
  10  * http://www.opensource.apple.com/apsl/ and read it before using this
  11  * file.
  12  *
  13  * The Original Code and all software distributed under the License are
  14  * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
  15  * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
  16  * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
  17  * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
  18  * Please see the License for the specific language governing rights and
  19  * limitations under the License.
  20  *
  21  * @APPLE_LICENSE_HEADER_END@
  22  */
  23
  24 #import <Foundation/Foundation.h>
  25 #import <XCTest/XCTest.h>
  26 #import "keychain/ckks/RateLimiter.h"
  27
  28 @interface TestObject : NSObject
  29 @property NSString *uuid;
  30 - (NSString *)invalid;
  31 @end
  32
  33 @implementation TestObject
  34 - (instancetype)init {
  35     self = [super init];
  36     if (self) {
  37         _uuid = [[NSUUID UUID] UUIDString];
  38     }
  39     return self;
  40 }
  41
  42 - (instancetype)initWithNilUuid {
  43     self = [super init];
  44     if (self) {
  45         _uuid = nil;
  46     }
  47     return self;
  48 }
  49
  50 // It's super illegal for this to get called because of property allowlisting
  51 - (NSString *)invalid {
  52     NSAssert(NO, @"'invalid' is not an approved property");
  53     return nil;
  54 }
  55 @end
  56
  57 @interface RateLimiterTests : XCTestCase
  58 @property NSDictionary *config;
  59 @property NSString *filepath;
  60 @property NSDate *time;
  61 @property RateLimiter *RL;
  62 @property TestObject *obj;
  63 @end
  64
  65 @implementation RateLimiterTests
  66
  67 - (void)setUp {
  68     [super setUp];
  69     // instantiate config, write to disk
  70     NSData *configData = [@"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\
  71 <!DOCTYPE plist PUBLIC \"-//Apple//DTD PLIST 1.0//EN\" \"http://www.apple.com/DTDs/PropertyList-1.0.dtd\">\
  72 <plist version=\"1.0\">\
  73 <dict>\
  74     <key>general</key>\
  75     <dict>\
  76         <key>maxStateSize</key>\
  77         <integer>250</integer>\
  78         <key>maxItemAge</key>\
  79         <integer>3600</integer>\
  80         <key>overloadDuration</key>\
  81         <integer>1800</integer>\
  82         <key>name</key>\
  83         <string>CKKS</string>\
  84         <key>MAType</key>\
  85         <string></string>\
  86     </dict>\
  87     <key>groups</key>\
  88     <array>\
  89         <dict>\
  90             <key>property</key>\
  91             <string>global</string>\
  92             <key>capacity</key>\
  93             <integer>20</integer>\
  94             <key>rate</key>\
  95             <integer>30</integer>\
  96             <key>badness</key>\
  97             <integer>1</integer>\
  98         </dict>\
  99         <dict>\
 100             <key>property</key>\
 101             <string>uuid</string>\
 102             <key>capacity</key>\
 103             <integer>3</integer>\
 104             <key>rate</key>\
 105             <integer>600</integer>\
 106             <key>badness</key>\
 107             <integer>3</integer>\
 108         </dict>\
 109         <dict>\
 110             <key>property</key>\
 111             <string>invalid</string>\
 112             <key>capacity</key>\
 113             <integer>0</integer>\
 114             <key>rate</key>\
 115             <integer>60000</integer>\
 116             <key>badness</key>\
 117             <integer>4</integer>\
 118             </dict>\
 119     </array>\
 120 </dict>\
 121 </plist>\
 122 " dataUsingEncoding:NSUTF8StringEncoding];
 123     NSError *err = nil;
 124     _config = [NSPropertyListSerialization propertyListWithData:configData options:NSPropertyListImmutable format:nil error:&err];
 125     if (!_config) {
 126         XCTFail(@"Could not deserialize property list: %@", err);
 127     }
 128     _filepath = [NSString stringWithFormat:@"/tmp/ratelimitertests_%@.plist", [[NSUUID UUID] UUIDString]];
 129     if (![configData writeToFile:_filepath atomically:NO]) {
 130         XCTFail(@"Could not write plist to %@", _filepath);
 131     }
 132     _RL = [[RateLimiter alloc] initWithConfig:_config];
 133     _obj = [TestObject new];
 134     _time = [NSDate date];
 135 }
 136
 137 - (void)tearDown {
 138     NSError *err = nil;
 139     if (![[NSFileManager defaultManager] removeItemAtPath:_filepath error:&err]) {
 140         XCTFail(@"Couldn't delete file %@: %@", _filepath, err);
 141     }
 142     [super tearDown];
 143 }
 144
 145 - (void)testInitWithConfig {
 146     self.RL = [[RateLimiter alloc] initWithConfig:self.config];
 147     XCTAssertNotNil(self.RL, @"RateLimiter with config succeeds");
 148     XCTAssertNil(self.RL.assetType, @"initWithConfig means no assetType");
 149     XCTAssertEqualObjects(self.config, self.RL.config, @"config was copied properly");
 150 }
 151
 152 - (void)testInitWithPlist {
 153     RateLimiter *RL = [[RateLimiter alloc] initWithPlistFromURL:[NSURL URLWithString:[NSString stringWithFormat:@"file://%@", self.filepath]]];
 154     XCTAssertNotNil(RL, @"RateLimiter with plist succeeds");
 155     XCTAssertNil(RL.assetType, @"initWithPlist means no assetType");
 156     XCTAssertEqualObjects(self.config, RL.config, @"config was loaded properly");
 157     RL = [[RateLimiter alloc] initWithPlistFromURL:[NSURL URLWithString:[NSString stringWithFormat:@"file://%@.nonexisting", self.filepath]]];
 158     XCTAssertNil(RL, "Cannot instantiate RateLimiter with invalid plist URL");
 159 }
 160
 161 - (void)testEncodingAndDecoding {
 162     NSDate* date = [NSDate date];
 163     NSDate* limit = nil;
 164     [self.RL judge:self.obj at:date limitTime:&limit];
 165
 166     NSMutableData *data = [NSMutableData new];
 167     NSKeyedArchiver *encoder = [[NSKeyedArchiver alloc] initForWritingWithMutableData:data];
 168     encoder.requiresSecureCoding = YES;
 169     [self.RL encodeWithCoder:encoder];
 170     [encoder finishEncoding];
 171     XCTAssertEqualObjects(self.config, self.RL.config, @"config unmodified after encoding");
 172     XCTAssertNil(self.RL.assetType, @"assetType still nil after encoding");
 173
 174     NSKeyedUnarchiver *decoder = [[NSKeyedUnarchiver alloc] initForReadingWithData:data];
 175     decoder.requiresSecureCoding = YES;
 176     RateLimiter *RL2 = [[RateLimiter alloc] initWithCoder:decoder];
 177     XCTAssertNotNil(RL2, @"Received an object from initWithCoder");
 178     XCTAssertEqualObjects(self.RL.config, RL2.config, @"config is the same after encoding and decoding");
 179     XCTAssertTrue([self.RL isEqual:RL2], @"RateLimiters believe they are the same");
 180     XCTAssertNil(RL2.assetType, @"assetType remains nil");
 181 }
 182
 183 - (void)testInitWithAssetType {
 184     // Not implemented yet, expect nil
 185     XCTAssertNil([[RateLimiter alloc] initWithAssetType:@"test"]);
 186 }
 187
 188 - (void)testReset {
 189     NSDate *limitTime = nil;
 190     [self.RL judge:[TestObject new] at:self.time limitTime:&limitTime];
 191     XCTAssertEqual([self.RL stateSize], 2ul, @"Single property judged once, state is 1 global plus 1 property");
 192     [self.RL reset];
 193     XCTAssertEqual([self.RL stateSize], 0ul, @"No buckets after reset");
 194     XCTAssertEqualObjects(self.config, self.RL.config);
 195 }
 196
 197 // Cause it to complain based on one item being hit repeatedly
 198 - (void)testJudgeSingleItem {
 199     NSDate *limitTime = nil;
 200     for (int idx = 0; idx < [self.config[@"groups"][1][@"capacity"] intValue]; ++idx) {
 201         XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Received RateLimiterBadnessClear");
 202         XCTAssertNil(limitTime, @"single object, clear to process right now");
 203     }
 204
 205     XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessGridlocked, @"Received RateLimiterBadnessGridlocked");
 206     XCTAssertNotNil(limitTime, @"After hammering same object need to wait now");
 207     XCTAssertEqualObjects(limitTime, [self.time dateByAddingTimeInterval:[self.config[@"groups"][1][@"rate"] intValue]], @"time: %@, process-OK time is time + rate (%d)", self.time, [self.config[@"groups"][1][@"rate"] intValue]);
 208 }
 209
 210 // Cause it to complain based on too many items in total
 211 - (void)testJudgeRandomItems {
 212     NSDate *limitTime = nil;
 213     TestObject *obj;
 214     for (int idx = 0; idx < [self.config[@"groups"][0][@"capacity"] intValue]; ++idx) {
 215         obj = [TestObject new];
 216         XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Received RateLimiterBadnessClear");
 217         XCTAssertNil(limitTime, @"single object, clear to process right now");
 218     }
 219
 220     XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessCongested, @"Received RateLimiterBadnessCongested");
 221     XCTAssertNotNil(limitTime, @"After hammering same object need to wait now");
 222     XCTAssertEqualObjects(limitTime, [self.time dateByAddingTimeInterval:[self.config[@"groups"][0][@"rate"] intValue]], @"time: %@, process-OK time is time + rate (%d)", self.time, [self.config[@"groups"][0][@"rate"] intValue]);
 223 }
 224
 225 - (void)testOverload {
 226     NSDate *limitTime = nil;
 227     while ([self.RL stateSize] <= [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue]) {
 228         TestObject *obj = [TestObject new];
 229         RateLimiterBadness rlb = [self.RL judge:obj at:self.time limitTime:&limitTime];
 230         XCTAssertTrue(rlb != RateLimiterBadnessOverloaded, @"No issues judging random objects under max state size");
 231     }
 232
 233     // While check is performed at the start of the loop, so now stateSize > maxStateSize. Judge should realize this right away, try to cope, fail and throw a fit
 234     XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessOverloaded, @"RateLimiter overloaded");
 235     XCTAssertEqualObjects(limitTime, [self.time dateByAddingTimeInterval:[self.config[@"general"][@"overloadDuration"] intValue]], @"Overload duration matches expectations");
 236 }
 237
 238 - (void)testTrimmingDueToTime {
 239     NSDate *limitTime = nil;
 240     for (int idx = 0; idx < [self.config[@"general"][@"maxStateSize"] intValue]/2; ++idx) {
 241         TestObject *obj = [TestObject new];
 242         [self.RL judge:obj at:self.time limitTime:&limitTime];
 243     }
 244     NSUInteger stateSize = [self.RL stateSize];
 245     XCTAssertEqual(stateSize, [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue] / 2 + 1, @"Number of objects added matches expectations");
 246     // Advance time enough to age out the existing objects
 247     NSDate *time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] + 1];
 248
 249     // It's been so long, judge should first trim and decide to throw away everything it has
 250     XCTAssertEqual([self.RL judge:self.obj at:time limitTime:&limitTime], RateLimiterBadnessClear, @"New judgment after long time goes fine");
 251     XCTAssertEqual([self.RL stateSize], 2ul, @"Old items gone, just global and one new item left");
 252 }
 253
 254 // RateLimiter is set to ignore properties that return nil
 255 - (void)testNilUuid {
 256     NSDate *limitTime = nil;
 257     TestObject *obj = [[TestObject alloc] initWithNilUuid];
 258     for (int idx = 0; idx < [self.config[@"groups"][0][@"capacity"] intValue]; ++idx) {
 259         XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Same object with nil property only judged on global rate");
 260         XCTAssertEqual([self.RL stateSize], 1ul, @"Nil property objects can't be added to state");
 261     }
 262 }
 263
 264 - (void)testTrimmingDueToSize {
 265     NSDate *limitTime = nil;
 266     // Put first half of items in
 267     for (int idx = 0; idx < [self.config[@"general"][@"maxStateSize"] intValue] / 2; ++idx) {
 268         TestObject *obj = [TestObject new];
 269         [self.RL judge:obj at:self.time limitTime:&limitTime];
 270     }
 271
 272     NSDate *time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] / 2];
 273
 274     // Put second half in later so trim has something to do afterwards
 275     while ([self.RL stateSize] <= [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue]) {
 276         TestObject *obj = [TestObject new];
 277         RateLimiterBadness rlb = [self.RL judge:obj at:time limitTime:&limitTime];
 278         XCTAssertTrue(rlb != RateLimiterBadnessOverloaded, @"No issues judging random objects under max state size");
 279     }
 280
 281     NSUInteger expectedStateSize = [self.RL stateSize] - [self.config[@"general"][@"maxStateSize"] intValue] / 2 + 1;
 282
 283     // Advance time past first batch but before second batch
 284     time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] + 1];
 285     // ...which requires adjusting for the fact that the token buckes will be almost full (i.e. further in the past)
 286     time = [time dateByAddingTimeInterval:-(([self.config[@"groups"][1][@"capacity"] integerValue] - 1) * [self.config[@"groups"][1][@"rate"] integerValue])];
 287
 288     XCTAssertNotEqual([self.RL judge:self.obj at:time limitTime:&limitTime], RateLimiterBadnessOverloaded, @"Judgment caused RL to trim out old items");
 289     XCTAssertEqual([self.RL stateSize], expectedStateSize);
 290 }
 291
 292 @end