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