keychain/ckks/CKKSRateLimiter.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 #if OCTAGON
  25 #import "CKKSRateLimiter.h"
  26 #import <utilities/debugging.h>
  27 #import <TargetConditionals.h>
  28
  29 typedef NS_ENUM(int, BucketType) {
  30     All,
  31     Group,
  32     UUID,
  33 };
  34
  35 @interface CKKSRateLimiter()
  36 @property (readwrite) NSDictionary<NSString *, NSNumber *> *config;
  37 @property NSMutableDictionary<NSString *, NSDate *> *buckets;
  38 @property NSDate *overloadUntil;
  39 @end
  40
  41 @implementation CKKSRateLimiter
  42
  43 - (instancetype)init {
  44     return [self initWithCoder:nil];
  45 }
  46
  47 - (instancetype)initWithCoder:(NSCoder *)coder {
  48     self = [super init];
  49     if (self) {
  50         if (coder) {
  51             _buckets = [coder decodeObjectOfClasses:[NSSet setWithObjects:[NSMutableDictionary class],
  52                                                                           [NSString class],
  53                                                                           [NSDate class],
  54                                                                           nil]
  55                                                                    forKey:@"buckets"];
  56         } else {
  57             _buckets = [NSMutableDictionary new];
  58         }
  59         _overloadUntil = nil;
  60         // this should be done from a downloadable plist, rdar://problem/29945628
  61         _config = [NSDictionary dictionaryWithObjectsAndKeys:
  62                    @30  , @"rateAll",
  63                    @120 , @"rateGroup",
  64                    @600 , @"rateUUID",
  65                    @20  , @"capacityAll",
  66                    @10  , @"capacityGroup",
  67                    @3   , @"capacityUUID",
  68                    @250 , @"trimSize",
  69                    @3600, @"trimTime",
  70                    @1800, @"overloadDuration", nil];
  71     }
  72     return self;
  73 }
  74
  75 - (BOOL)isEqual: (id) object {
  76     if(![object isKindOfClass:[CKKSRateLimiter class]]) {
  77         return NO;
  78     }
  79
  80     CKKSRateLimiter* obj = (CKKSRateLimiter*) object;
  81
  82     return ([self.config isEqual: obj.config] &&
  83             [self.buckets isEqual: obj.buckets] &&
  84             ((self.overloadUntil == nil && obj.overloadUntil == nil) || ([self.overloadUntil isEqual: obj.overloadUntil]))) ? YES : NO;
  85 }
  86
  87 - (int)rate:(enum BucketType)type {
  88     switch (type) {
  89         case All:
  90             return [self.config[@"rateAll"] intValue];
  91         case Group:
  92             return [self.config[@"rateGroup"] intValue];
  93         case UUID:
  94             return [self.config[@"rateUUID"] intValue];
  95     }
  96 }
  97
  98 - (int)capacity:(enum BucketType)type {
  99     switch (type) {
 100         case All:
 101             return [self.config[@"capacityAll"] intValue];
 102         case Group:
 103             return [self.config[@"capacityGroup"] intValue];
 104         case UUID:
 105             return [self.config[@"capacityUUID"] intValue];
 106     }
 107 }
 108
 109 - (NSDate *)consumeTokenFromBucket:(NSString *)name
 110                               type:(enum BucketType)type
 111                                 at:(NSDate *)time {
 112     NSDate *threshold = [time dateByAddingTimeInterval:-([self capacity:type] * [self rate:type])];
 113     NSDate *bucket = self.buckets[name];
 114
 115     if (!bucket || [bucket timeIntervalSinceDate:threshold] < 0) {
 116         bucket = threshold;
 117     }
 118
 119     // Implicitly track the number of tokens in the bucket.
 120     // "Would the token I need have been generated in the past or in the future?"
 121     bucket = [bucket dateByAddingTimeInterval:[self rate:type]];
 122     self.buckets[name] = bucket;
 123     return ([bucket timeIntervalSinceDate:time] <= 0) ? nil : [bucket copy];
 124 }
 125
 126 - (int)judge:(CKKSOutgoingQueueEntry * _Nonnull const)entry
 127           at:(NSDate * _Nonnull)time
 128    limitTime:(NSDate * _Nonnull __autoreleasing * _Nonnull) limitTime
 129 {
 130     if (self.overloadUntil) {
 131         if ([time timeIntervalSinceDate:self.overloadUntil] >= 0) {
 132             [self trim:time];
 133         }
 134         if (self.overloadUntil) {
 135             *limitTime = [self.overloadUntil copy];
 136             return 5;
 137         }
 138     }
 139
 140     NSDate *all = self.buckets[@"All"];
 141     if ((all && [time timeIntervalSinceDate:all] > [self.config[@"trimTime"] intValue]) ||
 142         self.buckets.count >= [self.config[@"trimSize"] unsignedIntValue]) {
 143         [self trim:time];
 144         if (self.overloadUntil) {
 145             *limitTime = self.overloadUntil;
 146             return 5;
 147         }
 148     }
 149
 150     int badness = 0;
 151     NSDate *sendTime = [self consumeTokenFromBucket:@"All" type:All at:time];
 152     if (sendTime) {
 153         badness = 1;
 154     }
 155     NSDate *backoff = [self consumeTokenFromBucket:[NSString stringWithFormat:@"G:%@", entry.accessgroup] type:Group at:time];
 156     if (backoff) {
 157         sendTime = sendTime == nil ? backoff : [sendTime laterDate:backoff];
 158         badness = ([backoff timeIntervalSinceDate:
 159                     [time dateByAddingTimeInterval:([self rate:Group] * 2)]] < 0) ? 2 : 3;
 160     }
 161     backoff = [self consumeTokenFromBucket:[NSString stringWithFormat:@"U:%@", entry.uuid] type:UUID at:time];
 162     if (backoff) {
 163         sendTime = sendTime == nil ? backoff : [sendTime laterDate:backoff];
 164         badness = 4;
 165     }
 166
 167     *limitTime = sendTime;
 168     return badness;
 169 }
 170
 171 - (NSUInteger)stateSize {
 172     return self.buckets.count;
 173 }
 174
 175 - (void)reset {
 176     self.buckets = [NSMutableDictionary new];
 177     self.overloadUntil = nil;
 178 }
 179
 180 - (void)trim:(NSDate *)time {
 181     int threshold = [self.config[@"trimTime"] intValue];
 182     NSSet *toRemove = [self.buckets keysOfEntriesPassingTest:^BOOL(NSString *key, NSDate *obj, BOOL *stop) {
 183         return [time timeIntervalSinceDate:obj] > threshold;
 184     }];
 185
 186     // Nothing to remove means everybody keeps being noisy. Tell them to go away.
 187     if ([toRemove count] == 0) {
 188         self.overloadUntil = [self.buckets[@"All"] dateByAddingTimeInterval:[self.config[@"overloadDuration"] intValue]];
 189         seccritical("RateLimiter overloaded until %@", self.overloadUntil);
 190     } else {
 191         self.overloadUntil = nil;
 192         [self.buckets removeObjectsForKeys:[toRemove allObjects]];
 193     }
 194 }
 195
 196 - (void)encodeWithCoder:(NSCoder *)coder {
 197     [coder encodeObject:self.buckets forKey:@"buckets"];
 198 }
 199
 200 - (NSString *)diagnostics {
 201     NSMutableString *diag = [NSMutableString stringWithFormat:@"RateLimiter config: %@\n", [self.config description]];
 202
 203     if (self.overloadUntil != nil) {
 204         [diag appendFormat:@"Overloaded until %@, %lu total buckets\n", self.overloadUntil, (unsigned long)[self.buckets count]];
 205     } else {
 206         [diag appendFormat:@"Not overloaded, %lu total buckets\n", (unsigned long)[self.buckets count]];
 207     }
 208
 209     NSArray *offenders = [self topOffendingAccessGroups:10];
 210     if (offenders) {
 211         [diag appendFormat:@"%lu congested buckets. Top offenders: \n%@ range %@ to %@\n",
 212          (unsigned long)[offenders count], offenders, self.buckets[offenders[0]], self.buckets[offenders[[offenders count] - 1]]];
 213     } else {
 214         [diag appendString:@"No buckets congested"];
 215     }
 216
 217     return diag;
 218 }
 219
 220 - (NSArray *)topOffendingAccessGroups:(NSUInteger)num {
 221     NSDate *now = [NSDate date];
 222     NSSet *congestedKeys = [self.buckets keysOfEntriesPassingTest:^BOOL(NSString *key, NSDate *obj, BOOL *stop) {
 223         if (![key hasPrefix:@"G:"]) {
 224             return NO;
 225         }
 226         return [now timeIntervalSinceDate:obj] <= 0 ? NO : YES;
 227     }];
 228
 229     if ([congestedKeys count] > 0) {
 230         // Marker must be type NSDate but can be anything since we know all objects will be in the dictionary
 231         NSDictionary *congested = [NSDictionary dictionaryWithObjects:[self.buckets objectsForKeys:[congestedKeys allObjects]
 232                                                                                     notFoundMarker:[NSDate date]]
 233                                                               forKeys:[congestedKeys allObjects]];
 234         NSArray *sortedKeys = [[[congested keysSortedByValueUsingSelector:@selector(compare:)] reverseObjectEnumerator] allObjects];
 235         if ([sortedKeys count] > num) {
 236             return [sortedKeys subarrayWithRange:NSMakeRange(0, num)];
 237         } else {
 238             return sortedKeys;
 239         }
 240     } else {
 241         return nil;
 242     }
 243 }
 244
 245 + (BOOL)supportsSecureCoding {
 246     return YES;
 247 }
 248
 249 @end
 250
 251 #endif // OCTAGON