ObjFW  Check-in [45452562b7]

@class OFDNSResolver;
@class OFDNSResolverContext;
@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
 *
 * @brief An enum describing why resolving a host failed.
 */

	OFUDPSocket *_IPv4Socket;
#ifdef OF_HAVE_IPV6
	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.
 *
 * This dictionary is checked before actually looking up a host.
 */

#import "OFData.h"
#import "OFDate.h"
#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"
#import "OFInitializationFailedException.h"
#import "OFInvalidArgumentException.h"
#import "OFInvalidFormatException.h"
#import "OFInvalidServerReplyException.h"
#import "OFOutOfRangeException.h"
#import "OFTruncatedDataException.h"

#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
 * immediately rejecting pointers to itself seems like a fair balance.
 */
#define MAX_ALLOWED_POINTERS 16

#define CNAME_RECURSION 3

/*
 * TODO:
 *
 *  - Fallback to TCP
 */

@interface OFDNSResolver () <OFUDPSocketDelegate>
@interface OFDNSResolver () <OFUDPSocketDelegate, OFTCPSocketDelegate>
- (void)of_contextTimedOut: (OFDNSResolverContext *)context;
@end

@interface OFDNSResolverContext: OFObject
{
@public
	OFDNSQuery *_query;
	OFNumber *_ID;
	OFDNSResolverSettings *_settings;
	size_t _nameServersIndex;
	unsigned int _attempt;
	id <OFDNSResolverQueryDelegate> _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
		     settings: (OFDNSResolverSettings *)settings
		     delegate: (id <OFDNSResolverQueryDelegate>)delegate;

- (void)dealloc
{
	[_query release];
	[_ID release];
	[_settings release];
	[_delegate release];
	[_queryData release];
	[_TCPSocket release];
	[_TCPQueryData release];
	[_cancelTimer release];

	[super dealloc];
}
@end

@implementation OFDNSResolver

- (instancetype)init
{
	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;
	}

	[_IPv4Socket cancelAsyncRequests];
	[_IPv4Socket release];
#ifdef OF_HAVE_IPV6
	[_IPv6Socket cancelAsyncRequests];
	[_IPv6Socket release];
#endif
	[_queries release];
	[_TCPQueries release];

	[super dealloc];
}

- (OFDictionary *)staticHosts
{
	return _settings->_staticHosts;

		      object: context
		     repeats: false];
	[[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:
		if (_IPv6Socket == nil) {
			of_socket_address_t address =

	objc_autoreleasePoolPop(pool);
}

- (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) {
	if (++context->_attempt < context->_settings->_maxAttempts) {
		context->_attempt++;
		context->_nameServersIndex = 0;
		[self of_sendQueryForContext: context
				 runLoopMode: runLoopMode];
		return;
	}

	context = [[context retain] autorelease];

	    @selector(resolver:didPerformQuery:response:exception:)])
		[context->_delegate resolver: self
			     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
- (bool)of_handleResponseBuffer: (unsigned char *)buffer
			 length: (size_t)length
			 sender: (const of_socket_address_t *)sender
	     exception: (id)exception
{
	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]];
	context = [[[_queries objectForKey: ID] retain] autorelease];

	if (context == nil)
		return true;

	if (context->_TCPSocket != nil) {
		if ([_TCPQueries objectForKey: context->_TCPSocket] != context)
			return true;
	if (!of_socket_address_equal(sender, &context->_usedNameServer))
	} else if (!of_socket_address_equal(sender, &context->_usedNameServer))
		return true;

	[context->_cancelTimer invalidate];
	[context->_cancelTimer release];
	context->_cancelTimer = nil;
	[_queries removeObjectForKey: ID];

	if ([context->_delegate respondsToSelector:
	    @selector(resolver:didPerformQuery:response:exception:)])
		[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 <OFDNSResolverHostDelegate>)delegate
{
	[self asyncResolveAddressesForHost: host