Adds websocket updates via the /updates endpoint

kordophone/Bridge/MBIMHTTPConnection.m

@@ -12,6 +12,7 @@
 #import "MBIMBridge_Private.h"
 #import "MBIMBridgeOperation.h"
 #import "MBIMAuthToken.h"
+#import "MBIMUpdateQueue.h"
 
 #import <Security/Security.h>
 #import <CocoaHTTPServer/HTTPMessage.h>
@@ -107,6 +108,7 @@
     if (operationClass != nil) {
         _currentOperation = [[operationClass alloc] initWithRequestURL:url completion:completion];
         _currentOperation.requestBodyData = _bodyData;
+        _currentOperation.request = self->request;
         
         [[[MBIMBridge sharedInstance] operationQueue] addOperation:_currentOperation];
         long status = dispatch_semaphore_wait(sema, dispatch_time(DISPATCH_TIME_NOW, (int64_t)(60.0 * NSEC_PER_SEC)));
@@ -124,6 +126,17 @@
     return response;
 }
 
+- (WebSocket *)webSocketForURI:(NSString *)path
+{
+    NSURL *url = [NSURL URLWithString:path];
+    NSString *endpointName = [url lastPathComponent];
+    if ([endpointName isEqualToString:@"updates"]) {
+        return [[MBIMUpdateQueue sharedInstance] vendUpdateWebSocketConsumerForRequest:request socket:asyncSocket];
+    }
+    
+    return [super webSocketForURI:path];
+}
+
 - (BOOL)expectsRequestBodyFromMethod:(NSString *)method atPath:(NSString *)path
 {
     if ([method isEqualToString:@"POST"]) {

kordophone/Bridge/MBIMUpdateQueue.h

@@ -10,6 +10,9 @@
 #import "IMFoundation_ClassDump.h"
 
 NS_ASSUME_NONNULL_BEGIN
+@class GCDAsyncSocket;
+@class HTTPMessage;
+@class WebSocket;
 
 @interface MBIMUpdateItem : NSObject
 @property (nonatomic, strong) IMChat *changedChat;
@@ -24,11 +27,13 @@ typedef void (^MBIMUpdateConsumer)(NSArray<MBIMUpdateItem *> *items);
 
 + (instancetype)sharedInstance;
 
-- (void)addConsumer:(MBIMUpdateConsumer)consumer withLastSyncedMessageSeq:(NSInteger)messageSeq;
-- (void)removeConsumer:(MBIMUpdateConsumer)consumer;
+- (void)addPollingConsumer:(MBIMUpdateConsumer)consumer withLastSyncedMessageSeq:(NSInteger)messageSeq;
+- (void)removePollingConsumer:(MBIMUpdateConsumer)consumer;
 
 - (void)enqueueUpdateItem:(MBIMUpdateItem *)item;
 
+- (WebSocket *)vendUpdateWebSocketConsumerForRequest:(HTTPMessage *)request socket:(GCDAsyncSocket *)socket;
+
 @end
 
 NS_ASSUME_NONNULL_END

kordophone/Bridge/MBIMUpdateQueue.m

@@ -9,20 +9,29 @@
 #import "MBIMUpdateQueue.h"
 #import "IMMessageItem+Encoded.h"
 #import "IMChat+Encoded.h"
+#import "MBIMHTTPConnection.h"
+#import "MBIMURLUtilities.h"
+
+#import <CocoaHTTPServer/GCDAsyncSocket.h>
+#import <CocoaHTTPServer/HTTPMessage.h>
+#import <CocoaHTTPServer/WebSocket.h>
 
 static const NSUInteger kUpdateItemsCullingLength = 100;
 
-@interface MBIMUpdateItem (/*INTERNAL*/)
+@interface MBIMUpdateItem (/*INTERNAL*/) <WebSocketDelegate>
 @property (nonatomic, assign) NSUInteger messageSequenceNumber;
 @end
 
 @implementation MBIMUpdateQueue {
     NSUInteger           _messageSequenceNumber;
     dispatch_queue_t     _accessQueue;
-    NSMutableArray      *_consumers;
+    NSMutableArray      *_longPollConsumers;
     
     // Maps message sequence number to update item
     NSMutableDictionary<NSNumber *, MBIMUpdateItem *> *_updateItemHistory;
+    
+    // WebSocket consumers
+    NSMutableDictionary<NSString *, MBIMUpdateConsumer> *_websocketConsumers;
 }
 
 + (instancetype)sharedInstance
@@ -41,20 +50,93 @@ static const NSUInteger kUpdateItemsCullingLength = 100;
     self = [super init];
     if (self) {
         _accessQueue = dispatch_queue_create("net.buzzert.MBIMUpdateQueue", DISPATCH_QUEUE_SERIAL);
-        _consumers = [[NSMutableArray alloc] init];
         _messageSequenceNumber = 0;
         _updateItemHistory = [[NSMutableDictionary alloc] init];
+        _websocketConsumers = [[NSMutableDictionary alloc] init];
+        _longPollConsumers = [[NSMutableArray alloc] init];
     }
     
     return self;
 }
 
-- (void)addConsumer:(MBIMUpdateConsumer)consumer withLastSyncedMessageSeq:(NSInteger)messageSeq
+- (void)addPollingConsumer:(MBIMUpdateConsumer)consumer withLastSyncedMessageSeq:(NSInteger)messageSeq
 {
-    __weak NSMutableArray *consumers = _consumers;
+    if (![self _syncConsumer:consumer fromLastMessageSeq:messageSeq]) {
+        __weak NSMutableArray *consumers = _longPollConsumers;
+        dispatch_async(_accessQueue, ^{
+            [consumers addObject:consumer];
+        });
+    }
+}
+
+- (void)removePollingConsumer:(MBIMUpdateConsumer)consumer
+{
+    __weak NSMutableArray *consumers = _longPollConsumers;
+    dispatch_async(_accessQueue, ^{
+        [consumers removeObject:consumer];
+    });
+}
+
+- (void)enqueueUpdateItem:(MBIMUpdateItem *)item
+{
+    __weak __auto_type pollingConsumers = _longPollConsumers;
+    __weak __auto_type websocketConsumers = _websocketConsumers;
     __weak NSMutableDictionary *updateItemHistory = _updateItemHistory;
     dispatch_async(_accessQueue, ^{
-        if ((messageSeq >= 0) && messageSeq < self->_messageSequenceNumber) {
+        self->_messageSequenceNumber++;
+        item.messageSequenceNumber = self->_messageSequenceNumber;
+        
+        // Notify polling consumers
+        for (MBIMUpdateConsumer consumer in pollingConsumers) {
+            consumer(@[ item ]);
+        }
+        [pollingConsumers removeAllObjects];
+        
+        // Notify websocket consumers
+        for (MBIMUpdateConsumer consumer in [websocketConsumers allValues]) {
+            consumer(@[ item ]);
+        }
+        
+        [updateItemHistory setObject:item forKey:@(item.messageSequenceNumber)];
+        
+        [self _cullUpdateItems];
+    });
+}
+
+- (WebSocket *)vendUpdateWebSocketConsumerForRequest:(HTTPMessage *)request socket:(GCDAsyncSocket *)gcdSocket
+{
+    WebSocket *socket = [[WebSocket alloc] initWithRequest:request socket:gcdSocket];
+    socket.delegate = self;
+    
+    MBIMUpdateConsumer consumer = ^(NSArray<MBIMUpdateItem *> *updates) {
+        NSMutableArray *encodedUpdates = [NSMutableArray array];
+        for (MBIMUpdateItem *item in updates) {
+            NSDictionary *updateDict = [item dictionaryRepresentation];
+            [encodedUpdates addObject:updateDict];
+        }
+        
+        NSData *data = [NSJSONSerialization dataWithJSONObject:encodedUpdates options:0 error:NULL];
+        [socket sendData:data];
+    };
+    
+    NSString *messageSeqString = [[request url] valueForQueryItemWithName:@"seq"];
+    [self _syncConsumer:consumer fromLastMessageSeq:(messageSeqString ? [messageSeqString integerValue] : -1)];
+    
+    __weak __auto_type websocketConsumers = _websocketConsumers;
+    dispatch_async(_accessQueue, ^{
+        NSString *websocketKey = [socket description];
+        [websocketConsumers setObject:consumer forKey:websocketKey];
+    });
+    
+    return socket;
+}
+
+- (BOOL)_syncConsumer:(MBIMUpdateConsumer)consumer fromLastMessageSeq:(NSInteger)messageSeq
+{
+    const BOOL needsSync = (messageSeq >= 0) && messageSeq < self->_messageSequenceNumber;
+    if (needsSync) {
+        __weak NSMutableDictionary *updateItemHistory = _updateItemHistory;
+        dispatch_async(_accessQueue, ^{
             NSMutableArray *batchedUpdates = [NSMutableArray array];
             for (NSUInteger seq = messageSeq + 1; seq <= self->_messageSequenceNumber; seq++) {
                 MBIMUpdateItem *item = [updateItemHistory objectForKey:@(seq)];
@@ -64,37 +146,10 @@ static const NSUInteger kUpdateItemsCullingLength = 100;
             }
             
             consumer(batchedUpdates);
-        } else {
-            [consumers addObject:consumer];
-        }
-    });
-}
-
-- (void)removeConsumer:(MBIMUpdateConsumer)consumer
-{
-    __weak NSMutableArray *consumers = _consumers;
-    dispatch_async(_accessQueue, ^{
-        [consumers removeObject:consumer];
-    });
-}
-
-- (void)enqueueUpdateItem:(MBIMUpdateItem *)item
-{
-    __weak NSMutableArray *consumers = _consumers;
-    __weak NSMutableDictionary *updateItemHistory = _updateItemHistory;
-    dispatch_async(_accessQueue, ^{
-        self->_messageSequenceNumber++;
-        item.messageSequenceNumber = self->_messageSequenceNumber;
-        
-        for (MBIMUpdateConsumer consumer in consumers) {
-            consumer(@[ item ]);
-        }
-        
-        [consumers removeAllObjects];
-        [updateItemHistory setObject:item forKey:@(item.messageSequenceNumber)];
-        
-        [self _cullUpdateItems];
-    });
+        });
+    }
+    
+    return needsSync;
 }
 
 - (void)_cullUpdateItems
