IPhoneUSB.m

/*!
 @source IPhoneUSB.m
 @project Pusher
 @author Alexander Maksimenko
 @copyright Copyright (c) 2009 Ripdev. All rights reserved.
 */

#import <unistd.h>

#import "AppleCRC32.h"
#import "IPhoneUSB.h"

#pragma mark C Callbacks

static void usbDeviceAdded(void *refCon, io_iterator_t iter)
{
	[(IPhoneUSB*)refCon usbDeviceAdded:iter];
}

static void usbDeviceRemoved(void *refCon, io_iterator_t iter)
{
	[(IPhoneUSB*)refCon usbDeviceRemoved:iter];
}

#pragma mark -

#define kProductIPhoneNormal 0x1290
#define kProductIPhone3GNormal 0x1292
#define kProductIPodNormal 0x1291
#define kProductIPhoneRecovery 0x1280
#define kProductIPhoneRecovery2 0x1281
#define kProductIPhoneDFU 0x1222
#define kProductIPhoneDFU2 0x1227

@implementation IPhoneUSB

-(id) init
{
	self = [super init];

	_ioKitNotificationPort = 0;
	_notificationRunLoopSource = 0;
	_enabled = NO;
    _newInterface = NO;
	
    _lock = [[NSLock alloc] init];
    
	return self;
}

-(void) dealloc
{
	[self ioKitTearDown];
    [_lock release];
    
	[super dealloc];
}

-(id) delegate
{
    return _delegate;
}

-(void) setDelegate:(id) delegate
{
    _delegate = delegate;
    
    if(_delegate != nil)
        [self restartService];
}

-(void) ioKitSetUp
{
	_ioKitNotificationPort = IONotificationPortCreate(kIOMasterPortDefault);
	_notificationRunLoopSource = IONotificationPortGetRunLoopSource(_ioKitNotificationPort);

	CFRunLoopAddSource(CFRunLoopGetCurrent(), _notificationRunLoopSource, kCFRunLoopDefaultMode);
}

-(void) ioKitTearDown
{
	if(_ioKitNotificationPort)
	{
        Log(@"%@ ioKitTearDown", self);

		CFRunLoopRemoveSource(CFRunLoopGetCurrent(), _notificationRunLoopSource, kCFRunLoopDefaultMode);
		IONotificationPortDestroy(_ioKitNotificationPort);
        _ioKitNotificationPort = 0;
	}
}

-(void) restartService
{
    [self ioKitTearDown];
    [self startService];
}

-(void) startService
{
    Log(@"%@ startService", self);

	_enabled = YES;

	if(_ioKitNotificationPort == 0)
	{
		[self ioKitSetUp];
		[self registerForUSBNotifications];
	}
}

-(void) stopService
{
    Log(@"%@ stopService", self);

	_enabled = NO;
}

-(BOOL) serviceStatus
{
	return _enabled;
}

-(void) registerForUSBNotifications
{
    Log(@"%@ registerForUSBNotifications", self);
    
	io_iterator_t addedIterator;
	io_iterator_t removedIterator;

	CFDictionaryRef matchingDictionary = IOServiceMatching(kIOUSBDeviceClassName);
	kern_return_t res = IOServiceAddMatchingNotification(_ioKitNotificationPort,
											kIOTerminatedNotification,
											matchingDictionary,
											usbDeviceRemoved,
											self,
											&removedIterator );

	if (res)
		[_delegate processIPhoneError:@"IOServiceAddMatchingNotification() failed.\n" :self];
	else
		[self usbDeviceRemoved:removedIterator];

	matchingDictionary = IOServiceMatching(kIOUSBDeviceClassName);
	res = IOServiceAddMatchingNotification(_ioKitNotificationPort,
															kIOMatchedNotification,
															matchingDictionary,
															usbDeviceAdded,
															self,
															&addedIterator );

	if (res)
		[_delegate processIPhoneError:@"IOServiceAddMatchingNotification() failed.\n" :self];
	else
		[self usbDeviceAdded:addedIterator];
}

