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