Index: src/OFDNSResolver.h ================================================================== --- src/OFDNSResolver.h +++ src/OFDNSResolver.h @@ -32,10 +32,11 @@ @class OFDNSResolverSettings; @class OFDate; @class OFDictionary OF_GENERIC(KeyType, ObjectType); @class OFMutableDictionary OF_GENERIC(KeyType, ObjectType); @class OFNumber; +@class OFTCPSocket; @class OFUDPSocket; /*! * @enum of_dns_resolver_error_t OFDNSResolver.h ObjFW/OFDNSResolver.h * @@ -129,10 +130,12 @@ OFUDPSocket *_IPv6Socket; #endif char _buffer[OF_DNS_RESOLVER_BUFFER_LENGTH]; OFMutableDictionary OF_GENERIC(OFNumber *, OFDNSResolverContext *) *_queries; + OFMutableDictionary OF_GENERIC(OFTCPSocket *, OFDNSResolverContext *) + *_TCPQueries; } /*! * @brief A dictionary of static hosts. * Index: src/OFDNSResolver.m ================================================================== --- src/OFDNSResolver.m +++ src/OFDNSResolver.m @@ -29,10 +29,11 @@ #import "OFDictionary.h" #import "OFHostAddressResolver.h" #import "OFNumber.h" #import "OFPair.h" #import "OFString.h" +#import "OFTCPSocket.h" #import "OFTimer.h" #import "OFUDPSocket.h" #import "OFUDPSocket+Private.h" #import "OFDNSQueryFailedException.h" @@ -46,10 +47,11 @@ #ifndef SOCK_DNS # define SOCK_DNS 0 #endif #define BUFFER_LENGTH OF_DNS_RESOLVER_BUFFER_LENGTH +#define MAX_DNS_RESPONSE_LENGTH 65536 /* * RFC 1035 doesn't specify if pointers to pointers are allowed, and if so how * many. Since it's unspecified, we have to assume that it might happen, but we * also want to limit it to avoid DoS. Limiting it to 16 levels of pointers and @@ -57,17 +59,11 @@ */ #define MAX_ALLOWED_POINTERS 16 #define CNAME_RECURSION 3 -/* - * TODO: - * - * - Fallback to TCP - */ - -@interface OFDNSResolver () +@interface OFDNSResolver () - (void)of_contextTimedOut: (OFDNSResolverContext *)context; @end @interface OFDNSResolverContext: OFObject { @@ -78,10 +74,14 @@ size_t _nameServersIndex; unsigned int _attempt; id _delegate; OFData *_queryData; of_socket_address_t _usedNameServer; + OFTCPSocket *_TCPSocket; + OFMutableData *_TCPQueryData; + void *_TCPBuffer; + size_t _responseLength; OFTimer *_cancelTimer; } - (instancetype)initWithQuery: (OFDNSQuery *)query ID: (OFNumber *)ID @@ -554,10 +554,12 @@ [_query release]; [_ID release]; [_settings release]; [_delegate release]; [_queryData release]; + [_TCPSocket release]; + [_TCPQueryData release]; [_cancelTimer release]; [super dealloc]; } @end @@ -585,10 +587,11 @@ self = [super init]; @try { _settings = [[OFDNSResolverSettings alloc] init]; _queries = [[OFMutableDictionary alloc] init]; + _TCPQueries = [[OFMutableDictionary alloc] init]; [_settings reload]; } @catch (id e) { [self release]; @throw e; @@ -607,10 +610,11 @@ #ifdef OF_HAVE_IPV6 [_IPv6Socket cancelAsyncRequests]; [_IPv6Socket release]; #endif [_queries release]; + [_TCPQueries release]; [super dealloc]; } - (OFDictionary *)staticHosts @@ -726,10 +730,25 @@ [[OFRunLoop currentRunLoop] addTimer: context->_cancelTimer forMode: runLoopMode]; nameServer = [context->_settings->_nameServers objectAtIndex: context->_nameServersIndex]; + + if (context->_settings->_usesTCP) { + OF_ENSURE(context->_TCPSocket == nil); + + context->_TCPSocket = [[OFTCPSocket alloc] init]; + [_TCPQueries setObject: context + forKey: context->_TCPSocket]; + + context->_TCPSocket.delegate = self; + [context->_TCPSocket asyncConnectToHost: nameServer + port: 53 + runLoopMode: runLoopMode]; + return; + } + context->_usedNameServer = of_socket_address_parse_ip(nameServer, 53); switch (context->_usedNameServer.family) { #ifdef OF_HAVE_IPV6 case OF_SOCKET_ADDRESS_FAMILY_IPV6: @@ -813,21 +832,30 @@ - (void)of_contextTimedOut: (OFDNSResolverContext *)context { of_run_loop_mode_t runLoopMode = [OFRunLoop currentRunLoop].currentMode; OFDNSQueryFailedException *exception; + + if (context->_TCPSocket != nil) { + context->_TCPSocket.delegate = nil; + [context->_TCPSocket cancelAsyncRequests]; + + [_TCPQueries removeObjectForKey: context->_TCPSocket]; + [context->_TCPSocket release]; + context->_TCPSocket = nil; + context->_responseLength = 0; + } if (context->_nameServersIndex + 1 < context->_settings->_nameServers.count) { context->_nameServersIndex++; [self of_sendQueryForContext: context runLoopMode: runLoopMode]; return; } - if (context->_attempt < context->_settings->_maxAttempts) { - context->_attempt++; + if (++context->_attempt < context->_settings->_maxAttempts) { context->_nameServersIndex = 0; [self of_sendQueryForContext: context runLoopMode: runLoopMode]; return; } @@ -858,26 +886,21 @@ didPerformQuery: context->_query response: nil exception: exception]; } -- (bool)socket: (OFUDPSocket *)sock - didReceiveIntoBuffer: (void *)buffer_ - length: (size_t)length - sender: (const of_socket_address_t *)sender - exception: (id)exception +- (bool)of_handleResponseBuffer: (unsigned char *)buffer + length: (size_t)length + sender: (const of_socket_address_t *)sender { - unsigned char *buffer = buffer_; OFDictionary *answerRecords = nil, *authorityRecords = nil; OFDictionary *additionalRecords = nil; OFDNSResponse *response = nil; + id exception = nil; OFNumber *ID; OFDNSResolverContext *context; - if (exception != nil) - return true; - if (length < 2) /* We can't get the ID to get the context. Ignore packet. */ return true; ID = [OFNumber numberWithUInt16: (buffer[0] << 8) | buffer[1]]; @@ -884,11 +907,14 @@ context = [[[_queries objectForKey: ID] retain] autorelease]; if (context == nil) return true; - if (!of_socket_address_equal(sender, &context->_usedNameServer)) + if (context->_TCPSocket != nil) { + if ([_TCPQueries objectForKey: context->_TCPSocket] != context) + return true; + } else if (!of_socket_address_equal(sender, &context->_usedNameServer)) return true; [context->_cancelTimer invalidate]; [context->_cancelTimer release]; context->_cancelTimer = nil; @@ -1010,10 +1036,152 @@ [context->_delegate resolver: self didPerformQuery: context->_query response: response exception: exception]; + return false; +} + +- (bool)socket: (OFUDPSocket *)sock + didReceiveIntoBuffer: (void *)buffer + length: (size_t)length + sender: (const of_socket_address_t *)sender + exception: (id)exception +{ + if (exception != nil) + return true; + + return [self of_handleResponseBuffer: buffer + length: length + sender: sender]; +} + +- (void)socket: (OFTCPSocket *)sock + didConnectToHost: (OFString *)host + port: (uint16_t)port + exception: (id)exception +{ + OFDNSResolverContext *context = [_TCPQueries objectForKey: sock]; + + OF_ENSURE(context != nil); + + if (exception != nil) { + /* + * TODO: Handle error immediately instead of waiting for the + * timer to try the next nameserver or to retry. + */ + [_TCPQueries removeObjectForKey: context->_TCPSocket]; + [context->_TCPSocket release]; + context->_TCPSocket = nil; + context->_responseLength = 0; + return; + } + + if (context->_TCPQueryData == nil) { + size_t queryDataCount = context->_queryData.count; + uint16_t tmp; + + if (queryDataCount > UINT16_MAX) + @throw [OFOutOfRangeException exception]; + + context->_TCPQueryData = [[OFMutableData alloc] + initWithCapacity: queryDataCount + 2]; + + tmp = OF_BSWAP16_IF_LE(queryDataCount); + [context->_TCPQueryData addItems: &tmp + count: sizeof(tmp)]; + [context->_TCPQueryData addItems: context->_queryData.items + count: queryDataCount]; + } + + [sock asyncWriteData: context->_TCPQueryData]; +} + +- (OFData *)stream: (OFStream *)stream + didWriteData: (OFData *)data + bytesWritten: (size_t)bytesWritten + exception: (id)exception +{ + OFTCPSocket *sock = (OFTCPSocket *)stream; + OFDNSResolverContext *context = [_TCPQueries objectForKey: sock]; + + OF_ENSURE(context != nil); + + if (exception != nil) { + /* + * TODO: Handle error immediately instead of waiting for the + * timer to try the next nameserver or to retry. + */ + [_TCPQueries removeObjectForKey: context->_TCPSocket]; + [context->_TCPSocket release]; + context->_TCPSocket = nil; + context->_responseLength = 0; + return nil; + } + + if (context->_TCPBuffer == nil) + context->_TCPBuffer = + [context allocMemoryWithSize: MAX_DNS_RESPONSE_LENGTH]; + + [sock asyncReadIntoBuffer: context->_TCPBuffer + exactLength: 2]; + return nil; +} + +- (bool)stream: (OFStream *)stream + didReadIntoBuffer: (void *)buffer + length: (size_t)length + exception: (id)exception +{ + OFTCPSocket *sock = (OFTCPSocket *)stream; + OFDNSResolverContext *context = [_TCPQueries objectForKey: sock]; + + OF_ENSURE(context != nil); + + if (exception != nil) { + /* + * TODO: Handle error immediately instead of waiting for the + * timer to try the next nameserver or to retry. + */ + goto done; + } + + if (context->_responseLength == 0) { + unsigned char *ucBuffer = buffer; + + OF_ENSURE(length == 2); + + context->_responseLength = (ucBuffer[0] << 8) | ucBuffer[1]; + + if (context->_responseLength > MAX_DNS_RESPONSE_LENGTH) + @throw [OFOutOfRangeException exception]; + + if (context->_responseLength == 0) + goto done; + + [sock asyncReadIntoBuffer: context->_TCPBuffer + exactLength: context->_responseLength]; + return false; + } + + if (length != context->_responseLength) + /* + * The connection was closed before we received the entire + * response. + */ + goto done; + + [self of_handleResponseBuffer: buffer + length: length + sender: NULL]; + +done: + [_TCPQueries removeObjectForKey: context->_TCPSocket]; + [context->_TCPSocket release]; + context->_TCPSocket = nil; + context->_responseLength = 0; + return false; } - (void)asyncResolveAddressesForHost: (OFString *)host delegate: (id )delegate