-(void) usbDeviceAdded:(io_iterator_t) iterator
{
	if(!_enabled)
		return;
		
	io_object_t	device;
	while ((device = IOIteratorNext(iterator)))
	{
		[self ifIPhoneAdded:device];
		IOObjectRelease(device);
	}
}

-(void) usbDeviceRemoved:(io_iterator_t) iterator
{
	if(!_enabled)
		return;

	io_object_t device;
	while ((device = IOIteratorNext(iterator)))
	{
		[self ifIPhoneRemoved:device];
		IOObjectRelease(device);
	}
}

-(void) ifIPhoneAdded:(io_object_t) device
{
    [_lock lock];

	IOCFPlugInInterface **iodev;
	SInt32 score;

	if(IOCreatePlugInInterfaceForService(device, kIOUSBDeviceUserClientTypeID, kIOCFPlugInInterfaceID, &iodev, &score) == 0)
	{
		if((*iodev)->QueryInterface(iodev, CFUUIDGetUUIDBytes(kIOUSBDeviceInterfaceID),	(LPVOID*)&_dev) == 0)
		{
			UInt16 vendorId = 0;
			UInt16 productId = 0;

			(*_dev)->GetDeviceVendor(_dev, &vendorId);
			(*_dev)->GetDeviceProduct(_dev, &productId);

			if(vendorId == kIOUSBVendorIDAppleComputer)
			{
				switch(productId)
				{
					case kProductIPodNormal:
                    
						_mode = IPodNormalMode;
						[_delegate processIPodNormalConnected:self];
						[self dealWithDevice];

                        break;
					case kProductIPhoneNormal:

						_mode = IPhoneNormalMode;
						[_delegate processIPhoneNormalConnected:self];
						[self dealWithDevice];

						break;
					case kProductIPhone3GNormal:

						_mode = IPhone3GNormalMode;
						[_delegate processIPhone3GNormalConnected:self];
						[self dealWithDevice];

						break;
					case kProductIPhoneRecovery:

                        _newInterface = NO;
						_mode = RecoveryMode;
						[_delegate processIPhoneRecoveryConnected:self];
						[self dealWithDevice];
						
						break;
					case kProductIPhoneRecovery2:
                    
                        _newInterface = YES;
						_mode = Recovery2Mode;
						[_delegate processIPhoneRecovery2Connected:self];
						[self dealWithDevice];
						
						break;
					case kProductIPhoneDFU:
						
						_mode = DFUMode;
						[_delegate processIPhoneDFUConnected:self];
						[self dealWithDevice];

						break;
					case kProductIPhoneDFU2:

						_mode = DFU2Mode;
						[_delegate processIPhoneDFU2Connected:self];
						[self dealWithDevice];

						break;
					default:
						break;
				}
			}
			(*_dev)->Release(_dev);
		}

		(*iodev)->Release(iodev);
	}
    
    [_lock unlock];
}

-(void) ifIPhoneRemoved:(io_object_t) device
{
	io_name_t deviceNameChars;
	if(IORegistryEntryGetName(device, deviceNameChars) == 0)
	{
		NSString *deviceName = [NSString stringWithCString: deviceNameChars];
	
		if([deviceName isEqualToString:@"iPhone"] || [deviceName isEqualToString:@"iPod"] || [deviceName isEqualToString:@"USB DFU Device"] || [deviceName isEqualToString:@"Apple Mobile Device (Recovery Mode)"] || [deviceName isEqualToString:@"Apple Mobile Device (DFU Mode)"])
		{
			if([deviceName isEqualToString:@"iPod"])
				[_delegate processIPodNormalDisconnected:self];
			else if([deviceName isEqualToString:@"iPhone"])
				[_delegate processIPhoneNormalDisconnected:self];
			else if([deviceName isEqualToString:@"Apple Mobile Device (Recovery Mode)"])
				[_delegate processIPhoneRecoveryDisconnected:self];
            else if([deviceName isEqualToString:@"Apple Mobile Device (DFU Mode)"])
				[_delegate processIPhoneDFU2Disconnected:self];
			else
				[_delegate processIPhoneDFUDisconnected:self];
		}
	}
}

