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