1/*
2 * Copyright (c) 2020 Apple Inc. All rights reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#import "unittest_common.h"
18#import "DNSHeuristicsInternal.h"
19#import <XCTest/XCTest.h>
20#import <OCMock/OCMock.h>
21
22@interface DNSHeuristicsTest : XCTestCase
23@end
24
25@implementation DNSHeuristicsTest
26
27#if (TARGET_OS_IPHONE && !TARGET_OS_MACCATALYST)
28- (void)testEmptyStateFailure {
29    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
30    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(@{});
31    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
32    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
33    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
34
35    NSURL *url = [NSURL URLWithString:@"https://example.com"];
36    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
37    OCMVerify(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
38}
39
40- (void)testStateFailureUnderThreshold {
41    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
42    NSUInteger now = [DNSHeuristics currentTimeMs];
43    NSURL *url = [NSURL URLWithString:@"https://example.com"];
44    NSDictionary *existingState = @{
45        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
46        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:1],
47        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
48        DNSHeuristicsFilterFlagKey: @(NO),
49    };
50    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
51    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
52    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
53    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
54
55    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
56    OCMVerify(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
57}
58
59- (void)testStateFailureOverThreshold {
60    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
61    NSUInteger now = [DNSHeuristics currentTimeMs];
62    NSURL *url = [NSURL URLWithString:@"https://example.com"];
63    NSDictionary *existingState = @{
64        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
65        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:DNSHeuristicDefaultLongCounterThreshold], // reporting an error will cause this count to exceed the threshold
66        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
67        DNSHeuristicsFilterFlagKey: @(NO),
68    };
69    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
70    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
71    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
72    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
73
74    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
75    OCMVerify(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
76}
77
78- (void)testStateFailureUnderThreshold_StickAfterFailure {
79    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
80    NSUInteger now = [DNSHeuristics currentTimeMs];
81    NSURL *url = [NSURL URLWithString:@"https://example.com"];
82    NSDictionary *existingState = @{
83        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
84        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:DNSHeuristicDefaultLongCounterThreshold],
85        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
86        DNSHeuristicsFilterFlagKey: @(YES),
87    };
88    OCMStub(ClassMethod([mockHeuristics getNetworkFilteredFlag:[OCMArg anyPointer]])).andReturn(YES);
89    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
90    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
91    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
92    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
93    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + DNSHeuristicDefaultLongCounterTimeWindow * 2); // two days pass, we should reset
94
95    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
96    OCMReject(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
97}
98
99- (void)testStateFailureUnderThreshold_ResetAfterSuccess {
100    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
101    NSUInteger now = [DNSHeuristics currentTimeMs];
102    NSDictionary *existingState = @{
103        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
104        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:DNSHeuristicDefaultLongCounterThreshold],
105        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
106        DNSHeuristicsFilterFlagKey: @(YES),
107    };
108    OCMStub(ClassMethod([mockHeuristics getNetworkFilteredFlag:[OCMArg anyPointer]])).andReturn(YES);
109    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
110    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
111    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
112    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
113    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + DNSHeuristicDefaultLongCounterTimeWindow * 2); // two days pass, we should reset
114
115    XCTAssertTrue([DNSHeuristics updateHeuristicState:YES isTimeout:NO]);
116    OCMVerify(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
117}
118
119- (void)testStateFailureDrainTokenBucket_NoReset {
120    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
121    NSUInteger now = [DNSHeuristics currentTimeMs];
122    NSURL *url = [NSURL URLWithString:@"https://example.com"];
123    NSDictionary *existingState = @{
124        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
125        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:0],
126        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:1],
127        DNSHeuristicsFilterFlagKey: @(NO),
128    };
129    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
130    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
131    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
132    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
133    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + 1); // within the same epoch -- overflow
134
135    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
136    OCMReject(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
137}
138
139- (void)testStateFailureDrainTokenBucket_Reset {
140    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
141    NSUInteger now = [DNSHeuristics currentTimeMs];
142    NSURL *url = [NSURL URLWithString:@"https://example.com"];
143    NSDictionary *existingState = @{
144        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
145        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:0],
146        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:1],
147        DNSHeuristicsFilterFlagKey: @(NO),
148    };
149    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
150    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
151    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
152    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
153    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + DNSHeuristicsDefaultBurstTokenBucketRefillTime + 1); // allow the bucket to replenish
154
155    XCTAssertTrue([DNSHeuristics reportResolutionFailure:url isTimeout:NO]);
156    OCMReject(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
157}
158
159- (void)testStateFailureFilteredThenSuccessBeforeWindow {
160    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
161    NSUInteger now = [DNSHeuristics currentTimeMs];
162    NSDictionary *existingState = @{
163        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:now],
164        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:0],
165        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
166        DNSHeuristicsFilterFlagKey: @(YES),
167    };
168    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
169    OCMStub(ClassMethod([mockHeuristics getNetworkFilteredFlag:[OCMArg anyPointer]])).andReturn(YES);
170    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
171    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
172    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
173    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + DNSHeuristicsDefaultBurstTokenBucketRefillTime + 1); // allow the bucket to replenish
174
175    XCTAssertTrue([DNSHeuristics updateHeuristicState:YES isTimeout:NO]);
176    OCMReject(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
177}
178
179- (void)testStateFailureFilteredThenSuccessAfterWindow {
180    id mockHeuristics = OCMClassMock([DNSHeuristics class]);
181    NSUInteger now = [DNSHeuristics currentTimeMs];
182    NSDictionary *existingState = @{
183        DNSHeuristicsLastFailureTimestamp: [NSNumber numberWithUnsignedInteger:(now - DNSHeuristicDefaultLongCounterTimeWindow)],
184        DNSHeuristicsLongCounterKey: [NSNumber numberWithInt:0],
185        DNSHeuristicsBurstCounterKey: [NSNumber numberWithInt:DNSHeuristicsDefaultBurstTokenBucketCapacity],
186        DNSHeuristicsFilterFlagKey: @(YES),
187    };
188    OCMStub(ClassMethod([mockHeuristics copyNetworkSettings:[OCMArg anyPointer]])).andReturn(existingState);
189    OCMStub(ClassMethod([mockHeuristics getNetworkFilteredFlag:[OCMArg anyPointer]])).andReturn(YES);
190    OCMStub(ClassMethod([mockHeuristics setNetworkSettings:[OCMArg anyPointer] network:[OCMArg anyPointer] value:[OCMArg any]])).andReturn(YES);
191    OCMStub(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
192    OCMStub(ClassMethod([mockHeuristics clearNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]])).andReturn(YES);
193    OCMStub(ClassMethod([mockHeuristics currentTimeMs])).andReturn(now + DNSHeuristicsDefaultBurstTokenBucketRefillTime + 1); // allow the bucket to replenish
194
195    XCTAssertTrue([DNSHeuristics updateHeuristicState:YES isTimeout:NO]);
196    OCMReject(ClassMethod([mockHeuristics setNetworkAsFiltered:[OCMArg anyPointer] network:[OCMArg anyPointer]]));
197}
198
199#endif // TARGET_OS_IPHONE
200
201@end
202