From fb0b879e7450c85472abed166bc4df5575ed4b14 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Fri, 3 Jul 2026 10:46:58 +0530 Subject: [PATCH] Use loadUnaligned/storeBytes for the DNS wire codec. The DNS wire codec writes and reads multi-byte integers (UInt16 type and class, UInt32 ttl) at offsets that follow a variable-length domain name, so they routinely land on unaligned addresses. copyIn/copyOut accessed that memory via advanced(by:).assumingMemoryBound(to:).pointee, which assumes the address is correctly aligned for T and treats the UInt8-bound buffer as if it were bound to T. Both assumptions are invalid here, so the accesses are undefined behavior even though arm64 happens to tolerate them. Switch the write path to UnsafeMutableRawBufferPointer.storeBytes and the read path to UnsafeRawBufferPointer.loadUnaligned, which make no alignment or type-binding assumptions for trivial types. Constrain the generic parameter to BitwiseCopyable so the trivial-type precondition is enforced at compile time; the only callers pass UInt8, UInt16, and UInt32. Add tests that round-trip UInt16 and UInt32 through copyIn/copyOut at every byte offset within a machine word, asserting correct big-endian layout, plus a Message serialize/deserialize round-trip whose fields land on odd offsets. --- Sources/DNSServer/Records/UInt8+Binding.swift | 11 +-- Tests/DNSServerTests/RecordsTests.swift | 91 +++++++++++++++++++ 2 files changed, 95 insertions(+), 7 deletions(-) diff --git a/Sources/DNSServer/Records/UInt8+Binding.swift b/Sources/DNSServer/Records/UInt8+Binding.swift index cd9b91b1a..6d2e78a12 100644 --- a/Sources/DNSServer/Records/UInt8+Binding.swift +++ b/Sources/DNSServer/Records/UInt8+Binding.swift @@ -24,29 +24,26 @@ import Foundation extension [UInt8] { /// Copy a value into the buffer at the given offset. /// - Returns: The new offset after writing, or nil if the buffer is too small. - package mutating func copyIn(as type: T.Type, value: T, offset: Int = 0) -> Int? { + package mutating func copyIn(as type: T.Type, value: T, offset: Int = 0) -> Int? { let size = MemoryLayout.size guard self.count >= size + offset else { return nil } return self.withUnsafeMutableBytes { - $0.baseAddress?.advanced(by: offset).assumingMemoryBound(to: T.self).pointee = value + $0.storeBytes(of: value, toByteOffset: offset, as: T.self) return offset + size } } /// Copy a value out of the buffer at the given offset. /// - Returns: A tuple of (new offset, value), or nil if the buffer is too small. - package func copyOut(as type: T.Type, offset: Int = 0) -> (Int, T)? { + package func copyOut(as type: T.Type, offset: Int = 0) -> (Int, T)? { let size = MemoryLayout.size guard self.count >= size + offset else { return nil } return self.withUnsafeBytes { - guard let value = $0.baseAddress?.advanced(by: offset).assumingMemoryBound(to: T.self).pointee else { - return nil - } - return (offset + size, value) + (offset + size, $0.loadUnaligned(fromByteOffset: offset, as: T.self)) } } diff --git a/Tests/DNSServerTests/RecordsTests.swift b/Tests/DNSServerTests/RecordsTests.swift index d1fc0cc8d..e09e4eb3f 100644 --- a/Tests/DNSServerTests/RecordsTests.swift +++ b/Tests/DNSServerTests/RecordsTests.swift @@ -23,6 +23,97 @@ import Testing @Suite("DNS Records Tests") struct RecordsTests { + // MARK: - Binding Tests + + /// The DNS wire codec writes multi-byte integers (`UInt16` type/class, + /// `UInt32` ttl) at variable offsets that follow a variable-length name, so + /// they routinely land on unaligned addresses. These tests exercise + /// `copyIn`/`copyOut` at every offset in a machine word to ensure the codec + /// reads and writes correct big-endian bytes regardless of alignment. + @Suite("Binding") + struct BindingTests { + @Test("UInt16 round-trips at every offset") + func uint16RoundTripEveryOffset() throws { + let value: UInt16 = 0x1234 + for offset in 0..<8 { + var buffer = [UInt8](repeating: 0, count: offset + MemoryLayout.size) + + let writeEnd = buffer.copyIn(as: UInt16.self, value: value.bigEndian, offset: offset) + #expect(writeEnd == offset + 2) + // Stored as big-endian bytes at the requested offset. + #expect(buffer[offset] == 0x12) + #expect(buffer[offset + 1] == 0x34) + + let (readEnd, raw) = try #require(buffer.copyOut(as: UInt16.self, offset: offset)) + #expect(readEnd == offset + 2) + #expect(UInt16(bigEndian: raw) == value) + } + } + + @Test("UInt32 round-trips at every offset") + func uint32RoundTripEveryOffset() throws { + let value: UInt32 = 0x1234_5678 + for offset in 0..<8 { + var buffer = [UInt8](repeating: 0, count: offset + MemoryLayout.size) + + let writeEnd = buffer.copyIn(as: UInt32.self, value: value.bigEndian, offset: offset) + #expect(writeEnd == offset + 4) + // Stored as big-endian bytes at the requested offset. + #expect(buffer[offset] == 0x12) + #expect(buffer[offset + 1] == 0x34) + #expect(buffer[offset + 2] == 0x56) + #expect(buffer[offset + 3] == 0x78) + + let (readEnd, raw) = try #require(buffer.copyOut(as: UInt32.self, offset: offset)) + #expect(readEnd == offset + 4) + #expect(UInt32(bigEndian: raw) == value) + } + } + + @Test("copyIn returns nil when the buffer is too small") + func copyInBufferTooSmall() { + var buffer = [UInt8](repeating: 0, count: 1) + #expect(buffer.copyIn(as: UInt16.self, value: 0, offset: 0) == nil) + #expect(buffer.copyIn(as: UInt16.self, value: 0, offset: 1) == nil) + } + + @Test("copyOut returns nil when the buffer is too small") + func copyOutBufferTooSmall() { + let buffer = [UInt8](repeating: 0, count: 1) + #expect(buffer.copyOut(as: UInt16.self, offset: 0) == nil) + #expect(buffer.copyOut(as: UInt16.self, offset: 1) == nil) + } + + @Test("Message round-trips with fields at unaligned offsets") + func messageRoundTripUnaligned() throws { + // "example.com." encodes to [7]example[3]com[0] = 13 bytes. With the + // 12-byte header, the question's type/class UInt16 fields start at + // offset 25 (odd), so serialize and deserialize both hit the + // unaligned path. + let original = Message( + id: 0xABCD, + type: .query, + recursionDesired: true, + questions: [Question(name: "example.com.", type: .host6)] + ) + + let data = try original.serialize() + let bytes = Array(data) + + // type (AAAA = 28) and class (IN = 1) are big-endian at odd offsets. + #expect(bytes[25] == 0x00) + #expect(bytes[26] == 28) + #expect(bytes[27] == 0x00) + #expect(bytes[28] == 1) + + let parsed = try Message(deserialize: data) + #expect(parsed.id == 0xABCD) + #expect(parsed.questions.count == 1) + #expect(parsed.questions[0].type == .host6) + #expect(parsed.questions[0].recordClass == .internet) + } + } + // MARK: - DNSName Tests @Suite("DNSName")