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