@@ -113,6 +168,18 @@ static const NSUInteger kUpdateItemsCullingLength = 100;
     });
 }
 
+#pragma mark - <WebSocketDelegate>
+
+- (void)webSocketDidClose:(WebSocket *)ws
+{
+    // xxx: not great, but works.
+    NSString *websocketKey = [ws description];
+    __weak __auto_type websocketConsumers = _websocketConsumers;
+    dispatch_async(_accessQueue, ^{
+        [websocketConsumers removeObjectForKey:websocketKey];
+    });
+}
+
 @end
 
 @implementation MBIMUpdateItem

kordophone/Bridge/Operations/MBIMBridgeOperation.h

@@ -7,6 +7,7 @@
 //
 
 #import <Foundation/Foundation.h>
+#import <CocoaHTTPServer/HTTPMessage.h>
 #import <CocoaHTTPServer/HTTPResponse.h>
 #import <CocoaHTTPServer/HTTPErrorResponse.h>
 
@@ -20,6 +21,7 @@ typedef void (^MBIMBridgeOperationCompletionBlock)(NSObject<HTTPResponse> * _Nul
 @property (class, nonatomic, readonly) NSString *endpointName;
 @property (class, nonatomic, readonly) BOOL requiresAuthentication; // default YES
 
+@property (nonatomic, strong) HTTPMessage *request;
 @property (nonatomic, strong) NSData *requestBodyData;
 @property (nonatomic, readonly) NSURL *requestURL;
 @property (nonatomic, readonly) MBIMBridgeOperationCompletionBlock serverCompletionBlock;

kordophone/Bridge/Operations/MBIMBridgeOperation.m

@@ -7,6 +7,7 @@
 //
 
 #import "MBIMBridgeOperation.h"
+#import "MBIMURLUtilities.h"
 
 @interface MBIMBridgeOperation (/*INTERNAL*/)
 @property (nonatomic, strong) NSURL *requestURL;
@@ -79,18 +80,7 @@
 
 - (NSString *)valueForQueryItemWithName:(NSString *)queryItemName
 {
-    NSURLComponents *urlComponents = [NSURLComponents componentsWithURL:self.requestURL
-                                                resolvingAgainstBaseURL:NO];
-    
-    NSString *value = nil;
-    for (NSURLQueryItem *queryItem in [urlComponents queryItems]) {
-        if ([[queryItem name] isEqualToString:queryItemName]) {
-            value = [queryItem value];
-            break;
-        }
-    }
-    
-    return value;
+    return [[self requestURL] valueForQueryItemWithName:queryItemName];
 }
 
 @end

kordophone/Bridge/Operations/MBIMUpdatePollOperation.m

@@ -40,13 +40,13 @@
         weakSelf.serverCompletionBlock(response);
     };
     
