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 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     _RL = [[RateLimiter alloc] initWithConfig:_config];
 129     _obj = [TestObject new];
 130     _time = [NSDate date];
 131 }
 132
 133 - (void)testInitWithConfig {
 134     self.RL = [[RateLimiter alloc] initWithConfig:self.config];
 135     XCTAssertNotNil(self.RL, @"RateLimiter with config succeeds");
 136     XCTAssertNil(self.RL.assetType, @"initWithConfig means no assetType");
 137     XCTAssertEqualObjects(self.config, self.RL.config, @"config was copied properly");
 138 }
 139
 140 - (void)testEncodingAndDecoding {
 141     NSDate* date = [NSDate date];
 142     NSDate* limit = nil;
 143     [self.RL judge:self.obj at:date limitTime:&limit];
 144
 145     NSKeyedArchiver *encoder = [[NSKeyedArchiver alloc] initRequiringSecureCoding:YES];
 146     [self.RL encodeWithCoder:encoder];
 147     NSData* data = encoder.encodedData;
 148
 149     XCTAssertEqualObjects(self.config, self.RL.config, @"config unmodified after encoding");
 150     XCTAssertNil(self.RL.assetType, @"assetType still nil after encoding");
 151
 152     NSKeyedUnarchiver *decoder = [[NSKeyedUnarchiver alloc] initForReadingFromData:data error:nil];
 153     RateLimiter *RL2 = [[RateLimiter alloc] initWithCoder:decoder];
 154     XCTAssertNotNil(RL2, @"Received an object from initWithCoder");
 155     XCTAssertEqualObjects(self.RL.config, RL2.config, @"config is the same after encoding and decoding");
 156     XCTAssertTrue([self.RL isEqual:RL2], @"RateLimiters believe they are the same");
 157     XCTAssertNil(RL2.assetType, @"assetType remains nil");
 158 }
 159
 160 - (void)testReset {
 161     NSDate *limitTime = nil;
 162     [self.RL judge:[TestObject new] at:self.time limitTime:&limitTime];
 163     XCTAssertEqual([self.RL stateSize], 2ul, @"Single property judged once, state is 1 global plus 1 property");
 164     [self.RL reset];
 165     XCTAssertEqual([self.RL stateSize], 0ul, @"No buckets after reset");
 166     XCTAssertEqualObjects(self.config, self.RL.config);
 167 }
 168
 169 // Cause it to complain based on one item being hit repeatedly
 170 - (void)testJudgeSingleItem {
 171     NSDate *limitTime = nil;
 172     for (int idx = 0; idx < [self.config[@"groups"][1][@"capacity"] intValue]; ++idx) {
 173         XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Received RateLimiterBadnessClear");
 174         XCTAssertNil(limitTime, @"single object, clear to process right now");
 175     }
 176
 177     XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessGridlocked, @"Received RateLimiterBadnessGridlocked");
 178     XCTAssertNotNil(limitTime, @"After hammering same object need to wait now");
 179     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]);
 180 }
 181
 182 // Cause it to complain based on too many items in total
 183 - (void)testJudgeRandomItems {
 184     NSDate *limitTime = nil;
 185     TestObject *obj;
 186     for (int idx = 0; idx < [self.config[@"groups"][0][@"capacity"] intValue]; ++idx) {
 187         obj = [TestObject new];
 188         XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Received RateLimiterBadnessClear");
 189         XCTAssertNil(limitTime, @"single object, clear to process right now");
 190     }
 191
 192     XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessCongested, @"Received RateLimiterBadnessCongested");
 193     XCTAssertNotNil(limitTime, @"After hammering same object need to wait now");
 194     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]);
 195 }
 196
 197 - (void)testOverload {
 198     NSDate *limitTime = nil;
 199     while ([self.RL stateSize] <= [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue]) {
 200         TestObject *obj = [TestObject new];
 201         RateLimiterBadness rlb = [self.RL judge:obj at:self.time limitTime:&limitTime];
 202         XCTAssertTrue(rlb != RateLimiterBadnessOverloaded, @"No issues judging random objects under max state size");
 203     }
 204
 205     // 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
 206     XCTAssertEqual([self.RL judge:self.obj at:self.time limitTime:&limitTime], RateLimiterBadnessOverloaded, @"RateLimiter overloaded");
 207     XCTAssertEqualObjects(limitTime, [self.time dateByAddingTimeInterval:[self.config[@"general"][@"overloadDuration"] unsignedIntValue]], @"Overload duration matches expectations");
 208 }
 209
 210 - (void)testTrimmingDueToTime {
 211     NSDate *limitTime = nil;
 212     for (int idx = 0; idx < [self.config[@"general"][@"maxStateSize"] intValue]/2; ++idx) {
 213         TestObject *obj = [TestObject new];
 214         [self.RL judge:obj at:self.time limitTime:&limitTime];
 215     }
 216     NSUInteger stateSize = [self.RL stateSize];
 217     XCTAssertEqual(stateSize, [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue] / 2 + 1, @"Number of objects added matches expectations");
 218     // Advance time enough to age out the existing objects
 219     NSDate *time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] + 1];
 220
 221     // It's been so long, judge should first trim and decide to throw away everything it has
 222     XCTAssertEqual([self.RL judge:self.obj at:time limitTime:&limitTime], RateLimiterBadnessClear, @"New judgment after long time goes fine");
 223     XCTAssertEqual([self.RL stateSize], 2ul, @"Old items gone, just global and one new item left");
 224 }
 225
 226 // RateLimiter is set to ignore properties that return nil
 227 - (void)testNilUuid {
 228     NSDate *limitTime = nil;
 229     TestObject *obj = [[TestObject alloc] initWithNilUuid];
 230     for (int idx = 0; idx < [self.config[@"groups"][0][@"capacity"] intValue]; ++idx) {
 231         XCTAssertEqual([self.RL judge:obj at:self.time limitTime:&limitTime], RateLimiterBadnessClear, @"Same object with nil property only judged on global rate");
 232         XCTAssertEqual([self.RL stateSize], 1ul, @"Nil property objects can't be added to state");
 233     }
 234 }
 235
 236 - (void)testTrimmingDueToSize {
 237     NSDate *limitTime = nil;
 238     // Put first half of items in
 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
 244     NSDate *time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] / 2];
 245
 246     // Put second half in later so trim has something to do afterwards
 247     while ([self.RL stateSize] <= [self.config[@"general"][@"maxStateSize"] unsignedIntegerValue]) {
 248         TestObject *obj = [TestObject new];
 249         RateLimiterBadness rlb = [self.RL judge:obj at:time limitTime:&limitTime];
 250         XCTAssertTrue(rlb != RateLimiterBadnessOverloaded, @"No issues judging random objects under max state size");
 251     }
 252
 253     NSUInteger expectedStateSize = [self.RL stateSize] - [self.config[@"general"][@"maxStateSize"] intValue] / 2 + 1;
 254
 255     // Advance time past first batch but before second batch
 256     time = [self.time dateByAddingTimeInterval:[self.config[@"general"][@"maxItemAge"] intValue] + 1];
 257     // ...which requires adjusting for the fact that the token buckes will be almost full (i.e. further in the past)
 258     time = [time dateByAddingTimeInterval:-(([self.config[@"groups"][1][@"capacity"] integerValue] - 1) * [self.config[@"groups"][1][@"rate"] integerValue])];
 259
 260     XCTAssertNotEqual([self.RL judge:self.obj at:time limitTime:&limitTime], RateLimiterBadnessOverloaded, @"Judgment caused RL to trim out old items");
 261     XCTAssertEqual([self.RL stateSize], expectedStateSize);
 262 }
 263
 264 @end