-(void) dealWithDevice
{
    if (_usbNotificationsOnly)
        return;

	UInt8 configIndex;
	
	switch(_mode)
	{
		case IPodNormalMode:
		case IPhoneNormalMode:
		case IPhone3GNormalMode:
			configIndex = 2;
			break;
		case RecoveryMode:
		case Recovery2Mode:
			configIndex = 0;
			break;
		case DFUMode:
		case DFU2Mode:
			configIndex = 0;
			break;
		default:
			return;
	}

    Log(@"%@ Deal with device in %d mode", self, _mode);
    
    if((*_dev)->USBDeviceOpen(_dev) == 0)
    {
	    IOUSBConfigurationDescriptorPtr	confDesc;
		
		if((*_dev)->GetConfigurationDescriptorPtr(_dev, configIndex, &confDesc) == 0)
		{
			if((*_dev)->SetConfiguration(_dev, confDesc->bConfigurationValue) == 0)
			{
				IOUSBFindInterfaceRequest interfaceRequest;

				interfaceRequest.bInterfaceClass = ((_mode == DFUMode || _mode == DFU2Mode) ? 254 : 255);
				interfaceRequest.bInterfaceSubClass = kIOUSBFindInterfaceDontCare;
				interfaceRequest.bInterfaceProtocol = kIOUSBFindInterfaceDontCare;
				interfaceRequest.bAlternateSetting = kIOUSBFindInterfaceDontCare;
    
				io_iterator_t iterator;
				if((*_dev)->CreateInterfaceIterator(_dev, &interfaceRequest, &iterator) == 0)
				{
					io_service_t usbInterfaceRef;
					if((usbInterfaceRef = IOIteratorNext(iterator)))
					{
						[self dealWithInterface:usbInterfaceRef];
						IOObjectRelease(usbInterfaceRef);
					}
					else
						[_delegate processIPhoneError:@"IOIteratorNext() failed.\n" :self];

					IOObjectRelease(iterator);
				}
				else
					[_delegate processIPhoneError:@"CreateInterfaceIterator() failed.\n" :self];
			}
			else
				[_delegate processIPhoneError:@"SetConfiguration() failed.\n" :self];
		}
		else
			[_delegate processIPhoneError:@"GetConfigurationDescriptorPtr() failed.\n" :self];

		(*_dev)->USBDeviceClose(_dev);
	}
	else
		[_delegate processIPhoneError:@"USBDeviceOpen() failed.\n" :self];

    Log(@"%@ Finish deal with device in %d mode", self, _mode);
}

