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