-    [[MBIMUpdateQueue sharedInstance] addConsumer:_updateConsumer withLastSyncedMessageSeq:messageSeq];
+    [[MBIMUpdateQueue sharedInstance] addPollingConsumer:_updateConsumer withLastSyncedMessageSeq:messageSeq];
 }
 
 - (void)cancel
 {
     [super cancel];
-    [[MBIMUpdateQueue sharedInstance] removeConsumer:_updateConsumer];
+    [[MBIMUpdateQueue sharedInstance] removePollingConsumer:_updateConsumer];
 }
 
 - (NSObject<HTTPResponse> *)cancelAndReturnTimeoutResponse

kordophone/Bridge/Operations/Utilities/MBIMURLUtilities.h

@@ -0,0 +1,17 @@
+//
+//  MBIMURLUtilities.h
+//  kordophoned
+//
+//  Created by James Magahern on 1/17/23.
+//  Copyright © 2023 James Magahern. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface NSURL (MBIMURLUtilities)
+- (nullable NSString *)valueForQueryItemWithName:(NSString *)queryItemName;
+@end
+
+NS_ASSUME_NONNULL_END

kordophone/Bridge/Operations/Utilities/MBIMURLUtilities.m

@@ -0,0 +1,28 @@
+//
+//  MBIMURLUtilities.m
+//  kordophoned
+//
+//  Created by James Magahern on 1/17/23.
+//  Copyright © 2023 James Magahern. All rights reserved.
+//
+
+#import "MBIMURLUtilities.h"
+
+@implementation NSURL (MBIMURLUtilities)
+
+- (nullable NSString *)valueForQueryItemWithName:(NSString *)queryItemName
+{
+    NSURLComponents *urlComponents = [NSURLComponents componentsWithURL:self resolvingAgainstBaseURL:NO];
+    
+    NSString *value = nil;
+    for (NSURLQueryItem *queryItem in [urlComponents queryItems]) {
+        if ([[queryItem name] isEqualToString:queryItemName]) {
+            value = [queryItem value];
+            break;
+        }
+    }
+    
+    return value;
+}
+
+@end