Fix MTProto-MTContext retain cycle

This commit is contained in:
Ali 2023-11-02 17:31:33 +04:00
parent fb2d16a75e
commit 6e6424729a
2 changed files with 149 additions and 72 deletions

View File

@ -675,12 +675,10 @@ private final class NotificationServiceHandler {
private let notificationKeyDisposable = MetaDisposable()
private let pollDisposable = MetaDisposable()
init?(queue: Queue, updateCurrentContent: @escaping (NotificationContent) -> Void, completed: @escaping () -> Void, payload: [AnyHashable: Any]) {
init?(queue: Queue, episode: String, updateCurrentContent: @escaping (NotificationContent) -> Void, completed: @escaping () -> Void, payload: [AnyHashable: Any]) {
//debug_linker_fail_test()
self.queue = queue
let episode = String(UInt32.random(in: 0 ..< UInt32.max), radix: 16)
guard let appBundleIdentifier = Bundle.main.bundleIdentifier, let lastDotRange = appBundleIdentifier.range(of: ".", options: [.backwards]) else {
return nil
}
@ -1558,7 +1556,8 @@ private final class NotificationServiceHandler {
if !shouldSynchronizeState {
pollSignal = .complete()
} else {
stateManager.network.shouldKeepConnection.set(.single(true))
let shouldKeepConnection = stateManager.network.shouldKeepConnection
shouldKeepConnection.set(.single(true))
if peerId.namespace == Namespaces.Peer.CloudChannel {
Logger.shared.log("NotificationService \(episode)", "Will poll channel \(peerId)")
@ -1569,6 +1568,9 @@ private final class NotificationServiceHandler {
peerId: peerId,
stateManager: stateManager
)
|> afterDisposed { [weak shouldKeepConnection] in
shouldKeepConnection?.set(.single(false))
}
} else {
Logger.shared.log("NotificationService \(episode)", "Will perform non-specific getDifference")
enum ControlError {
@ -1584,6 +1586,9 @@ private final class NotificationServiceHandler {
}
}
|> restartIfError
|> afterDisposed { [weak shouldKeepConnection] in
shouldKeepConnection?.set(.single(false))
}
pollSignal = signal
}
@ -1670,6 +1675,9 @@ private final class NotificationServiceHandler {
return
}
let shouldKeepConnection = stateManager.network.shouldKeepConnection
shouldKeepConnection.set(.single(true))
var fetchStoriesSignal: Signal<Void, NoError> = .single(Void())
fetchStoriesSignal = _internal_pollPeerStories(postbox: stateManager.postbox, network: stateManager.network, accountPeerId: stateManager.accountPeerId, peerId: peerId)
|> map { _ -> Void in
@ -1760,6 +1768,9 @@ private final class NotificationServiceHandler {
}
}
)
|> afterDisposed { [weak shouldKeepConnection] in
shouldKeepConnection?.set(.single(false))
}
let fetchMediaSignal: Signal<Data?, NoError> = .single(nil)
@ -1865,11 +1876,6 @@ private final class NotificationServiceHandler {
}
}
let pollSignal: Signal<Never, NoError>
pollSignal = .complete()
stateManager.network.shouldKeepConnection.set(.single(true))
let pollWithUpdatedContent: Signal<NotificationContent, NoError>
if interactionAuthorId != nil || messageId != nil {
pollWithUpdatedContent = stateManager.postbox.transaction { transaction -> NotificationContent in
@ -1899,13 +1905,8 @@ private final class NotificationServiceHandler {
return content
}
|> then(
pollSignal
|> map { _ -> NotificationContent in }
)
} else {
pollWithUpdatedContent = pollSignal
|> map { _ -> NotificationContent in }
pollWithUpdatedContent = .complete()
}
var updatedContent = initialContent
@ -1987,10 +1988,7 @@ private final class NotificationServiceHandler {
}
let completeRemoval: () -> Void = {
guard let strongSelf = self else {
return
}
var content = NotificationContent(isLockedMessage: nil)
let content = NotificationContent(isLockedMessage: nil)
Logger.shared.log("NotificationService \(episode)", "Updating content to \(content)")
updateCurrentContent(content)
@ -2126,12 +2124,16 @@ final class NotificationService: UNNotificationServiceExtension {
private var initialContent: UNNotificationContent?
private let content = Atomic<NotificationContent?>(value: nil)
private var contentHandler: ((UNNotificationContent) -> Void)?
private var episode: String?
override init() {
super.init()
}
override func didReceive(_ request: UNNotificationRequest, withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
let episode = String(UInt32.random(in: 0 ..< UInt32.max), radix: 16)
self.episode = episode
self.initialContent = request.content
self.contentHandler = contentHandler
@ -2142,6 +2144,7 @@ final class NotificationService: UNNotificationServiceExtension {
self.impl = QueueLocalObject(queue: queue, generate: { [weak self] in
return BoxedNotificationServiceHandler(value: NotificationServiceHandler(
queue: queue,
episode: episode,
updateCurrentContent: { value in
let _ = content.swap(value)
},
@ -2152,15 +2155,17 @@ final class NotificationService: UNNotificationServiceExtension {
strongSelf.impl = nil
if let contentHandler = strongSelf.contentHandler {
Logger.shared.log("NotificationService \(episode)", "Complete handling notification")
strongSelf.contentHandler = nil
if let content = content.with({ $0 }) {
/*let request = UNNotificationRequest(identifier: UUID().uuidString, content: content.generate(), trigger: .none)
UNUserNotificationCenter.current().add(request)
contentHandler(UNMutableNotificationContent())*/
contentHandler(content.generate())
} else if let initialContent = strongSelf.initialContent {
contentHandler(initialContent)
}
} else {
Logger.shared.log("NotificationService \(episode)", "Attempted to repeatedly complete handling notification")
}
},
payload: request.content.userInfo
@ -2170,6 +2175,10 @@ final class NotificationService: UNNotificationServiceExtension {
override func serviceExtensionTimeWillExpire() {
if let contentHandler = self.contentHandler {
self.contentHandler = nil
Logger.shared.log("NotificationService \(self.episode ?? "???")", "Completing due to serviceExtensionTimeWillExpire")
if let content = self.content.with({ $0 }) {
contentHandler(content.generate())
} else if let initialContent = self.initialContent {

View File

@ -141,6 +141,24 @@ static MTDatacenterAuthInfoMapKeyStruct parseAuthInfoMapKeyInteger(NSNumber *key
return parseAuthInfoMapKey([key longLongValue]);
}
@interface MTWeakContextChangeListener : NSObject
@property (nonatomic, weak) id<MTContextChangeListener> target;
@end
@implementation MTWeakContextChangeListener
- (instancetype)initWithTarget:(id<MTContextChangeListener>)target {
self = [super init];
if (self != nil) {
_target = target;
}
return self;
}
@end
@interface MTContext () <MTDiscoverDatacenterAddressActionDelegate, MTDatacenterTransferAuthActionDelegate>
{
int64_t _uniqueId;
@ -160,7 +178,7 @@ static MTDatacenterAuthInfoMapKeyStruct parseAuthInfoMapKeyInteger(NSNumber *key
NSMutableDictionary *_authTokenById;
NSMutableArray *_changeListeners;
NSMutableArray<MTWeakContextChangeListener *> *_changeListeners;
MTSignal *_discoverBackupAddressListSignal;
@ -463,9 +481,17 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
{
[[MTContext contextQueue] dispatchOnQueue:^
{
if (![_changeListeners containsObject:changeListener])
{
[_changeListeners addObject:changeListener];
bool alreadyContains = false;
for (MTWeakContextChangeListener *value in _changeListeners) {
id<MTContextChangeListener> target = value.target;
if (target == changeListener) {
alreadyContains = true;
break;
}
}
if (!alreadyContains) {
[_changeListeners addObject:[[MTWeakContextChangeListener alloc] initWithTarget:changeListener]];
}
}];
}
@ -474,7 +500,15 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
{
[[MTContext contextQueue] dispatchOnQueue:^
{
[_changeListeners removeObject:changeListener];
for (NSInteger i = 0; i < _changeListeners.count; i++) {
MTWeakContextChangeListener *value = _changeListeners[i];
id<MTContextChangeListener> target = value.target;
if (target == changeListener) {
[_changeListeners removeObjectAtIndex:i];
break;
}
}
} synchronous:true];
}
@ -544,19 +578,25 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
_datacenterAddressSetById[@(datacenterId)] = addressSet;
[_keychain setObject:_datacenterAddressSetById forKey:@"datacenterAddressSetById" group:@"persistent"];
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
NSArray<MTWeakContextChangeListener *> *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
if ([listener respondsToSelector:@selector(contextDatacenterAddressSetUpdated:datacenterId:addressSet:)])
[listener contextDatacenterAddressSetUpdated:self datacenterId:datacenterId addressSet:addressSet];
for (MTWeakContextChangeListener *listenerWrapper in changeListeners) {
id<MTContextChangeListener> listener = listenerWrapper.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterAddressSetUpdated:datacenterId:addressSet:)]) {
[listener contextDatacenterAddressSetUpdated:self datacenterId:datacenterId addressSet:addressSet];
}
}
}
if (true) {
bool shouldReset = previousAddressSetWasEmpty || updateSchemes;
for (id<MTContextChangeListener> listener in currentListeners) {
if ([listener respondsToSelector:@selector(contextDatacenterTransportSchemesUpdated:datacenterId:shouldReset:)]) {
[listener contextDatacenterTransportSchemesUpdated:self datacenterId:datacenterId shouldReset:shouldReset];
for (MTWeakContextChangeListener *listenerWrapper in changeListeners) {
id<MTContextChangeListener> listener = listenerWrapper.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterTransportSchemesUpdated:datacenterId:shouldReset:)]) {
[listener contextDatacenterTransportSchemesUpdated:self datacenterId:datacenterId shouldReset:shouldReset];
}
}
}
} else {
@ -633,12 +673,14 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
_datacenterAddressSetById[@(datacenterId)] = addressSet;
[_keychain setObject:_datacenterAddressSetById forKey:@"datacenterAddressSetById" group:@"persistent"];
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
if ([listener respondsToSelector:@selector(contextDatacenterAddressSetUpdated:datacenterId:addressSet:)])
[listener contextDatacenterAddressSetUpdated:self datacenterId:datacenterId addressSet:addressSet];
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterAddressSetUpdated:datacenterId:addressSet:)]) {
[listener contextDatacenterAddressSetUpdated:self datacenterId:datacenterId addressSet:addressSet];
}
}
}
}
}];
@ -671,12 +713,13 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
[_keychain setObject:_datacenterAuthInfoById forKey:@"datacenterAuthInfoById" group:@"persistent"];
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
if ([listener respondsToSelector:@selector(contextDatacenterAuthInfoUpdated:datacenterId:authInfo:selector:)])
[listener contextDatacenterAuthInfoUpdated:self datacenterId:datacenterId authInfo:authInfo selector:selector];
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterAuthInfoUpdated:datacenterId:authInfo:selector:)])
[listener contextDatacenterAuthInfoUpdated:self datacenterId:datacenterId authInfo:authInfo selector:selector];
}
}
if (wasNil && authInfo != nil && selector == MTDatacenterAuthInfoSelectorPersistent) {
@ -712,12 +755,14 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
{
[[MTContext contextQueue] dispatchOnQueue:^
{
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
if ([listener respondsToSelector:@selector(contextIsPasswordRequiredUpdated:datacenterId:)])
[listener contextIsPasswordRequiredUpdated:self datacenterId:datacenterId];
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextIsPasswordRequiredUpdated:datacenterId:)]) {
[listener contextIsPasswordRequiredUpdated:self datacenterId:datacenterId];
}
}
}
}];
}
@ -738,15 +783,18 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
return current;
}];
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
if (MTLogEnabled()) {
MTLog(@"[MTContext#%" PRIxPTR ": %@ transport scheme updated for %d: %@]", (intptr_t)self, media ? @"media" : @"generic", datacenterId, transportScheme);
}
for (id<MTContextChangeListener> listener in currentListeners) {
if ([listener respondsToSelector:@selector(contextDatacenterTransportSchemesUpdated:datacenterId:shouldReset:)])
[listener contextDatacenterTransportSchemesUpdated:self datacenterId:datacenterId shouldReset:false];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterTransportSchemesUpdated:datacenterId:shouldReset:)])
[listener contextDatacenterTransportSchemesUpdated:self datacenterId:datacenterId shouldReset:false];
}
}
}
}];
@ -1011,9 +1059,12 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
_datacenterPublicKeysById[@(datacenterId)] = publicKeys;
[_keychain setObject:_datacenterPublicKeysById forKey:@"datacenterPublicKeysById" group:@"ephemeral"];
for (id<MTContextChangeListener> listener in _changeListeners) {
if ([listener respondsToSelector:@selector(contextDatacenterPublicKeysUpdated:datacenterId:publicKeys:)]) {
[listener contextDatacenterPublicKeysUpdated:self datacenterId:datacenterId publicKeys:publicKeys];
for (MTWeakContextChangeListener *value in _changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (listener) {
if ([listener respondsToSelector:@selector(contextDatacenterPublicKeysUpdated:datacenterId:publicKeys:)]) {
[listener contextDatacenterPublicKeysUpdated:self datacenterId:datacenterId publicKeys:publicKeys];
}
}
}
}
@ -1023,7 +1074,11 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
- (void)publicKeysForDatacenterWithIdRequired:(NSInteger)datacenterId {
[[MTContext contextQueue] dispatchOnQueue:^{
if (_fetchPublicKeysActions[@(datacenterId)] == nil) {
for (id<MTContextChangeListener> listener in _changeListeners) {
for (MTWeakContextChangeListener *value in _changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (!listener) {
continue;
}
if ([listener respondsToSelector:@selector(fetchContextDatacenterPublicKeys:datacenterId:)]) {
MTSignal *signal = [listener fetchContextDatacenterPublicKeys:self datacenterId:datacenterId];
if (signal != nil) {
@ -1136,7 +1191,11 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
MTDatacenterAddressSet *addressSet = [[MTDatacenterAddressSet alloc] initWithAddressList:addressList];
MTSignal *discoverSignal = [MTDiscoverConnectionSignals discoverSchemeWithContext:self datacenterId:datacenterId addressList:addressSet.addressList media:media isProxy:isProxy];
MTSignal *conditionSignal = [MTSignal single:@(true)];
for (id<MTContextChangeListener> listener in _changeListeners) {
for (MTWeakContextChangeListener *value in _changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (!listener) {
continue;
}
if ([listener respondsToSelector:@selector(isContextNetworkAccessAllowed:)]) {
MTSignal *signal = [listener isContextNetworkAccessAllowed:self];
if (signal != nil) {
@ -1338,10 +1397,12 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
[_authTokenById removeObjectForKey:@(datacenterId)];
[_keychain setObject:_authTokenById forKey:@"authTokenById" group:@"persistent"];
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (!listener) {
continue;
}
if ([listener respondsToSelector:@selector(contextDatacenterAuthTokenUpdated:datacenterId:authToken:)])
[listener contextDatacenterAuthTokenUpdated:self datacenterId:datacenterId authToken:authToken];
}
@ -1467,9 +1528,12 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
if (apiEnvironment != nil) {
_apiEnvironment = apiEnvironment;
NSArray *currentListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners)
{
NSArray *changeListeners = [[NSArray alloc] initWithArray:_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (!listener) {
continue;
}
if ([listener respondsToSelector:@selector(contextApiEnvironmentUpdated:apiEnvironment:)]) {
[listener contextApiEnvironmentUpdated:self apiEnvironment:apiEnvironment];
}
@ -1503,8 +1567,12 @@ static void copyKeychainDictionaryKey(NSString * _Nonnull group, NSString * _Non
}
if ([isRemoved boolValue]) {
NSArray *currentListeners = [[NSArray alloc] initWithArray:strongSelf->_changeListeners];
for (id<MTContextChangeListener> listener in currentListeners) {
NSArray *changeListeners = [[NSArray alloc] initWithArray:strongSelf->_changeListeners];
for (MTWeakContextChangeListener *value in changeListeners) {
id<MTContextChangeListener> listener = value.target;
if (!listener) {
continue;
}
if ([listener respondsToSelector:@selector(contextLoggedOut:)])
[listener contextLoggedOut:strongSelf];
}