DnsProxy class: DNS proxy for DoH, implement with Objective-C

DnsProxy.h

//
//  DnsProxy.h
//

#import <Foundation/Foundation.h>

/**
 * DNS Proxy
 *
 * This class is responsible for setting up a DNS proxy server that listens on a local port and
 * forwards DNS queries to a remote DoH server, and then forwards the responses back to the client.
 *
 * The default DoH server is Cloudflare's with URL https://1.1.1.1/dns-query
 * The default listen port is 53
 */

@interface DnsProxy : NSObject
@property (nonatomic, copy) NSString* dohServerUrl;
@property (nonatomic, copy) NSString* listenIpAddress;
@property (nonatomic, assign) uint16_t listenPort;
@property (nonatomic, copy) void (^logHandler)(NSString* logMessage);
- (void)start;
@end

DnsProxy.m

//
//  DnsProxy.m
//

#import "DnsProxy.h"
#import <CocoaAsyncSocket/GCDAsyncUdpSocket.h>
#include <arpa/inet.h>
#include <netinet/in.h>

@interface DnsProxy () <GCDAsyncUdpSocketDelegate>
@end

@implementation DnsProxy {
    GCDAsyncUdpSocket* _udpSocket;
    NSURL* _url;
}

- (instancetype)init {
    self = [super init];
    if (self) {
        self->_listenIpAddress = @"127.0.0.1";
        self->_listenPort = 53;
        self->_dohServerUrl = @"https://1.1.1.1/dns-query";
    }
    return self;
}

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

- (void)start {
    _url = [NSURL URLWithString:self.dohServerUrl];
    if (!_url) {
        if (self.logHandler) {
            self.logHandler([NSString stringWithFormat:@"Invalid URL: %@", self.dohServerUrl]);
        }
        return;
    }

    _udpSocket = [[GCDAsyncUdpSocket alloc] initWithDelegate:self delegateQueue:dispatch_get_main_queue()];

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_len = sizeof(addr);
    addr.sin_family = AF_INET;
    addr.sin_port = htons(self.listenPort);
    addr.sin_addr.s_addr = inet_addr([self.listenIpAddress UTF8String]);

    NSData* address = [NSData dataWithBytes:&addr length:sizeof(addr)];

    NSError* error = nil;
    [_udpSocket bindToAddress:address error:&error];
    if (error) {
        if (self.logHandler) {
            self.logHandler([NSString stringWithFormat:@"bindToAddress error: %@", error]);
        }
        return;
    }

    [_udpSocket beginReceiving:&error];
    if (error) {
        if (self.logHandler) {
            self.logHandler([NSString stringWithFormat:@"beginReceiving error: %@", error]);
        }
        return;
    }
}

- (void)udpSocket:(GCDAsyncUdpSocket*)sock
       didReceiveData:(NSData*)data
          fromAddress:(NSData*)address
    withFilterContext:(nullable id)filterContext {
    NSMutableURLRequest* request = [NSMutableURLRequest requestWithURL:self->_url];
    request.HTTPMethod = @"POST";
    [request setValue:@"application/dns-message" forHTTPHeaderField:@"Content-Type"];
    request.HTTPBody = data;

    NSString* domain = [self extractDomainFromData:data];
    if (self.logHandler) {
        self.logHandler([NSString stringWithFormat:@"DNS query for domain: %@ starting", domain]);
    }

    NSURLSessionDataTask* task = [[NSURLSession sharedSession]
        dataTaskWithRequest:request
          completionHandler:^(NSData* _Nullable data, NSURLResponse* _Nullable response, NSError* _Nullable error) {
              if (error) {
                  if (self.logHandler) {
                      self.logHandler([NSString stringWithFormat:@"DNS query for domain: %@ failed", domain]);
                  }
                  return;
              }
              [sock sendData:data toAddress:address withTimeout:-1 tag:0];
              if (self.logHandler) {
                  self.logHandler([NSString stringWithFormat:@"DNS query for domain: %@ succeeded", domain]);
              }
          }];
    [task resume];
}

- (NSString*)extractDomainFromData:(NSData*)data {
    uint8_t* bytes = (uint8_t*)[data bytes];
    NSMutableString* domain = [NSMutableString string];
    int i = 12; // Skip the header
    while (i < [data length]) {
        uint8_t len = bytes[i++];
        while (len--) {
            uint8_t byte = bytes[i++];
            [domain appendFormat:@"%c", byte];
        }
        if (bytes[i] != 0) {
            [domain appendString:@"."];
        } else {
            break;
        }
    }
    return [domain copy];
}

@end

Usage.m

    DnsProxy* _dndProxy;

    _dndProxy = [[DnsProxy alloc] init];
    _dndProxy.dohServerUrl = @"https://1.1.1.1/dns-query";
    _dndProxy.listenPort = 53;
    _dndProxy.listenIpAddress = @"127.0.0.1";
    __weak typeof(self) weakSelf2 = self;
    _dndProxy.logHandler = ^(NSString* message) { [weakSelf2 postWormholeMessage:message verbose:3]; };
    [_dndProxy start];