From cc97fc2fa2a6fd875c2af18cbad65a365e9ffdd0 Mon Sep 17 00:00:00 2001
From: Peter <>
Date: Wed, 19 Dec 2018 02:51:05 +0300
Subject: [PATCH] Support for alternative scheme variants

---
 MTDiscoverConnectionSignals.m |  2 +-
 MTProtoKit/MTContext.m        | 77 ++++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/MTDiscoverConnectionSignals.m b/MTDiscoverConnectionSignals.m
index 3f86c5828d..2b68317280 100644
--- a/MTDiscoverConnectionSignals.m
+++ b/MTDiscoverConnectionSignals.m
@@ -218,7 +218,7 @@
             }];
             [bestTcp4Signals addObject:signal];
             
-            NSArray *alternatePorts = @[@80, /*@5222*/];
+            NSArray *alternatePorts = @[@80, @5222];
             for (NSNumber *nPort in alternatePorts) {
                 NSSet *ipsWithPort = tcpIpsByPort[nPort];
                 if (![ipsWithPort containsObject:address.ip]) {
diff --git a/MTProtoKit/MTContext.m b/MTProtoKit/MTContext.m
index bfe790773a..ef5b4145fa 100644
--- a/MTProtoKit/MTContext.m
+++ b/MTProtoKit/MTContext.m
@@ -66,6 +66,63 @@
 
 @end
 
+@interface MTTransportSchemeKey : NSObject<NSCoding, NSCopying>
+
+@property (nonatomic, readonly) NSInteger datacenterId;
+@property (nonatomic, readonly) bool isProxy;
+@property (nonatomic, readonly) bool isMedia;
+
+@end
+
+@implementation MTTransportSchemeKey
+
+- (instancetype)initWithDatacenterId:(NSInteger)datacenterId isProxy:(bool)isProxy isMedia:(bool)isMedia {
+    self = [super init];
+    if (self != nil) {
+        _datacenterId = datacenterId;
+        _isProxy = isProxy;
+        _isMedia = isMedia;
+    }
+    return self;
+}
+
+- (instancetype)initWithCoder:(NSCoder *)aDecoder {
+    return [self initWithDatacenterId:[aDecoder decodeIntegerForKey:@"datacenterId"] isProxy:[aDecoder decodeBoolForKey:@"isProxy"] isMedia:[aDecoder decodeBoolForKey:@"isMedia"]];
+}
+
+- (void)encodeWithCoder:(NSCoder *)aCoder {
+    [aCoder encodeInteger:_datacenterId forKey:@"datacenterId"];
+    [aCoder encodeBool:_isProxy forKey:@"isProxy"];
+    [aCoder encodeBool:_isMedia forKey:@"isMedia"];
+}
+
+- (instancetype)copyWithZone:(NSZone *)__unused zone {
+    return self;
+}
+
+- (NSUInteger)hash {
+    return _datacenterId * 31 * 31 + (_isProxy ? 1 : 0) * 31 + (_isMedia ? 1 : 0);
+}
+
+- (BOOL)isEqual:(id)object {
+    if (![object isKindOfClass:[MTTransportSchemeKey class]]) {
+        return false;
+    }
+    MTTransportSchemeKey *other = (MTTransportSchemeKey *)object;
+    if (_datacenterId != other->_datacenterId) {
+        return false;
+    }
+    if (_isProxy != other->_isProxy) {
+        return false;
+    }
+    if (_isMedia != other->_isMedia) {
+        return false;
+    }
+    return true;
+}
+
+@end
+
 @interface MTContext () <MTDiscoverDatacenterAddressActionDelegate, MTDatacenterAuthActionDelegate, MTDatacenterTransferAuthActionDelegate>
 {
     int64_t _uniqueId;
@@ -74,6 +131,7 @@
     
     NSMutableDictionary *_datacenterSeedAddressSetById;
     NSMutableDictionary *_datacenterAddressSetById;
+    NSMutableDictionary<MTTransportSchemeKey *, MTTransportScheme *> *_datacenterManuallySelectedSchemeById;
     
     NSMutableDictionary<NSNumber *, NSMutableDictionary<MTDatacenterAddress *, MTTransportSchemeStats *> *> *_transportSchemeStats;
     MTTimer *_schemeStatsSyncTimer;
@@ -143,6 +201,7 @@
         _datacenterSeedAddressSetById = [[NSMutableDictionary alloc] init];
         
         _datacenterAddressSetById = [[NSMutableDictionary alloc] init];
+        _datacenterManuallySelectedSchemeById = [[NSMutableDictionary alloc] init];
         
         _transportSchemeStats = [[NSMutableDictionary alloc] init];
         
@@ -273,6 +332,15 @@
                 }
             }
             
+            NSDictionary *datacenterManuallySelectedSchemeById = [keychain objectForKey:@"datacenterManuallySelectedSchemeById_v1" group:@"persistent"];
+            if (datacenterManuallySelectedSchemeById != nil) {
+                _datacenterManuallySelectedSchemeById = [[NSMutableDictionary alloc] initWithDictionary:datacenterManuallySelectedSchemeById];
+                if (MTLogEnabled()) {
+                    MTLog(@"[MTContext loaded datacenterManuallySelectedSchemeById: %@]", _datacenterManuallySelectedSchemeById);
+                }
+            }
+            
+            
             [_apiEnvironment.datacenterAddressOverrides enumerateKeysAndObjectsUsingBlock:^(NSNumber *nDatacenterId, MTDatacenterAddress *address, __unused BOOL *stop) {
                 _datacenterAddressSetById[nDatacenterId] = [[MTDatacenterAddressSet alloc] initWithAddressList:@[address]];
             }];
@@ -561,6 +629,9 @@
 - (void)updateTransportSchemeForDatacenterWithId:(NSInteger)datacenterId transportScheme:(MTTransportScheme *)transportScheme media:(bool)media isProxy:(bool)isProxy {
     [[MTContext contextQueue] dispatchOnQueue:^ {
         if (transportScheme != nil && datacenterId != 0) {
+            _datacenterManuallySelectedSchemeById[[[MTTransportSchemeKey alloc] initWithDatacenterId:datacenterId isProxy:isProxy isMedia:media]] = transportScheme;
+            [_keychain setObject:_datacenterManuallySelectedSchemeById forKey:@"datacenterManuallySelectedSchemeById_v1" group:@"persistent"];
+            
             [self reportTransportSchemeSuccessForDatacenterId:datacenterId transportScheme:transportScheme];
             [self _withTransportSchemeStatsForDatacenterId:datacenterId transportScheme:transportScheme process:^MTTransportSchemeStats *(MTTransportSchemeStats *current) {
                 current = [current withUpdatedLastFailureTimestamp:0];
@@ -779,6 +850,10 @@
         } else {
             [results addObjectsFromArray:[self _allTransportSchemesForDatacenterWithId:datacenterId]];
         }
+        MTTransportScheme *manualScheme = _datacenterManuallySelectedSchemeById[[[MTTransportSchemeKey alloc] initWithDatacenterId:datacenterId isProxy:isProxy isMedia:media]];
+        if (manualScheme != nil && ![results containsObject:manualScheme]) {
+            [results addObject:manualScheme];
+        }
     } synchronous:true];
     
     for (int i = (int)(results.count - 1); i >= 0; i--) {
@@ -950,7 +1025,7 @@
                         [strongSelf->_transportSchemeDisposableByDatacenterId removeObjectForKey:@(datacenterId)];
                     }];
                 }
-            }] startWithNext:^(id next)
+            }] startWithNext:^(MTTransportScheme *next)
             {
                 if (MTLogEnabled()) {
                     MTLog(@"scheme: %@", next);