-(void) dealWithInterface:(io_service_t) usbInterfaceRef
{
    if (_usbNotificationsOnly)
        return;

    Log(@"%@ Deal with interface in %d mode, _dev = 0x%08X", self, _mode, _dev);

    IOCFPlugInInterface **iodev;
    SInt32 score;
	
	if(IOCreatePlugInInterfaceForService(usbInterfaceRef, kIOUSBInterfaceUserClientTypeID, kIOCFPlugInInterfaceID, &iodev, &score) == 0)
	{
		if((*iodev)->QueryInterface(iodev, CFUUIDGetUUIDBytes(kIOUSBInterfaceInterfaceID), (LPVOID)&_intf) == 0)
		{
			(*iodev)->Release(iodev);
			
			if((*_intf)->USBInterfaceOpen(_intf) == 0)
			{
			
				_ctrlIn = 0;
				_ctrlOut = 0;
				_serialIn = 0;
				_serialOut = 0;
				_fileIn = 0;
				_fileOut = 0;
                _transaction = 0;

				switch(_mode)
				{
					case IPodNormalMode:

						_ctrlOut = 1; // 0x04
						_ctrlIn = 2; // 0x85

						[_delegate processIPodNormal:self];

						break;
					case IPhoneNormalMode:

						_ctrlOut = 1; // 0x04
						_ctrlIn = 2; // 0x85

						[_delegate processIPhoneNormal:self];

						break;
					case IPhone3GNormalMode:

						_ctrlOut = 1; // 0x04
						_ctrlIn = 2; // 0x85

						[_delegate processIPhone3GNormal:self];

						break;
					case RecoveryMode:

                        _serialIn = 1; // 0x81
                        _serialOut = 2; // 0x02
                        _ctrlIn = 3; // 0x83
                        
                        _ctrlOut = 4; // 0x04
                        _fileIn = 5; // 0x85
                        _fileOut = 6; // 0x05
                    
                        [_delegate processIPhoneRecovery:self];
                        
                        break;
					case Recovery2Mode:
                    
                        if((*_intf)->SetAlternateInterface(_intf, 1) == 0)
                        {
                            _serialIn = 1; // 0x81
                            _serialOut = 2; // 0x02
                            
                            [_delegate processIPhoneRecovery2:self];
                        }

						break;
					case DFUMode:

						[_delegate processIPhoneDFU:self];
						
						break;
					case DFU2Mode:
						
						[_delegate processIPhoneDFU2:self];
						
						break;
					default:
						break;
				}

				(*_intf)->USBInterfaceClose(_intf);
			}
			else
				[_delegate processIPhoneError:@"USBInterfaceOpen() failed.\n" :self];

			(*_intf)->Release(_intf);
		}
		else
			[_delegate processIPhoneError:@"QueryInterface() failed.\n" :self];
	}
	else
		[_delegate processIPhoneError:@"IOCreatePlugInInterfaceForService() failed.\n" :self];

    Log(@"%@ Finish deal with interface in %d mode, _dev = 0x%08X", self, _mode, _dev);
}

- (void)enableUSBNotificationOnlyMode
{
    _usbNotificationsOnly = YES;
}

@end

@implementation IPhoneUSB (IPhoneIO)

-(IOReturn) resetDevice
{
    if(_dev == 0)
    {
        [_delegate processIPhoneError:@"resetDevice() failed, _dev = NULL" :self];
        return kIOReturnError;
    }
    
    IOReturn res = (*_dev)->ResetDevice(_dev);

    if(res != kIOReturnSuccess)
        [_delegate processIPhoneError:[NSString stringWithFormat:@"resetDevice() failed with error %X", res] :self];
    
    return res;
}

-(IOReturn) writeSerial:(void*) data :(unsigned int) size
{
    if(_intf == 0)
    {
        [_delegate processIPhoneError:@"writeSerial() failed, _intf = NULL" :self];
        return kIOReturnError;
    }
    
    IOReturn res = (*_intf)->WritePipe(_intf, _serialOut, data, size);

    if(res != kIOReturnSuccess)
        [_delegate processIPhoneError:[NSString stringWithFormat:@"writeSerial() failed with error %X", res] :self];
    
    return res;
}

-(IOReturn) readSerial:(void*) data :(unsigned int*) pSize
{
    if(_intf == 0)
    {
        [_delegate processIPhoneError:@"readSerial() failed, _intf = NULL" :self];
        return kIOReturnError;
    }
    
	IOReturn res = (*_intf)->ReadPipe(_intf, _serialIn, data, (UInt32*)pSize);

    if(res != kIOReturnSuccess)
        [_delegate processIPhoneError:[NSString stringWithFormat:@"readSerial() failed with error %X", res] :self];
        
    return res;
}

