From 5f9c696ac75e5e7a32e956337b8fd4051dbac84e Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 09:39:03 -0600 Subject: [PATCH 1/9] Add TCP DNS transport with connection limits, idle timeout, and frame guard --- Sources/DNSServer/DNSServer+Handle.swift | 23 +- Sources/DNSServer/DNSServer+TCPHandle.swift | 127 ++++++++++ Sources/DNSServer/DNSServer.swift | 137 +++++++++- .../Helpers/APIServer/APIServer+Start.swift | 28 +- Tests/DNSServerTests/TCPHandleTest.swift | 239 ++++++++++++++++++ 5 files changed, 528 insertions(+), 26 deletions(-) create mode 100644 Sources/DNSServer/DNSServer+TCPHandle.swift create mode 100644 Tests/DNSServerTests/TCPHandleTest.swift diff --git a/Sources/DNSServer/DNSServer+Handle.swift b/Sources/DNSServer/DNSServer+Handle.swift index e258fb482..8ba4dbfe1 100644 --- a/Sources/DNSServer/DNSServer+Handle.swift +++ b/Sources/DNSServer/DNSServer+Handle.swift @@ -19,10 +19,7 @@ import NIOCore import NIOPosix extension DNSServer { - /// Handles the DNS request. - /// - Parameters: - /// - outbound: The NIOAsyncChannelOutboundWriter for which to respond. - /// - packet: The request packet. + /// Handles a UDP DNS request and writes the response back to the sender. func handle( outbound: NIOAsyncChannelOutboundWriter>, packet: inout AddressedEnvelope @@ -37,11 +34,20 @@ extension DNSServer { } } + self.log?.debug("sending response for request") + let responseData = try await self.processRaw(data: data) + let rData = ByteBuffer(bytes: responseData) + try? await outbound.write(AddressedEnvelope(remoteAddress: packet.remoteAddress, data: rData)) + self.log?.debug("processing done") + } + + /// Deserializes a raw DNS query, runs it through the handler chain, and returns + /// the serialized response bytes. Used by both UDP and TCP transports. + func processRaw(data: Data) async throws -> Data { self.log?.debug("deserializing message") let query = try Message(deserialize: data) self.log?.debug("processing query: \(query.questions)") - // always send response let responseData: Data do { self.log?.debug("awaiting processing") @@ -76,11 +82,6 @@ extension DNSServer { responseData = try response.serialize() } - self.log?.debug("sending response for \(query.id)") - let rData = ByteBuffer(bytes: responseData) - try? await outbound.write(AddressedEnvelope(remoteAddress: packet.remoteAddress, data: rData)) - - self.log?.debug("processing done") - + return responseData } } diff --git a/Sources/DNSServer/DNSServer+TCPHandle.swift b/Sources/DNSServer/DNSServer+TCPHandle.swift new file mode 100644 index 000000000..e501d9d0b --- /dev/null +++ b/Sources/DNSServer/DNSServer+TCPHandle.swift @@ -0,0 +1,127 @@ +//===----------------------------------------------------------------------===// +// Copyright © 2025-2026 Apple Inc. and the container project authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +import Foundation +import NIOCore +import NIOPosix +import Synchronization + +/// Tracks the timestamp of the last fully-processed DNS message on a TCP connection. +/// The idle watchdog reads this to decide when to close a stale connection. +private actor ConnectionActivity { + private var lastSeen: ContinuousClock.Instant + + init() { + lastSeen = ContinuousClock.now + } + + func ping() { + lastSeen = ContinuousClock.now + } + + func idle(after duration: Duration) -> Bool { + ContinuousClock.now - lastSeen >= duration + } +} + +extension DNSServer { + /// Handles a single inbound TCP DNS connection. + /// + /// DNS over TCP (RFC 1035 §4.2.2) prefixes every message with a 2-byte big-endian + /// length field. Bytes are accumulated until a complete message is available, processed + /// through the handler chain, and a length-prefixed response is written back. + /// Multiple messages on the same connection (pipelining) are handled in order. + /// + /// The connection is closed when: + /// - the client sends no data for `tcpIdleTimeout` between complete messages + /// - a frame exceeding `maxTCPMessageSize` is received (indicates framing desync) + /// - an unrecoverable I/O error occurs + /// + /// This method is `async` (not `async throws`) so per-connection failures never + /// propagate up to the accept loop. + func handleTCP(channel: NIOAsyncChannel) async { + do { + try await channel.executeThenClose { inbound, outbound in + let activity = ConnectionActivity() + try await withThrowingTaskGroup(of: Void.self) { group in + // Watchdog: checks every second if the connection has gone idle. + // The timeout resets after each fully-processed message, so a client + // sending pipelined queries spaced a few seconds apart is unaffected. + group.addTask { + while await !activity.idle(after: self.tcpIdleTimeout) { + try await Task.sleep(for: .seconds(1)) + } + self.log?.debug("TCP DNS: idle timeout, closing connection") + throw CancellationError() + } + + // Read loop: accumulates bytes and dispatches complete DNS messages. + group.addTask { + var buffer = ByteBuffer() + for try await chunk in inbound { + await activity.ping() + buffer.writeImmutableBuffer(chunk) + + // A single TCP segment may contain multiple back-to-back DNS messages. + while buffer.readableBytes >= 2 { + // Peek at the 2-byte length prefix without advancing the reader. + guard let msgLen = buffer.getInteger( + at: buffer.readerIndex, as: UInt16.self + ) else { break } + + // A value above our ceiling almost certainly means the framing + // is desynced — bail now rather than trying to read 60K bytes. + guard msgLen > 0, msgLen <= DNSServer.maxTCPMessageSize else { + self.log?.error( + "TCP DNS: unexpected frame size \(msgLen) bytes, closing connection") + return + } + + // Wait until the full payload has arrived before dispatching. + let needed = 2 + Int(msgLen) + guard buffer.readableBytes >= needed else { break } + + buffer.moveReaderIndex(forwardBy: 2) + let msgSlice = buffer.readSlice(length: Int(msgLen))! + let msgData = Data(msgSlice.readableBytesView) + + let responseData = try await self.processRaw(data: msgData) + + // Reset the idle clock now that we have completed a round trip. + await activity.ping() + + var out = ByteBuffer() + out.writeInteger(UInt16(responseData.count)) + out.writeBytes(responseData) + try await outbound.write(out) + } + } + } + + // First task to finish wins. The other is cancelled — any resulting + // CancellationError is intentionally ignored because no further work + // should happen on a connection that has already ended. + try await group.next() + group.cancelAll() + } + } + } catch is CancellationError { + // Expected on timeout — nothing to report. + } catch { + log?.warning("TCP DNS: connection error: \(error)") + } + } +} diff --git a/Sources/DNSServer/DNSServer.swift b/Sources/DNSServer/DNSServer.swift index d38e55459..7b832e82a 100644 --- a/Sources/DNSServer/DNSServer.swift +++ b/Sources/DNSServer/DNSServer.swift @@ -18,25 +18,64 @@ import Foundation import Logging import NIOCore import NIOPosix +import Synchronization -/// Provides a DNS server. -/// - Parameters: -/// - host: The host address on which to listen. -/// - port: The port for the server to listen. -public struct DNSServer { +/// Thread-safe active-connection counter used to enforce `maxConcurrentConnections`. +/// +/// Wrapped in a `final class` so it can be stored as a reference in the +/// `Copyable` `DNSServer` struct without running into `~Copyable` restrictions. +private final class ConnectionCounter: Sendable { + private let storage: Mutex = .init(0) + + func tryIncrement(limit: Int) -> Bool { + storage.withLock { count in + guard count < limit else { return false } + count += 1 + return true + } + } + + func decrement() { + storage.withLock { $0 -= 1 } + } +} + +/// Provides a DNS server over UDP and TCP. +public struct DNSServer: @unchecked Sendable { public var handler: DNSHandler let log: Logger? + // Maximum number of concurrent TCP connections. Connections beyond this limit + // are closed immediately. The UDP path is unaffected. + static let maxConcurrentConnections = 128 + + // Generous upper bound on a DNS message received over TCP. Legitimate messages + // are well under this limit; anything larger almost certainly indicates a framing + // desync rather than a valid query, so we close the connection instead of trying + // to allocate a huge buffer. + static let maxTCPMessageSize: UInt16 = 4096 + + // How long a TCP connection may be idle between fully-processed messages before + // it is closed. Exposed as an instance property so tests can inject a short value. + let tcpIdleTimeout: Duration + + // Active TCP connection counter. + private let connections: ConnectionCounter + public init( handler: DNSHandler, - log: Logger? = nil + log: Logger? = nil, + tcpIdleTimeout: Duration = .seconds(30) ) { self.handler = handler self.log = log + self.tcpIdleTimeout = tcpIdleTimeout + self.connections = ConnectionCounter() } + // MARK: - UDP + public func run(host: String, port: Int) async throws { - // TODO: TCP server let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .channelOption(.socketOption(.so_reuseaddr), value: 1) .bind(host: host, port: port) @@ -59,7 +98,6 @@ public struct DNSServer { } public func run(socketPath: String) async throws { - // TODO: TCP server let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .bind(unixDomainSocketPath: socketPath, cleanupExistingSocketFile: true) .flatMapThrowing { channel in @@ -82,5 +120,88 @@ public struct DNSServer { } } + // MARK: - TCP + + /// Runs a TCP DNS listener on the given host and port. + /// + /// DNS over TCP is required by RFC 1035 for responses that exceed 512 bytes (the UDP + /// wire limit). Clients signal truncation via the TC bit and retry over TCP. + /// This method runs a separate TCP server on the same port number as the UDP listener; + /// the OS distinguishes them by socket type (SOCK_STREAM vs SOCK_DGRAM). + public func runTCP(host: String, port: Int) async throws { + let server = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) + .bind( + host: host, + port: port + ) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + + try await server.executeThenClose { inbound in + try await withThrowingDiscardingTaskGroup { group in + for try await child in inbound { + guard connections.tryIncrement(limit: Self.maxConcurrentConnections) else { + log?.warning( + "TCP DNS: connection limit (\(Self.maxConcurrentConnections)) reached, dropping connection") + continue + } + + group.addTask { + defer { self.connections.decrement() } + await self.handleTCP(channel: child) + } + } + } + } + } + + /// Runs a TCP DNS listener on a Unix domain socket. + /// + /// Mirrors `runTCP(host:port:)` but binds to a socket file path instead of a + /// TCP/IP address. Any existing socket file at the path is removed before binding. + public func runTCP(socketPath: String) async throws { + try? FileManager.default.removeItem(atPath: socketPath) + + let address = try SocketAddress(unixDomainSocketPath: socketPath) + let server = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .bind(to: address) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + + try await server.executeThenClose { inbound in + try await withThrowingDiscardingTaskGroup { group in + for try await child in inbound { + guard connections.tryIncrement(limit: Self.maxConcurrentConnections) else { + log?.warning( + "TCP DNS: connection limit (\(Self.maxConcurrentConnections)) reached, dropping connection") + continue + } + + group.addTask { + defer { self.connections.decrement() } + await self.handleTCP(channel: child) + } + } + } + } + } + public func stop() async throws {} } diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index 62fa0296f..0a49111c1 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -97,22 +97,36 @@ extension APIServer { try await server.listen() } - // start up host table DNS + // start up host table DNS (UDP) + let hostsResolver = ContainerDNSHandler(networkService: networkService) + let nxDomainResolver = NxDomainResolver() + let compositeResolver = CompositeResolver(handlers: [hostsResolver, nxDomainResolver]) + let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) + let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) + group.addTask { - let hostsResolver = ContainerDNSHandler(networkService: networkService) - let nxDomainResolver = NxDomainResolver() - let compositeResolver = CompositeResolver(handlers: [hostsResolver, nxDomainResolver]) - let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) - let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) log.info( - "starting DNS resolver for container hostnames", + "starting DNS resolver for container hostnames (UDP)", metadata: [ "host": "\(Self.listenAddress)", "port": "\(Self.dnsPort)", ] ) try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) + } + // TCP server on the same port. Required by RFC 1035 for responses that + // exceed the 512-byte UDP wire limit. The OS distinguishes the two by + // socket type (SOCK_DGRAM vs SOCK_STREAM), so there is no port conflict. + group.addTask { + log.info( + "starting DNS resolver for container hostnames (TCP)", + metadata: [ + "host": "\(Self.listenAddress)", + "port": "\(Self.dnsPort)", + ] + ) + try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) } // start up realhost DNS diff --git a/Tests/DNSServerTests/TCPHandleTest.swift b/Tests/DNSServerTests/TCPHandleTest.swift new file mode 100644 index 000000000..5f50b2593 --- /dev/null +++ b/Tests/DNSServerTests/TCPHandleTest.swift @@ -0,0 +1,239 @@ +//===----------------------------------------------------------------------===// +// Copyright © 2025-2026 Apple Inc. and the container project authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +import DNS +import Foundation +import NIOCore +import NIOPosix +import Testing + +@testable import DNSServer + +// MARK: - Unit tests for the shared processRaw helper + +struct ProcessRawTest { + @Test func testProcessRawReturnsSerializedResponse() async throws { + let handler = HostTableResolver(hosts4: ["myhost.": IPv4("10.0.0.1")!]) + let server = DNSServer(handler: StandardQueryValidator(handler: handler)) + + let query = Message( + id: 99, + type: .query, + questions: [Question(name: "myhost.", type: .host)] + ) + let responseData = try await server.processRaw(data: query.serialize()) + let response = try Message(deserialize: responseData) + + #expect(99 == response.id) + #expect(.response == response.type) + #expect(.noError == response.returnCode) + #expect(1 == response.answers.count) + let record = response.answers.first as? HostRecord + #expect(IPv4("10.0.0.1") == record?.ip) + } + + @Test func testProcessRawReturnsNxdomainForUnknownHost() async throws { + let handler = HostTableResolver(hosts4: ["known.": IPv4("10.0.0.1")!]) + let composite = CompositeResolver(handlers: [handler, NxDomainResolver()]) + let server = DNSServer(handler: StandardQueryValidator(handler: composite)) + + let query = Message( + id: 55, + type: .query, + questions: [Question(name: "unknown.", type: .host)] + ) + let responseData = try await server.processRaw(data: query.serialize()) + let response = try Message(deserialize: responseData) + + #expect(55 == response.id) + #expect(.nonExistentDomain == response.returnCode) + #expect(0 == response.answers.count) + } +} + +// MARK: - TCP integration tests + +struct TCPHandleTest { + /// Serializes a DNS query and wraps it in the 2-byte length prefix required by TCP. + private func makeTCPFrame(hostname: String, id: UInt16 = 1) throws -> ByteBuffer { + let bytes = try Message( + id: id, + type: .query, + questions: [Question(name: hostname, type: .host)] + ).serialize() + var buf = ByteBuffer() + buf.writeInteger(UInt16(bytes.count)) + buf.writeBytes(bytes) + return buf + } + + @Test func testTCPRoundTrip() async throws { + let handler = HostTableResolver(hosts4: ["tcp-test.": IPv4("5.6.7.8")!]) + let server = DNSServer( + handler: StandardQueryValidator(handler: handler), + tcpIdleTimeout: .seconds(2) + ) + + // Bind on an ephemeral port using the same pattern as DNSServer.runTCP(). + let listener = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) + .bind( + host: "127.0.0.1", + port: 0 + ) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + let port = listener.channel.localAddress!.port! + + // Server task: accept one connection, handle it, then we close the listener. + async let serverDone: Void = listener.executeThenClose { inbound in + for try await child in inbound { + await server.handleTCP(channel: child) + break + } + } + + // Client: connect, send a query, read the response. + let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) + .connect( + host: "127.0.0.1", + port: port + ) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + + var response: Message? + try await client.executeThenClose { inbound, outbound in + try await outbound.write(try makeTCPFrame(hostname: "tcp-test.", id: 77)) + for try await var chunk in inbound { + let len = chunk.readInteger(as: UInt16.self)! + let slice = chunk.readSlice(length: Int(len))! + response = try Message(deserialize: Data(slice.readableBytesView)) + break + } + } + + // Unblock the server accept loop. + try await listener.channel.close() + try? await serverDone + + #expect(77 == response?.id) + #expect(.noError == response?.returnCode) + #expect(1 == response?.answers.count) + let record = response?.answers.first as? HostRecord + #expect(IPv4("5.6.7.8") == record?.ip) + } + + @Test func testTCPPipelinedQueries() async throws { + let handler = HostTableResolver(hosts4: [ + "first.": IPv4("1.1.1.1")!, + "second.": IPv4("2.2.2.2")!, + ]) + let server = DNSServer( + handler: StandardQueryValidator(handler: handler), + tcpIdleTimeout: .seconds(2) + ) + + let listener = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) + .bind( + host: "127.0.0.1", + port: 0 + ) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + let port = listener.channel.localAddress!.port! + + async let serverDone: Void = listener.executeThenClose { inbound in + for try await child in inbound { + await server.handleTCP(channel: child) + break + } + } + + let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) + .connect( + host: "127.0.0.1", + port: port + ) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init( + inboundType: ByteBuffer.self, + outboundType: ByteBuffer.self + ) + ) + } + } + + var responses: [Message] = [] + try await client.executeThenClose { inbound, outbound in + // Write two queries in a single write to simulate pipelining. + var combined = try makeTCPFrame(hostname: "first.", id: 10) + combined.writeImmutableBuffer(try makeTCPFrame(hostname: "second.", id: 20)) + try await outbound.write(combined) + + var accumulator = ByteBuffer() + for try await chunk in inbound { + accumulator.writeImmutableBuffer(chunk) + while accumulator.readableBytes >= 2 { + guard let len = accumulator.getInteger(at: accumulator.readerIndex, as: UInt16.self) else { break } + guard accumulator.readableBytes >= 2 + Int(len) else { break } + accumulator.moveReaderIndex(forwardBy: 2) + let slice = accumulator.readSlice(length: Int(len))! + responses.append(try Message(deserialize: Data(slice.readableBytesView))) + } + if responses.count == 2 { break } + } + } + + try await listener.channel.close() + try? await serverDone + + #expect(2 == responses.count) + #expect(10 == responses[0].id) + #expect(20 == responses[1].id) + let a1 = responses[0].answers.first as? HostRecord + let a2 = responses[1].answers.first as? HostRecord + #expect(IPv4("1.1.1.1") == a1?.ip) + #expect(IPv4("2.2.2.2") == a2?.ip) + } +} From 127fb6ddad4d5a287814bc7b2c63ba774cda8333 Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:03:11 -0600 Subject: [PATCH 2/9] Standardize SwiftDoc and inline comments for TCP DNS implementation --- Sources/DNSServer/DNSServer.swift | 20 ++++++++++--------- .../Helpers/APIServer/APIServer+Start.swift | 9 ++++++--- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/Sources/DNSServer/DNSServer.swift b/Sources/DNSServer/DNSServer.swift index 7b832e82a..db1e4900a 100644 --- a/Sources/DNSServer/DNSServer.swift +++ b/Sources/DNSServer/DNSServer.swift @@ -45,21 +45,21 @@ public struct DNSServer: @unchecked Sendable { public var handler: DNSHandler let log: Logger? - // Maximum number of concurrent TCP connections. Connections beyond this limit - // are closed immediately. The UDP path is unaffected. + /// Maximum number of concurrent TCP connections. + /// Connections beyond this limit are closed immediately. The UDP path is unaffected. static let maxConcurrentConnections = 128 - // Generous upper bound on a DNS message received over TCP. Legitimate messages - // are well under this limit; anything larger almost certainly indicates a framing - // desync rather than a valid query, so we close the connection instead of trying - // to allocate a huge buffer. + /// Generous upper bound on a DNS message received over TCP. + /// Legitimate messages are well under this limit; anything larger almost certainly + /// indicates a framing desync rather than a valid query, so we close the connection + /// instead of trying to allocate a huge buffer. static let maxTCPMessageSize: UInt16 = 4096 - // How long a TCP connection may be idle between fully-processed messages before - // it is closed. Exposed as an instance property so tests can inject a short value. + /// How long a TCP connection may be idle between fully-processed messages before + /// it is closed. Exposed as an instance property so tests can inject a short value. let tcpIdleTimeout: Duration - // Active TCP connection counter. + /// Active TCP connection counter. private let connections: ConnectionCounter public init( @@ -75,6 +75,7 @@ public struct DNSServer: @unchecked Sendable { // MARK: - UDP + /// Runs a UDP DNS listener on the given host and port. public func run(host: String, port: Int) async throws { let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .channelOption(.socketOption(.so_reuseaddr), value: 1) @@ -97,6 +98,7 @@ public struct DNSServer: @unchecked Sendable { } } + /// Runs a UDP DNS listener on a Unix domain socket. public func run(socketPath: String) async throws { let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .bind(unixDomainSocketPath: socketPath, cleanupExistingSocketFile: true) diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index 0a49111c1..687d26049 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -115,9 +115,12 @@ extension APIServer { try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) } - // TCP server on the same port. Required by RFC 1035 for responses that - // exceed the 512-byte UDP wire limit. The OS distinguishes the two by - // socket type (SOCK_DGRAM vs SOCK_STREAM), so there is no port conflict. + // start up host table DNS (TCP) + // + // Required by RFC 1035 for responses that exceed the 512-byte UDP wire + // limit. The OS distinguishes the two by socket type (SOCK_DGRAM vs + // SOCK_STREAM), so there is no port conflict on the same bind address. + group.addTask { log.info( "starting DNS resolver for container hostnames (TCP)", From 699613f0e7504f4b890b8ddf4ca27c4aaad8e4c7 Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:07:05 -0600 Subject: [PATCH 3/9] Remove descriptive inline and doc comments to match project style --- Sources/DNSServer/DNSServer+TCPHandle.swift | 30 ------------------- Sources/DNSServer/DNSServer.swift | 26 ---------------- .../Helpers/APIServer/APIServer+Start.swift | 4 --- Tests/DNSServerTests/TCPHandleTest.swift | 8 ----- 4 files changed, 68 deletions(-) diff --git a/Sources/DNSServer/DNSServer+TCPHandle.swift b/Sources/DNSServer/DNSServer+TCPHandle.swift index e501d9d0b..2d9156b9e 100644 --- a/Sources/DNSServer/DNSServer+TCPHandle.swift +++ b/Sources/DNSServer/DNSServer+TCPHandle.swift @@ -19,8 +19,6 @@ import NIOCore import NIOPosix import Synchronization -/// Tracks the timestamp of the last fully-processed DNS message on a TCP connection. -/// The idle watchdog reads this to decide when to close a stale connection. private actor ConnectionActivity { private var lastSeen: ContinuousClock.Instant @@ -38,28 +36,11 @@ private actor ConnectionActivity { } extension DNSServer { - /// Handles a single inbound TCP DNS connection. - /// - /// DNS over TCP (RFC 1035 §4.2.2) prefixes every message with a 2-byte big-endian - /// length field. Bytes are accumulated until a complete message is available, processed - /// through the handler chain, and a length-prefixed response is written back. - /// Multiple messages on the same connection (pipelining) are handled in order. - /// - /// The connection is closed when: - /// - the client sends no data for `tcpIdleTimeout` between complete messages - /// - a frame exceeding `maxTCPMessageSize` is received (indicates framing desync) - /// - an unrecoverable I/O error occurs - /// - /// This method is `async` (not `async throws`) so per-connection failures never - /// propagate up to the accept loop. func handleTCP(channel: NIOAsyncChannel) async { do { try await channel.executeThenClose { inbound, outbound in let activity = ConnectionActivity() try await withThrowingTaskGroup(of: Void.self) { group in - // Watchdog: checks every second if the connection has gone idle. - // The timeout resets after each fully-processed message, so a client - // sending pipelined queries spaced a few seconds apart is unaffected. group.addTask { while await !activity.idle(after: self.tcpIdleTimeout) { try await Task.sleep(for: .seconds(1)) @@ -68,29 +49,23 @@ extension DNSServer { throw CancellationError() } - // Read loop: accumulates bytes and dispatches complete DNS messages. group.addTask { var buffer = ByteBuffer() for try await chunk in inbound { await activity.ping() buffer.writeImmutableBuffer(chunk) - // A single TCP segment may contain multiple back-to-back DNS messages. while buffer.readableBytes >= 2 { - // Peek at the 2-byte length prefix without advancing the reader. guard let msgLen = buffer.getInteger( at: buffer.readerIndex, as: UInt16.self ) else { break } - // A value above our ceiling almost certainly means the framing - // is desynced — bail now rather than trying to read 60K bytes. guard msgLen > 0, msgLen <= DNSServer.maxTCPMessageSize else { self.log?.error( "TCP DNS: unexpected frame size \(msgLen) bytes, closing connection") return } - // Wait until the full payload has arrived before dispatching. let needed = 2 + Int(msgLen) guard buffer.readableBytes >= needed else { break } @@ -100,7 +75,6 @@ extension DNSServer { let responseData = try await self.processRaw(data: msgData) - // Reset the idle clock now that we have completed a round trip. await activity.ping() var out = ByteBuffer() @@ -111,15 +85,11 @@ extension DNSServer { } } - // First task to finish wins. The other is cancelled — any resulting - // CancellationError is intentionally ignored because no further work - // should happen on a connection that has already ended. try await group.next() group.cancelAll() } } } catch is CancellationError { - // Expected on timeout — nothing to report. } catch { log?.warning("TCP DNS: connection error: \(error)") } diff --git a/Sources/DNSServer/DNSServer.swift b/Sources/DNSServer/DNSServer.swift index db1e4900a..7a6065888 100644 --- a/Sources/DNSServer/DNSServer.swift +++ b/Sources/DNSServer/DNSServer.swift @@ -20,10 +20,6 @@ import NIOCore import NIOPosix import Synchronization -/// Thread-safe active-connection counter used to enforce `maxConcurrentConnections`. -/// -/// Wrapped in a `final class` so it can be stored as a reference in the -/// `Copyable` `DNSServer` struct without running into `~Copyable` restrictions. private final class ConnectionCounter: Sendable { private let storage: Mutex = .init(0) @@ -40,26 +36,16 @@ private final class ConnectionCounter: Sendable { } } -/// Provides a DNS server over UDP and TCP. public struct DNSServer: @unchecked Sendable { public var handler: DNSHandler let log: Logger? - /// Maximum number of concurrent TCP connections. - /// Connections beyond this limit are closed immediately. The UDP path is unaffected. static let maxConcurrentConnections = 128 - /// Generous upper bound on a DNS message received over TCP. - /// Legitimate messages are well under this limit; anything larger almost certainly - /// indicates a framing desync rather than a valid query, so we close the connection - /// instead of trying to allocate a huge buffer. static let maxTCPMessageSize: UInt16 = 4096 - /// How long a TCP connection may be idle between fully-processed messages before - /// it is closed. Exposed as an instance property so tests can inject a short value. let tcpIdleTimeout: Duration - /// Active TCP connection counter. private let connections: ConnectionCounter public init( @@ -75,7 +61,6 @@ public struct DNSServer: @unchecked Sendable { // MARK: - UDP - /// Runs a UDP DNS listener on the given host and port. public func run(host: String, port: Int) async throws { let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .channelOption(.socketOption(.so_reuseaddr), value: 1) @@ -98,7 +83,6 @@ public struct DNSServer: @unchecked Sendable { } } - /// Runs a UDP DNS listener on a Unix domain socket. public func run(socketPath: String) async throws { let srv = try await DatagramBootstrap(group: NIOSingletons.posixEventLoopGroup) .bind(unixDomainSocketPath: socketPath, cleanupExistingSocketFile: true) @@ -124,12 +108,6 @@ public struct DNSServer: @unchecked Sendable { // MARK: - TCP - /// Runs a TCP DNS listener on the given host and port. - /// - /// DNS over TCP is required by RFC 1035 for responses that exceed 512 bytes (the UDP - /// wire limit). Clients signal truncation via the TC bit and retry over TCP. - /// This method runs a separate TCP server on the same port number as the UDP listener; - /// the OS distinguishes them by socket type (SOCK_STREAM vs SOCK_DGRAM). public func runTCP(host: String, port: Int) async throws { let server = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) @@ -166,10 +144,6 @@ public struct DNSServer: @unchecked Sendable { } } - /// Runs a TCP DNS listener on a Unix domain socket. - /// - /// Mirrors `runTCP(host:port:)` but binds to a socket file path instead of a - /// TCP/IP address. Any existing socket file at the path is removed before binding. public func runTCP(socketPath: String) async throws { try? FileManager.default.removeItem(atPath: socketPath) diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index 687d26049..019352d2e 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -116,10 +116,6 @@ extension APIServer { } // start up host table DNS (TCP) - // - // Required by RFC 1035 for responses that exceed the 512-byte UDP wire - // limit. The OS distinguishes the two by socket type (SOCK_DGRAM vs - // SOCK_STREAM), so there is no port conflict on the same bind address. group.addTask { log.info( diff --git a/Tests/DNSServerTests/TCPHandleTest.swift b/Tests/DNSServerTests/TCPHandleTest.swift index 5f50b2593..bb5656ed2 100644 --- a/Tests/DNSServerTests/TCPHandleTest.swift +++ b/Tests/DNSServerTests/TCPHandleTest.swift @@ -22,7 +22,6 @@ import Testing @testable import DNSServer -// MARK: - Unit tests for the shared processRaw helper struct ProcessRawTest { @Test func testProcessRawReturnsSerializedResponse() async throws { @@ -64,10 +63,8 @@ struct ProcessRawTest { } } -// MARK: - TCP integration tests struct TCPHandleTest { - /// Serializes a DNS query and wraps it in the 2-byte length prefix required by TCP. private func makeTCPFrame(hostname: String, id: UInt16 = 1) throws -> ByteBuffer { let bytes = try Message( id: id, @@ -87,7 +84,6 @@ struct TCPHandleTest { tcpIdleTimeout: .seconds(2) ) - // Bind on an ephemeral port using the same pattern as DNSServer.runTCP(). let listener = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) .bind( @@ -106,7 +102,6 @@ struct TCPHandleTest { } let port = listener.channel.localAddress!.port! - // Server task: accept one connection, handle it, then we close the listener. async let serverDone: Void = listener.executeThenClose { inbound in for try await child in inbound { await server.handleTCP(channel: child) @@ -114,7 +109,6 @@ struct TCPHandleTest { } } - // Client: connect, send a query, read the response. let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) .connect( host: "127.0.0.1", @@ -142,7 +136,6 @@ struct TCPHandleTest { } } - // Unblock the server accept loop. try await listener.channel.close() try? await serverDone @@ -206,7 +199,6 @@ struct TCPHandleTest { var responses: [Message] = [] try await client.executeThenClose { inbound, outbound in - // Write two queries in a single write to simulate pipelining. var combined = try makeTCPFrame(hostname: "first.", id: 10) combined.writeImmutableBuffer(try makeTCPFrame(hostname: "second.", id: 20)) try await outbound.write(combined) From 1295c60f87509f3977b6d3760c9ddbc784a1c9c7 Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:33:39 -0600 Subject: [PATCH 4/9] Test: refine TCP drop connection test assertions --- Sources/DNSServer/DNSServer+Handle.swift | 10 +- Sources/DNSServer/DNSServer+TCPHandle.swift | 11 +- Tests/DNSServerTests/TCPHandleTest.swift | 115 +++++++++++++++++++- 3 files changed, 117 insertions(+), 19 deletions(-) diff --git a/Sources/DNSServer/DNSServer+Handle.swift b/Sources/DNSServer/DNSServer+Handle.swift index 8ba4dbfe1..899e1dd85 100644 --- a/Sources/DNSServer/DNSServer+Handle.swift +++ b/Sources/DNSServer/DNSServer+Handle.swift @@ -24,15 +24,7 @@ extension DNSServer { outbound: NIOAsyncChannelOutboundWriter>, packet: inout AddressedEnvelope ) async throws { - let chunkSize = 512 - var data = Data() - - self.log?.debug("reading data") - while packet.data.readableBytes > 0 { - if let chunk = packet.data.readBytes(length: min(chunkSize, packet.data.readableBytes)) { - data.append(contentsOf: chunk) - } - } + let data = Data(packet.data.readableBytesView) self.log?.debug("sending response for request") let responseData = try await self.processRaw(data: data) diff --git a/Sources/DNSServer/DNSServer+TCPHandle.swift b/Sources/DNSServer/DNSServer+TCPHandle.swift index 2d9156b9e..8d36c49cc 100644 --- a/Sources/DNSServer/DNSServer+TCPHandle.swift +++ b/Sources/DNSServer/DNSServer+TCPHandle.swift @@ -42,8 +42,10 @@ extension DNSServer { let activity = ConnectionActivity() try await withThrowingTaskGroup(of: Void.self) { group in group.addTask { - while await !activity.idle(after: self.tcpIdleTimeout) { - try await Task.sleep(for: .seconds(1)) + let pollInterval = min(Duration.seconds(1), self.tcpIdleTimeout) + while true { + try await Task.sleep(for: pollInterval) + if await activity.idle(after: self.tcpIdleTimeout) { break } } self.log?.debug("TCP DNS: idle timeout, closing connection") throw CancellationError() @@ -52,7 +54,6 @@ extension DNSServer { group.addTask { var buffer = ByteBuffer() for try await chunk in inbound { - await activity.ping() buffer.writeImmutableBuffer(chunk) while buffer.readableBytes >= 2 { @@ -60,7 +61,7 @@ extension DNSServer { at: buffer.readerIndex, as: UInt16.self ) else { break } - guard msgLen > 0, msgLen <= DNSServer.maxTCPMessageSize else { + guard msgLen > 0, msgLen <= Self.maxTCPMessageSize else { self.log?.error( "TCP DNS: unexpected frame size \(msgLen) bytes, closing connection") return @@ -70,7 +71,7 @@ extension DNSServer { guard buffer.readableBytes >= needed else { break } buffer.moveReaderIndex(forwardBy: 2) - let msgSlice = buffer.readSlice(length: Int(msgLen))! + guard let msgSlice = buffer.readSlice(length: Int(msgLen)) else { break } let msgData = Data(msgSlice.readableBytesView) let responseData = try await self.processRaw(data: msgData) diff --git a/Tests/DNSServerTests/TCPHandleTest.swift b/Tests/DNSServerTests/TCPHandleTest.swift index bb5656ed2..279bb58d7 100644 --- a/Tests/DNSServerTests/TCPHandleTest.swift +++ b/Tests/DNSServerTests/TCPHandleTest.swift @@ -129,14 +129,15 @@ struct TCPHandleTest { try await client.executeThenClose { inbound, outbound in try await outbound.write(try makeTCPFrame(hostname: "tcp-test.", id: 77)) for try await var chunk in inbound { - let len = chunk.readInteger(as: UInt16.self)! - let slice = chunk.readSlice(length: Int(len))! + // Safe in-process: the entire response will arrive in one chunk + guard let len = chunk.readInteger(as: UInt16.self) else { break } + guard let slice = chunk.readSlice(length: Int(len)) else { break } response = try Message(deserialize: Data(slice.readableBytesView)) break } } - try await listener.channel.close() + try? await listener.channel.close() try? await serverDone #expect(77 == response?.id) @@ -210,14 +211,14 @@ struct TCPHandleTest { guard let len = accumulator.getInteger(at: accumulator.readerIndex, as: UInt16.self) else { break } guard accumulator.readableBytes >= 2 + Int(len) else { break } accumulator.moveReaderIndex(forwardBy: 2) - let slice = accumulator.readSlice(length: Int(len))! + guard let slice = accumulator.readSlice(length: Int(len)) else { break } responses.append(try Message(deserialize: Data(slice.readableBytesView))) } if responses.count == 2 { break } } } - try await listener.channel.close() + try? await listener.channel.close() try? await serverDone #expect(2 == responses.count) @@ -228,4 +229,108 @@ struct TCPHandleTest { #expect(IPv4("1.1.1.1") == a1?.ip) #expect(IPv4("2.2.2.2") == a2?.ip) } + + @Test func testTCPDropsOversizedFrame() async throws { + let handler = HostTableResolver(hosts4: ["oversize.": IPv4("1.1.1.1")!]) + let server = DNSServer( + handler: StandardQueryValidator(handler: handler), + tcpIdleTimeout: .seconds(2) + ) + + let listener = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) + .bind(host: "127.0.0.1", port: 0) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + } + } + let port = listener.channel.localAddress!.port! + + async let serverDone: Void = listener.executeThenClose { inbound in + for try await child in inbound { + await server.handleTCP(channel: child) + break + } + } + + let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) + .connect(host: "127.0.0.1", port: port) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + } + } + + var receivedChunks = 0 + var connectionError: Error? + do { + try await client.executeThenClose { inbound, outbound in + var buf = ByteBuffer() + buf.writeInteger(UInt16(DNSServer.maxTCPMessageSize + 1)) + // The server inspects only the 2-byte length prefix before closing. + // The payload bytes that follow are never read. + buf.writeBytes([UInt8](repeating: 0, count: 10)) + try await outbound.write(buf) + + for try await _ in inbound { + receivedChunks += 1 + } + } + } catch { + connectionError = error + } + + try? await listener.channel.close() + try? await serverDone + + #expect(receivedChunks == 0, "Expected server to drop connection without responding to oversized frame") + } + + @Test func testTCPIdleTimeoutDropsConnection() async throws { + let handler = HostTableResolver(hosts4: ["idle.": IPv4("1.1.1.1")!]) + let server = DNSServer( + handler: StandardQueryValidator(handler: handler), + tcpIdleTimeout: .milliseconds(100) + ) + + let listener = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup) + .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) + .bind(host: "127.0.0.1", port: 0) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + } + } + let port = listener.channel.localAddress!.port! + + async let serverDone: Void = listener.executeThenClose { inbound in + for try await child in inbound { + await server.handleTCP(channel: child) + break + } + } + + let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) + .connect(host: "127.0.0.1", port: port) { channel in + channel.eventLoop.makeCompletedFuture { + try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + } + } + + var receivedChunks = 0 + var connectionError: Error? + do { + try await client.executeThenClose { inbound, outbound in + try await Task.sleep(for: .milliseconds(300)) + for try await _ in inbound { + receivedChunks += 1 + } + } + } catch { + connectionError = error + } + + try? await listener.channel.close() + try? await serverDone + + #expect(receivedChunks == 0, "Expected server to drop connection due to idle timeout") + } } From 74d00c88272b9843284eeac11049ed0e2c7f842e Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:41:24 -0600 Subject: [PATCH 5/9] Test: remove unused connectionError variable in TCP tests --- Tests/DNSServerTests/TCPHandleTest.swift | 28 ++++++++++++++++-------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/Tests/DNSServerTests/TCPHandleTest.swift b/Tests/DNSServerTests/TCPHandleTest.swift index 279bb58d7..5434ff73f 100644 --- a/Tests/DNSServerTests/TCPHandleTest.swift +++ b/Tests/DNSServerTests/TCPHandleTest.swift @@ -241,7 +241,10 @@ struct TCPHandleTest { .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) .bind(host: "127.0.0.1", port: 0) { channel in channel.eventLoop.makeCompletedFuture { - try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self) + ) } } let port = listener.channel.localAddress!.port! @@ -256,12 +259,14 @@ struct TCPHandleTest { let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) .connect(host: "127.0.0.1", port: port) { channel in channel.eventLoop.makeCompletedFuture { - try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self) + ) } } var receivedChunks = 0 - var connectionError: Error? do { try await client.executeThenClose { inbound, outbound in var buf = ByteBuffer() @@ -276,7 +281,7 @@ struct TCPHandleTest { } } } catch { - connectionError = error + print("testTCPDropsOversizedFrame: connection closed with error (expected): \(error)") } try? await listener.channel.close() @@ -296,7 +301,10 @@ struct TCPHandleTest { .serverChannelOption(.socketOption(.so_reuseaddr), value: 1) .bind(host: "127.0.0.1", port: 0) { channel in channel.eventLoop.makeCompletedFuture { - try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self) + ) } } let port = listener.channel.localAddress!.port! @@ -311,12 +319,14 @@ struct TCPHandleTest { let client = try await ClientBootstrap(group: NIOSingletons.posixEventLoopGroup) .connect(host: "127.0.0.1", port: port) { channel in channel.eventLoop.makeCompletedFuture { - try NIOAsyncChannel(wrappingChannelSynchronously: channel, configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self)) + try NIOAsyncChannel( + wrappingChannelSynchronously: channel, + configuration: .init(inboundType: ByteBuffer.self, outboundType: ByteBuffer.self) + ) } } var receivedChunks = 0 - var connectionError: Error? do { try await client.executeThenClose { inbound, outbound in try await Task.sleep(for: .milliseconds(300)) @@ -325,7 +335,7 @@ struct TCPHandleTest { } } } catch { - connectionError = error + print("testTCPIdleTimeoutDropsConnection: connection closed with error (expected): \(error)") } try? await listener.channel.close() @@ -333,4 +343,4 @@ struct TCPHandleTest { #expect(receivedChunks == 0, "Expected server to drop connection due to idle timeout") } -} +} \ No newline at end of file From 35aad40f491d92c928eec5d1fe552a6670a53c57 Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:47:17 -0600 Subject: [PATCH 6/9] resolving the merge conflicts for apiserver+start --- .../Helpers/APIServer/APIServer+Start.swift | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index 019352d2e..cb00a7a66 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -115,8 +115,6 @@ extension APIServer { try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) } - // start up host table DNS (TCP) - group.addTask { log.info( "starting DNS resolver for container hostnames (TCP)", @@ -127,32 +125,6 @@ extension APIServer { ) try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) } - - // start up realhost DNS - /* - group.addTask { - let localhostResolver = LocalhostDNSHandler(log: log) - do { - try localhostResolver.monitorResolvers() - } catch { - log.error("could not initialize resolver monitor", metadata: ["error": "\(error)"]) - throw error - } - - let nxDomainResolver = NxDomainResolver() - let compositeResolver = CompositeResolver(handlers: [localhostResolver, nxDomainResolver]) - let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) - let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) - log.info( - "starting DNS resolver for localhost", - metadata: [ - "host": "\(Self.listenAddress)", - "port": "\(Self.localhostDNSPort)", - ] - ) - try await dnsServer.run(host: Self.listenAddress, port: Self.localhostDNSPort) - } - */ } } catch { log.error( From e9a90af7201befce64d946f9a5762b7c616b4f3a Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 10:58:01 -0600 Subject: [PATCH 7/9] Fix: correctly resolve APIServer+Start merge conflict --- .../Helpers/APIServer/APIServer+Start.swift | 56 +++++++++++++++++-- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index cb00a7a66..b9110f9a8 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -91,10 +91,15 @@ extension APIServer { $0[$1.key.rawValue] = $1.value }), log: log) - await withThrowingTaskGroup(of: Void.self) { group in + await withTaskGroup(of: Result.self) { group in group.addTask { log.info("starting XPC server") - try await server.listen() + do { + try await server.listen() + return .success(()) + } catch { + return .failure(error) + } } // start up host table DNS (UDP) @@ -112,7 +117,12 @@ extension APIServer { "port": "\(Self.dnsPort)", ] ) - try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) + do { + try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) + return .success(()) + } catch { + return .failure(error) + } } group.addTask { @@ -123,7 +133,45 @@ extension APIServer { "port": "\(Self.dnsPort)", ] ) - try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) + do { + try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) + return .success(()) + } catch { + return .failure(error) + } + } + + // start up realhost DNS + group.addTask { + do { + let localhostResolver = LocalhostDNSHandler(log: log) + try localhostResolver.monitorResolvers() + + let nxDomainResolver = NxDomainResolver() + let compositeResolver = CompositeResolver(handlers: [localhostResolver, nxDomainResolver]) + let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) + let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) + log.info( + "starting DNS resolver for localhost", + metadata: [ + "host": "\(Self.listenAddress)", + "port": "\(Self.localhostDNSPort)", + ] + ) + try await dnsServer.run(host: Self.listenAddress, port: Self.localhostDNSPort) + return .success(()) + } catch { + return .failure(error) + } + } + + for await result in group { + switch result { + case .success(): + continue + case .failure(let error): + log.error("API server task failed: \(error)") + } } } } catch { From 47c831a4be21187269b03f9a00c67c611560d74a Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Thu, 26 Feb 2026 21:12:32 -0600 Subject: [PATCH 8/9] Fix: strip trailing DNS dot for external lookups and remove invalid try Co-authored-by: renish avaiya --- Sources/DNSServer/DNSServer+TCPHandle.swift | 10 +-- .../Helpers/APIServer/APIServer+Start.swift | 68 ++++++------------- .../APIServer/ContainerDNSHandler.swift | 20 ++++-- Tests/DNSServerTests/TCPHandleTest.swift | 6 +- 4 files changed, 40 insertions(+), 64 deletions(-) diff --git a/Sources/DNSServer/DNSServer+TCPHandle.swift b/Sources/DNSServer/DNSServer+TCPHandle.swift index 8d36c49cc..494992ec2 100644 --- a/Sources/DNSServer/DNSServer+TCPHandle.swift +++ b/Sources/DNSServer/DNSServer+TCPHandle.swift @@ -1,5 +1,5 @@ //===----------------------------------------------------------------------===// -// Copyright © 2025-2026 Apple Inc. and the container project authors. +// Copyright © 2026 Apple Inc. and the container project authors. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -57,9 +57,11 @@ extension DNSServer { buffer.writeImmutableBuffer(chunk) while buffer.readableBytes >= 2 { - guard let msgLen = buffer.getInteger( - at: buffer.readerIndex, as: UInt16.self - ) else { break } + guard + let msgLen = buffer.getInteger( + at: buffer.readerIndex, as: UInt16.self + ) + else { break } guard msgLen > 0, msgLen <= Self.maxTCPMessageSize else { self.log?.error( diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index b9110f9a8..486ef168d 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -91,18 +91,13 @@ extension APIServer { $0[$1.key.rawValue] = $1.value }), log: log) - await withTaskGroup(of: Result.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { log.info("starting XPC server") - do { - try await server.listen() - return .success(()) - } catch { - return .failure(error) - } + try await server.listen() } - // start up host table DNS (UDP) + // start up host table DNS (UDP and TCP) let hostsResolver = ContainerDNSHandler(networkService: networkService) let nxDomainResolver = NxDomainResolver() let compositeResolver = CompositeResolver(handlers: [hostsResolver, nxDomainResolver]) @@ -117,12 +112,7 @@ extension APIServer { "port": "\(Self.dnsPort)", ] ) - do { - try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) - return .success(()) - } catch { - return .failure(error) - } + try await dnsServer.run(host: Self.listenAddress, port: Self.dnsPort) } group.addTask { @@ -133,45 +123,25 @@ extension APIServer { "port": "\(Self.dnsPort)", ] ) - do { - try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) - return .success(()) - } catch { - return .failure(error) - } + try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) } - // start up realhost DNS group.addTask { - do { - let localhostResolver = LocalhostDNSHandler(log: log) - try localhostResolver.monitorResolvers() - - let nxDomainResolver = NxDomainResolver() - let compositeResolver = CompositeResolver(handlers: [localhostResolver, nxDomainResolver]) - let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) - let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) - log.info( - "starting DNS resolver for localhost", - metadata: [ - "host": "\(Self.listenAddress)", - "port": "\(Self.localhostDNSPort)", - ] - ) - try await dnsServer.run(host: Self.listenAddress, port: Self.localhostDNSPort) - return .success(()) - } catch { - return .failure(error) - } - } + let localhostResolver = LocalhostDNSHandler(log: log) + await localhostResolver.monitorResolvers() - for await result in group { - switch result { - case .success(): - continue - case .failure(let error): - log.error("API server task failed: \(error)") - } + let nxDomainResolver = NxDomainResolver() + let compositeResolver = CompositeResolver(handlers: [localhostResolver, nxDomainResolver]) + let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) + let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) + log.info( + "starting DNS resolver for localhost", + metadata: [ + "host": "\(Self.listenAddress)", + "port": "\(Self.localhostDNSPort)", + ] + ) + try await dnsServer.run(host: Self.listenAddress, port: Self.localhostDNSPort) } } } catch { diff --git a/Sources/Helpers/APIServer/ContainerDNSHandler.swift b/Sources/Helpers/APIServer/ContainerDNSHandler.swift index 4fcf8a2c4..a764162f0 100644 --- a/Sources/Helpers/APIServer/ContainerDNSHandler.swift +++ b/Sources/Helpers/APIServer/ContainerDNSHandler.swift @@ -18,7 +18,6 @@ import ContainerAPIService import DNS import DNSServer -/// Handler that uses table lookup to resolve hostnames. struct ContainerDNSHandler: DNSHandler { private let networkService: NetworksService private let ttl: UInt32 @@ -37,10 +36,6 @@ struct ContainerDNSHandler: DNSHandler { case ResourceRecordType.host6: let result = try await answerHost6(question: question) if result.record == nil && result.hostnameExists { - // Return NODATA (noError with empty answers) when hostname exists but has no IPv6. - // This is required because musl libc has issues when A record exists but AAAA returns NXDOMAIN. - // musl treats NXDOMAIN on AAAA as "domain doesn't exist" and fails DNS resolution entirely. - // NODATA correctly indicates "no IPv6 address available, but domain exists". return Message( id: query.id, type: .response, @@ -91,9 +86,16 @@ struct ContainerDNSHandler: DNSHandler { } private func answerHost(question: Question) async throws -> ResourceRecord? { - guard let ipAllocation = try await networkService.lookup(hostname: question.name) else { + var hostname = question.name + if hostname.hasSuffix(".") { + hostname.removeLast() + } + print("DEBUG: ContainerDNSHandler looking up hostname: '\(hostname)'") + guard let ipAllocation = try await networkService.lookup(hostname: hostname) else { + print("DEBUG: ContainerDNSHandler lookup failed for '\(hostname)'") return nil } + print("DEBUG: ContainerDNSHandler found IP: \(ipAllocation.ipv4Address.address.description)") let ipv4 = ipAllocation.ipv4Address.address.description guard let ip = IPv4(ipv4) else { throw DNSResolverError.serverError("failed to parse IP address: \(ipv4)") @@ -103,7 +105,11 @@ struct ContainerDNSHandler: DNSHandler { } private func answerHost6(question: Question) async throws -> (record: ResourceRecord?, hostnameExists: Bool) { - guard let ipAllocation = try await networkService.lookup(hostname: question.name) else { + var hostname = question.name + if hostname.hasSuffix(".") { + hostname.removeLast() + } + guard let ipAllocation = try await networkService.lookup(hostname: hostname) else { return (nil, false) } guard let ipv6Address = ipAllocation.ipv6Address else { diff --git a/Tests/DNSServerTests/TCPHandleTest.swift b/Tests/DNSServerTests/TCPHandleTest.swift index 5434ff73f..61b793af1 100644 --- a/Tests/DNSServerTests/TCPHandleTest.swift +++ b/Tests/DNSServerTests/TCPHandleTest.swift @@ -1,5 +1,5 @@ //===----------------------------------------------------------------------===// -// Copyright © 2025-2026 Apple Inc. and the container project authors. +// Copyright © 2026 Apple Inc. and the container project authors. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -22,7 +22,6 @@ import Testing @testable import DNSServer - struct ProcessRawTest { @Test func testProcessRawReturnsSerializedResponse() async throws { let handler = HostTableResolver(hosts4: ["myhost.": IPv4("10.0.0.1")!]) @@ -63,7 +62,6 @@ struct ProcessRawTest { } } - struct TCPHandleTest { private func makeTCPFrame(hostname: String, id: UInt16 = 1) throws -> ByteBuffer { let bytes = try Message( @@ -343,4 +341,4 @@ struct TCPHandleTest { #expect(receivedChunks == 0, "Expected server to drop connection due to idle timeout") } -} \ No newline at end of file +} From eddc3ed91a8ed9f4ed1a9c67e89a4dbcc770023c Mon Sep 17 00:00:00 2001 From: Ronit Sabhaya Date: Fri, 27 Feb 2026 12:15:54 -0600 Subject: [PATCH 9/9] added logs to see the result for empty return --- Sources/DNSServer/DNSServer+TCPHandle.swift | 10 ++++++++++ .../Helpers/APIServer/APIServer+Start.swift | 18 ++---------------- .../APIServer/ContainerDNSHandler.swift | 3 --- 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/Sources/DNSServer/DNSServer+TCPHandle.swift b/Sources/DNSServer/DNSServer+TCPHandle.swift index d15c74761..5a10c98be 100644 --- a/Sources/DNSServer/DNSServer+TCPHandle.swift +++ b/Sources/DNSServer/DNSServer+TCPHandle.swift @@ -19,6 +19,8 @@ import NIOCore import NIOPosix import Synchronization +/// Tracks the activity timestamps for a specific TCP connection. +/// Used to enforce idle timeouts and disconnect inactive clients to prevent resource exhaustion. private actor ConnectionActivity { private var lastSeen: ContinuousClock.Instant @@ -36,6 +38,14 @@ private actor ConnectionActivity { } extension DNSServer { + /// Handles a single active TCP connection from an inbound client. + /// + /// This method manages the lifecycle of the connection, reading length-prefixed DNS queries + /// iteratively and executing the underlying `processRaw` logic for each query concurrently + /// using Swift Concurrency. It enforces strict idle timeouts to prevent stale clients + /// from holding connections open indefinitely. + /// + /// - Parameter channel: The connected asynchronous TCP channel containing the message buffer streams. func handleTCP(channel: NIOAsyncChannel) async { do { try await channel.executeThenClose { inbound, outbound in diff --git a/Sources/Helpers/APIServer/APIServer+Start.swift b/Sources/Helpers/APIServer/APIServer+Start.swift index 0c01eb0f5..3c71d5d07 100644 --- a/Sources/Helpers/APIServer/APIServer+Start.swift +++ b/Sources/Helpers/APIServer/APIServer+Start.swift @@ -126,22 +126,8 @@ extension APIServer { try await dnsServer.runTCP(host: Self.listenAddress, port: Self.dnsPort) } - group.addTask { - let localhostResolver = LocalhostDNSHandler(log: log) - await localhostResolver.monitorResolvers() - - let nxDomainResolver = NxDomainResolver() - let compositeResolver = CompositeResolver(handlers: [localhostResolver, nxDomainResolver]) - let hostsQueryValidator = StandardQueryValidator(handler: compositeResolver) - let dnsServer: DNSServer = DNSServer(handler: hostsQueryValidator, log: log) - log.info( - "starting DNS resolver for localhost", - metadata: [ - "host": "\(Self.listenAddress)", - "port": "\(Self.localhostDNSPort)", - ] - ) - try await dnsServer.run(host: Self.listenAddress, port: Self.localhostDNSPort) + for try await _ in group { + continue } } } catch { diff --git a/Sources/Helpers/APIServer/ContainerDNSHandler.swift b/Sources/Helpers/APIServer/ContainerDNSHandler.swift index a764162f0..d63df836c 100644 --- a/Sources/Helpers/APIServer/ContainerDNSHandler.swift +++ b/Sources/Helpers/APIServer/ContainerDNSHandler.swift @@ -90,12 +90,9 @@ struct ContainerDNSHandler: DNSHandler { if hostname.hasSuffix(".") { hostname.removeLast() } - print("DEBUG: ContainerDNSHandler looking up hostname: '\(hostname)'") guard let ipAllocation = try await networkService.lookup(hostname: hostname) else { - print("DEBUG: ContainerDNSHandler lookup failed for '\(hostname)'") return nil } - print("DEBUG: ContainerDNSHandler found IP: \(ipAllocation.ipv4Address.address.description)") let ipv4 = ipAllocation.ipv4Address.address.description guard let ip = IPv4(ipv4) else { throw DNSResolverError.serverError("failed to parse IP address: \(ipv4)")