WJLSocket: Objective-C wrapper for BSD Sockets on iOS / Mac

NSData+DebugOutput.h

#import <Foundation/Foundation.h>

@interface NSData (DebugOutput)

- (NSString *)debugOutput;

@end

NSData+DebugOutput.m

#import "NSData+DebugOutput.h"

@implementation NSData (DebugOutput)

- (NSString *)debugOutput
{
    int rowlength = 16;
    const uint8_t *data = [self bytes];
    NSMutableString *str = [NSMutableString new];
    for (NSInteger i = 0; i < (NSInteger) [self length]; ) {
        [str appendFormat:@"\n"];
        for (NSInteger j = 0; j < rowlength; j++, i++) {
            if (i >= [self length] || i < 0) {
                [str appendString:@"  "];
            } else {
                [str appendFormat:@"%02x", data[i]];
            }
            if (i % 2) [str appendString:@" "];
        }
        i -= rowlength;
        [str appendString:@" "];
        for (NSInteger j = 0; j < rowlength; j++, i++) {
            if (i >= [self length] || i < 0) {
                [str appendString:@" "];
            } else if (data[i] < 33 || data[i] > '~') {
                [str appendString:@"."];
            } else {
                [str appendFormat:@"%c", data[i]];
            }
        }
    }
    return str;
}

@end

WJLSocket.h

#import <Foundation/Foundation.h>

/**
 * A simple Objective-C wrapper around BSD Sockets for OS X and iOS.
 */
@interface WJLSocket : NSObject

@property (assign) BOOL debug;

+ (instancetype)socketWithHostname:(NSString *)hostname port:(NSUInteger)port;

/**
 * Connect to the hostname and port specified in the constructor. The handler is
 * not called until the entire connection process is complete. Once the handler
 is called (indicating success), the socket is ready to be used.
 */
- (void)connect:(void(^)(BOOL success))handler;

/**
 * Write <i>data</i> to the socket.
 *
 * The handler being called with <i>success == YES</i> indicates that the entire
 * buffer was written successfully. Otherwise, <i>success</i> will be <i>NO</i>
 * and the socket will be closed.
 */
- (void)write:(NSData *)data withHandler:(void(^)(BOOL success))handler;

/**
 * Attempt to read up to <i>length</i> bytes from the socket.
 *
 * If the read is successful, the data that was read will be passed to the
 * handler. If the read times out, the handler will be called with a non-nil,
 * zero-length <i>data</i>. A nil <i>data</i> indicates a socket error, and the
 * socket will be closed.
 */
- (void)readWithLength:(NSUInteger)length withHandler:(void(^)(NSData *data))handler;


/**
 * Attempt to peek at up to <i>length</i> bytes from the socket. This data is
 * not removed and will be included in the next read call.
 *
 * If the peek is successful, the data that was peeked will be passed to the
 * handler. If the read times out, the handler will be called with a non-nil,
 * zero-length <i>data</i>. A nil <i>data</i> indicates a socket error, and the
 * socket will be closed.
 */
- (void)peekWithLength:(NSUInteger)length withHandler:(void(^)(NSData *data))handler;

@end

WJLSocket.m

#import "WJLSocket.h"
#import "NSData+DebugOutput.h"
#import <netdb.h>

@interface WJLSocket ()

@property (strong) dispatch_queue_t socketQueue;
@property (copy)   NSString *hostname;
@property (assign) NSUInteger port;
@property (assign) int socketDescriptor;
@property (assign) int state;

@end

@implementation WJLSocket

+ (instancetype)socketWithHostname:(NSString *)hostname port:(NSUInteger)port
{
    return [[self alloc] initWithHostname:hostname port:port];
}

- (instancetype)initWithHostname:(NSString *)hostname port:(NSUInteger)port
{
    self = [super init];
    if (self) {
        _socketQueue = dispatch_queue_create("net.wjlafrance.socketqueue", DISPATCH_QUEUE_CONCURRENT);
        _hostname = hostname;
        _port = port;
    }
    return self;
}

- (void)dealloc
{
    [self close];
}

- (void)close
{
    close(self.socketDescriptor);
    self.socketDescriptor = 0;
    self.state = 0;
}

- (void)connect:(void(^)(BOOL success))handler
{
    NSAssert(NULL != handler, @"Handler must not be NULL");

    NSAssert(0 == self.state, @"-connect may only be called once.");
    self.state = 1;

    self.socketDescriptor = socket(AF_INET, SOCK_STREAM, 0);
    NSAssert(0 != self.socketDescriptor, @"Socket allocation failed!");

    dispatch_async(_socketQueue, ^{
        struct hostent *host = gethostbyname([self.hostname UTF8String]);
        struct sockaddr_in sockaddr;
        memset(&sockaddr, 0, sizeof(sockaddr));
        sockaddr.sin_len = sizeof(sockaddr);
        sockaddr.sin_family = AF_INET;
        sockaddr.sin_port = htons(self.port);
        sockaddr.sin_addr = *(struct in_addr *) host->h_addr_list[0];

        BOOL success = (-1 != connect(self.socketDescriptor, (struct sockaddr *) &sockaddr, sizeof(sockaddr)));
        if (success) {
            self.state = 2;
        } else {
            [self close];
        }
        dispatch_async(dispatch_get_main_queue(), ^{
            handler(success);
        });
    });
}

- (void)write:(NSData *)data withHandler:(void(^)(BOOL success))handler
{
    NSAssert(nil != data, @"data must not be nil");

    NSAssert(2 == self.state, @"Socket must be connected before calling -write");

    dispatch_async(_socketQueue, ^{
        if (self.debug) {
            NSLog(@"C -> S %@:%lu (%ld bytes) %@", self.hostname, (unsigned long) self.port, (unsigned long) [data length], [data debugOutput]);
        }

        ssize_t bytesWritten = send(self.socketDescriptor, [data bytes], [data length], 0);

        BOOL success = (bytesWritten == (ssize_t) [data length]);
        if (!success) {
            [self close];
        }

        if (NULL != handler) {
            dispatch_async(dispatch_get_main_queue(), ^{
                handler(success);
            });
        }
    });
}

- (void)readWithLength:(NSUInteger)length flags:(int)flags withHandler:(void(^)(NSData *data))handler
{
    NSAssert(NULL != handler, @"Handler must not be NULL");

    NSAssert(2 == self.state, @"Socket must be connected before calling -receiveWithMaxLength:withHandler:");

    dispatch_async(_socketQueue, ^{
        void *buffer = malloc(length);
        NSData *data;

        ssize_t bytesRead = recv(self.socketDescriptor, buffer, length, flags);

        BOOL success = (bytesRead == (ssize_t) length);
        if (!success) {
            [self close];
        } else {
            data = [NSData dataWithBytes:buffer length:bytesRead];
        }

        free(buffer);

        if (success && self.debug && !(flags & MSG_PEEK)) {
            NSLog(@"S -> C %@:%lu (%ld bytes) %@", self.hostname, (unsigned long) self.port, (unsigned long) [data length], [data debugOutput]);
        }

        dispatch_async(dispatch_get_main_queue(), ^{
            handler(data);
        });
    });
}

- (void)readWithLength:(NSUInteger)length withHandler:(void(^)(NSData *data))handler
{
    [self readWithLength:length flags:MSG_WAITALL withHandler:handler];
}

- (void)peekWithLength:(NSUInteger)length withHandler:(void(^)(NSData *data))handler
{
    [self readWithLength:length flags:MSG_PEEK | MSG_WAITALL withHandler:handler];
}

@end