-(IOReturn) sendRequest:(IOUSBDevRequest*) request
{
    if(_dev == 0)
    {
        [_delegate processIPhoneError:@"DeviceRequest() failed, _dev = NULL" :self];
        return kIOReturnError;
    }

//    Log(@"Send request:\nbmRequestType = %X\nbRequest = %X\nwValue = %X\nwIndex = %X\nwLength = %X",request->bmRequestType,request->bRequest,request->wValue,request->wIndex,request->wLength);

    IOReturn res = (*_dev)->DeviceRequest(_dev, request);
	
    if(res != kIOReturnSuccess)
        [_delegate processIPhoneError:[NSString stringWithFormat:@"DeviceRequest() failed with error %X", res] :self];
	else if(request->wLength != request->wLenDone)
        [_delegate processIPhoneError:[NSString stringWithFormat:@"DeviceRequest() failed, wLength != wLenDone (%d != %d)", request->wLength, request->wLenDone] :self];

	return res;
}

-(IOReturn) sendCommand:(NSString*) cmd
{
    cmd = [cmd stringByAppendingString:@"\n"];
    
    return [self writeSerial:(unsigned char*)[cmd UTF8String] :[cmd length]];
}

-(void) uploadFile:(NSString*) file
{
    dfu_status status;

    status = [self dfuStatus];
    if(status.bStatus != DFU_STATUS_OK)
        return;

    NSMutableData* data = [NSMutableData dataWithContentsOfFile:file];
    
    if(data != nil)
    {
        if(_mode == DFUMode || _mode == DFU2Mode)
        {
            unsigned int crc = 0xFFFFFFFF;
            unsigned char hash[] = {0xff, 0xff, 0xff, 0xff, 0xac, 0x05, 0x00, 0x01, 0x55, 0x46, 0x44, 0x10};
            [data appendBytes:hash length:sizeof(hash)];
            AppleCRC32Checksum(&crc, [data mutableBytes], [data length]);
            [data appendBytes:&crc length:sizeof(crc)];
        }

        void* p = [data mutableBytes];
        unsigned int size = [data length];
        unsigned int sent = 0;
        unsigned short len = 0x800;
        
        [_delegate processBeginUpload:size];
        [_delegate processIPhoneMsg:[NSString stringWithFormat:@">> Uploading file %@ to device", file] :self];

        while(sent < size)
        {
            if(size - sent < 0x800)
                len = size - sent;

            status = [self uploadFileChunk:p + sent :len];
            
            sent += len;
            
            if(status.bStatus != DFU_STATUS_OK)
            {
                [_delegate processIPhoneError:@"Couldn't send file chunk to device" :self];
                return;
            }
        }
        
        status = [self uploadFileChunk:0 :0];
        
        if(status.bState == DFU_STATE_DFU_MANIFEST_SYNC)
            status = [self dfuStatus];

        if(status.bState == DFU_STATE_DFU_MANIFEST)
        {
            if(_mode == DFU2Mode)
            {
                status = [self dfuStatus];

                if(status.bState == DFU_STATE_DFU_MANIFEST_WAIT_RESET)
                    [self resetDevice];
            }
                
            [_delegate processIPhoneMsg:@">> Upload success" :self];
        }
        else
            [_delegate processIPhoneError:@"Upload failed" :self];
    }
    else
        [_delegate processIPhoneError:[NSString stringWithFormat:@"Couldn't open file %@", file] :self];
}

-(dfu_status) uploadFileChunk:(void*) buf :(unsigned short) length
{
    dfu_status status;
    IOUSBDevRequest request;

    status.bStatus = DFU_STATUS_ERROR_UNKNOWN;
    status.bwPollTimeout = 0;
    status.bState = DFU_STATE_DFU_ERROR;
    status.iString = 0;

	request.bmRequestType = USBmakebmRequestType(kUSBOut, kUSBClass, kUSBInterface);
    request.bRequest = DFU_DNLOAD;
    request.wValue = _transaction++;
    request.wIndex = 0;
    request.wLength = length;
    request.pData = buf;
	
	IOReturn ret = [self sendRequest:&request];
    
    if(ret == kIOReturnSuccess)
    {
        do
        {
            status = [self dfuStatus];
            usleep(5000);
        }
        while(status.bState == DFU_STATE_DFU_DOWNLOAD_BUSY);
    }
    
    return status;
}

