Kordophone/kordophone/Bridge/MBIMPingPongWebSocket.m

//
//  MBIMPingPongWebSocket.m
//  kordophoned
//
//  Created by James Magahern on 6/13/25.
//  Copyright © 2025 James Magahern. All rights reserved.
//

#import "MBIMPingPongWebSocket.h"
#import "GCDAsyncSocket.h"

// WebSocket opcodes (from the base class)
#define TIMEOUT_NONE          -1
#define TIMEOUT_REQUEST_BODY  10

#define TAG_HTTP_REQUEST_BODY      100
#define TAG_HTTP_RESPONSE_HEADERS  200
#define TAG_HTTP_RESPONSE_BODY     201

#define TAG_PREFIX                 300
#define TAG_MSG_PLUS_SUFFIX        301
#define TAG_MSG_WITH_LENGTH        302
#define TAG_MSG_MASKING_KEY        303
#define TAG_PAYLOAD_PREFIX         304
#define TAG_PAYLOAD_LENGTH         305
#define TAG_PAYLOAD_LENGTH16       306
#define TAG_PAYLOAD_LENGTH64       307

#define WS_OP_CONTINUATION_FRAME   0
#define WS_OP_TEXT_FRAME           1
#define WS_OP_BINARY_FRAME         2
#define WS_OP_CONNECTION_CLOSE     8
#define WS_OP_PING                 9
#define WS_OP_PONG                 10

// Additional tags for ping/pong handling
#define TAG_PING_PAYLOAD           400
#define TAG_PONG_SENT             401

@interface WebSocket ()
- (void)socket:(GCDAsyncSocket *)asyncSocket didReadData:(NSData *)data withTag:(long)tag;
@end

@implementation MBIMPingPongWebSocket
{
    NSData *pendingPingPayload;
}

#pragma mark - Ping/Pong Frame Construction

- (NSData *)createPongFrameWithPayload:(NSData *)pingPayload {
    NSMutableData *frame = [NSMutableData data];
    
    // First byte: FIN (1) + RSV (000) + Opcode (1010 = 0xA for Pong)
    uint8_t firstByte = 0x8A; // 10001010 in binary
    [frame appendBytes:&firstByte length:1];
    
    // Second byte: MASK (0) + Payload Length
    NSUInteger payloadLength = [pingPayload length];
    
    if (payloadLength <= 125) {
        // Short payload length (0-125 bytes)
        uint8_t secondByte = (uint8_t)payloadLength; // MASK bit is 0 for server-to-client
        [frame appendBytes:&secondByte length:1];
    }
    else if (payloadLength <= 65535) {
        // Medium payload length (126-65535 bytes)
        uint8_t secondByte = 126;
        [frame appendBytes:&secondByte length:1];
        
        // Add 16-bit length in network byte order (big-endian)
        uint16_t extendedLength = htons((uint16_t)payloadLength);
        [frame appendBytes:&extendedLength length:2];
    }
    else {
        // Large payload length (65536+ bytes) - rarely needed for pings
        uint8_t secondByte = 127;
        [frame appendBytes:&secondByte length:1];
        
        // Add 64-bit length in network byte order (big-endian)
        uint64_t extendedLength = OSSwapHostToBigInt64(payloadLength);
        [frame appendBytes:&extendedLength length:8];
    }
    
    // Append the payload (copied exactly from ping)
    if (payloadLength > 0) {
        [frame appendData:pingPayload];
    }
    
    return [frame copy];
}

- (void)sendPongWithPayload:(NSData *)payload socket:(GCDAsyncSocket *)asyncSocket {
    NSData *pongFrame = [self createPongFrameWithPayload:payload];
    
    // Send the pong frame directly through the socket
    [asyncSocket writeData:pongFrame withTimeout:TIMEOUT_NONE tag:TAG_PONG_SENT];
    
    NSLog(@"Sent pong frame with payload length: %lu", (unsigned long)[payload length]);
}

#pragma mark - Override AsyncSocket Delegate

- (void)socket:(GCDAsyncSocket *)sock didReadData:(NSData *)data withTag:(long)tag {
    BOOL nextFrameMasked = [[self valueForKey:@"nextFrameMasked"]  boolValue];
    NSUInteger nextOpCode = [[self valueForKey:@"nextOpCode"] unsignedIntValue]; // sheesh.
    
    NSLog(@"tag: %ld, nextOpCode: %lu", tag, (unsigned long)nextOpCode);
    
    // Handle our custom ping payload tag
    if (tag == TAG_PING_PAYLOAD) {
        // We've read the ping payload, now send the pong response
        NSData *payload = data;
        
        // If the frame was masked, we need to unmask it
        if (nextFrameMasked && pendingPingPayload) {
            NSMutableData *unmaskedPayload = [data mutableCopy];
            UInt8 *payloadBytes = [unmaskedPayload mutableBytes];
            UInt8 *maskBytes = (UInt8 *)[pendingPingPayload bytes];
            
            for (NSUInteger i = 0; i < [data length]; i++) {
                payloadBytes[i] ^= maskBytes[i % 4];
            }
            payload = [unmaskedPayload copy];
        }
        
        [self sendPongWithPayload:payload socket:sock];
        
        // Continue reading the next frame
        [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
        return;
    }
    
    // Let the parent handle the frame parsing logic first
    [super socket:sock didReadData:data withTag:tag];
    
    // After the parent has processed, check if we need to handle ping/pong
    if (tag == TAG_MSG_WITH_LENGTH) {
        // The parent just finished reading a complete message
        if (nextOpCode == WS_OP_PING) {
            // This was a ping frame - the parent would have tried to process it as text
            // We need to send a pong response
            NSData *payload = data;
            
            // If the frame was masked, unmask it (parent already did this)
            [self sendPongWithPayload:payload socket:sock];
            
            // Don't call didReceiveMessage for ping frames
            return;
        }
        else if (nextOpCode == WS_OP_PONG) {
            // This was a pong frame - just log it and continue
            NSLog(@"Received pong frame with payload length: %lu", (unsigned long)[data length]);
            return;
        }
    }
}


#pragma mark - Helper Methods

// Expose the isValidWebSocketFrame method since it's private in the parent
- (BOOL)isValidWebSocketFrame:(UInt8)frame {
    NSUInteger rsv = frame & 0x70;
    NSUInteger opcode = frame & 0x0F;
    if (rsv || (3 <= opcode && opcode <= 7) || (0xB <= opcode && opcode <= 0xF)) {
        return NO;
    }
    return YES;
}

// Override didReceiveMessage to add logging
- (void)didReceiveMessage:(NSString *)msg {
    NSLog(@"Received message: %@", msg);
    [super didReceiveMessage:msg];
}

// Override didOpen to add logging
- (void)didOpen {
    NSLog(@"WebSocket opened with ping/pong support");
    [super didOpen];
}

- (void)didClose {
    NSLog(@"WebSocket closed");
    [super didClose];
}

@end