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
{
    NSUInteger currentPayloadLength;
}

#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 {
    NSData *pongFrame = [self createPongFrameWithPayload:payload];
    
    // Send the pong frame directly through the socket
    [asyncSocket writeData:pongFrame withTimeout:TIMEOUT_NONE tag:TAG_PONG_SENT];
}

#pragma mark - Override AsyncSocket Delegate

- (void)socket:(GCDAsyncSocket *)asyncSocket didReadData:(NSData *)data withTag:(long)tag {
    // xxx: sheesh: really need to get off of CococaHTTPServer.
    NSData *maskingKey = [self valueForKey:@"maskingKey"];
    BOOL nextFrameMasked = [[self valueForKey:@"nextFrameMasked"]  boolValue];
    NSUInteger nextOpCode = [[self valueForKey:@"nextOpCode"] unsignedIntValue];
    
    // 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 && maskingKey) {
            NSMutableData *unmaskedPayload = [data mutableCopy];
            UInt8 *payloadBytes = [unmaskedPayload mutableBytes];
            UInt8 *maskBytes = (UInt8 *)[maskingKey bytes];
            
            for (NSUInteger i = 0; i < [data length]; i++) {
                payloadBytes[i] ^= maskBytes[i % 4];
            }
            payload = [unmaskedPayload copy];
        }
        
        [self sendPongWithPayload:payload];
        
        // Continue reading the next frame
        [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
        return;
    }
    
    // Let parent handle TAG_PAYLOAD_PREFIX first to set nextOpCode
    if (tag == TAG_PAYLOAD_PREFIX) {
        [super socket:asyncSocket didReadData:data withTag:tag];
        return;
    }
    
    // Now intercept ping/pong handling after nextOpCode is set
    if (nextOpCode == WS_OP_PING || nextOpCode == WS_OP_PONG) {
        if (tag == TAG_PAYLOAD_LENGTH) {
            UInt8 frame = *(UInt8 *)[data bytes];
            BOOL masked = (frame & 0x80) != 0;
            NSUInteger length = frame & 0x7F;
            nextFrameMasked = masked;
            currentPayloadLength = length;
            
            if (length <= 125) {
                if (nextFrameMasked) {
                    [asyncSocket readDataToLength:4 withTimeout:TIMEOUT_NONE tag:TAG_MSG_MASKING_KEY];
                } else if (length > 0) {
                    [asyncSocket readDataToLength:length withTimeout:TIMEOUT_NONE tag:TAG_PING_PAYLOAD];
                } else {
                    // Empty payload, no masking - handle immediately
                    [self sendPongWithPayload:[NSData data]];
                    [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
                }
            }
            else if (length == 126) {
                [asyncSocket readDataToLength:2 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH16];
            }
            else {
                [asyncSocket readDataToLength:8 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH64];
            }
            return;
        }
        
        if (tag == TAG_MSG_MASKING_KEY) {
            // Store the masking key
            maskingKey = [data copy];
            
            // Now read the payload (or handle empty payload)
            if (currentPayloadLength > 0) {
                [asyncSocket readDataToLength:currentPayloadLength withTimeout:TIMEOUT_NONE tag:TAG_PING_PAYLOAD];
            } else {
                // Empty payload - send pong immediately
                [self sendPongWithPayload:[NSData data]];
                [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
            }
            return;
        }
        
        if (tag == TAG_PAYLOAD_LENGTH16) {
            UInt8 *pFrame = (UInt8 *)[data bytes];
            NSUInteger length = ((NSUInteger)pFrame[0] << 8) | (NSUInteger)pFrame[1];
            
            if (nextFrameMasked) {
                [asyncSocket readDataToLength:4 withTimeout:TIMEOUT_NONE tag:TAG_MSG_MASKING_KEY];
            }
            [asyncSocket readDataToLength:length withTimeout:TIMEOUT_NONE tag:TAG_PING_PAYLOAD];
            return;
        }
        
        if (tag == TAG_PAYLOAD_LENGTH64) {
            // For 64-bit length, this is too complex for ping frames - just close
            [self didClose];
            return;
        }
    }
    
    // For all other cases, call the parent implementation
    [super socket:asyncSocket didReadData:data withTag:tag];
}

#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;
}

@end