-(IOReturn) dfuAbort
{
    IOUSBDevRequest request;

	request.bmRequestType = USBmakebmRequestType(kUSBOut, kUSBClass, kUSBInterface);
    request.bRequest = DFU_ABORT;
    request.wValue = 0;
    request.wIndex = 0;
    request.wLength = 0;
    request.pData = 0;
	
	return [self sendRequest:&request];
}

-(dfu_status) dfuStatus
{
    char buffer[6] = {0};
    
    dfu_status status;
    IOUSBDevRequest request;

    status.bStatus = DFU_STATUS_ERROR_UNKNOWN;
    status.bwPollTimeout = 0;
    status.bState = DFU_STATE_DFU_ERROR;
    status.iString = 0;

	request.bmRequestType = USBmakebmRequestType(kUSBIn, kUSBClass, kUSBInterface);
    request.bRequest = DFU_GETSTATUS;
    request.wValue = 0;
    request.wIndex = 0;
    request.wLength = 6;
    request.pData = buffer;
    
    if([self sendRequest:&request] == kIOReturnSuccess)
    {
        status.bStatus = buffer[0];
        status.bwPollTimeout = ((0xff & buffer[3]) << 16) |
                                ((0xff & buffer[2]) << 8)  |
                                (0xff & buffer[1]);

        status.bState  = buffer[4];
        status.iString = buffer[5];
    }

    if(status.bStatus != DFU_STATUS_OK)
    {
        [_delegate processIPhoneError:[self stringForDFUStatus:status.bStatus] :self];
        [self dfuClearStatus];
    }
    
    return status;
}

-(IOReturn) dfuClearStatus
{
    IOUSBDevRequest request;

	request.bmRequestType = USBmakebmRequestType(kUSBOut, kUSBClass, kUSBInterface);
    request.bRequest = DFU_CLRSTATUS;
    request.wValue = 0;
    request.wIndex = 0;
    request.wLength = 0;
    request.pData = 0;
    
    return [self sendRequest:&request];
}

-(NSString*) stringForDFUStatus:(int) status
{
    switch(status)
    {
        case DFU_STATUS_OK:
            return @"No error condition is present";
        case DFU_STATUS_ERROR_TARGET:
            return @"File is not targeted for use by this device";
        case DFU_STATUS_ERROR_FILE:
            return @"File is for this device but fails some vendor-specific test";
        case DFU_STATUS_ERROR_WRITE:
            return @"Device is unable to write memory";
        case DFU_STATUS_ERROR_ERASE:
            return @"Memory erase function failed";
        case DFU_STATUS_ERROR_CHECK_ERASED:
            return @"Memory erase check failed";
        case DFU_STATUS_ERROR_PROG:
            return @"Program memory function failed";
        case DFU_STATUS_ERROR_VERIFY:
            return @"Program memory failed verification";
        case DFU_STATUS_ERROR_ADDRESS:
            return @"Cannot program memory due to received address that is out of range";
        case DFU_STATUS_ERROR_NOTDONE:
            return @"Received DFU_DNLOAD with wLength = 0, but device does not think that it has all data yet";
        case DFU_STATUS_ERROR_FIRMWARE:
            return @"Device's firmware is corrupt. It cannot return to run-time (non-DFU) operations";
        case DFU_STATUS_ERROR_VENDOR:
            return @"iString indicates a vendor specific error";
        case DFU_STATUS_ERROR_USBR:
            return @"Device detected unexpected USB reset signalling";
        case DFU_STATUS_ERROR_POR:
            return @"Device detected unexpected power on reset";
        case DFU_STATUS_ERROR_UNKNOWN:
            return @"Something went wrong, but the device does not know what it was";
        case DFU_STATUS_ERROR_STALLEDPKT:
            return @"Device stalled an unexpected request";
    }

    return nil;
}

@end