From b7079b75cd604bc675af080edef56b486decb155 Mon Sep 17 00:00:00 2001
From: Andre Dieb <andre.dieb@gmail.com>
Date: Wed, 21 Feb 2024 18:51:52 -0300
Subject: [PATCH] Fix continuation memory leak in Ares.query (#31)

* Fix continuation memory leak in Ares.query

* Use class for QueryReplyHandler

* Deallocate QueryReplyHandler for DNSSD

* Move defer block after allocate/initialize, use class

Move defer deallocation block to after initialization.

Use class instead of struct for DNSSD.QueryReplyHandler.
---
 .../c-ares/DNSResolver_c-ares.swift           | 15 +++++++-------
 .../dnssd/DNSResolver_dnssd.swift             | 20 ++++++++++---------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/Sources/AsyncDNSResolver/c-ares/DNSResolver_c-ares.swift b/Sources/AsyncDNSResolver/c-ares/DNSResolver_c-ares.swift
index 7953969..7babb5c 100644
--- a/Sources/AsyncDNSResolver/c-ares/DNSResolver_c-ares.swift
+++ b/Sources/AsyncDNSResolver/c-ares/DNSResolver_c-ares.swift
@@ -160,8 +160,12 @@ class Ares {
                             preconditionFailure("'arg' is nil. This is a bug.")
                         }
 
-                        let handler = QueryReplyHandler(pointer: handlerPointer)
-                        defer { handlerPointer.deallocate() }
+                        let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
+                        let handler = pointer.pointee
+                        defer {
+                            pointer.deinitialize(count: 1)
+                            pointer.deallocate()
+                        }
 
                         handler.handle(status: status, buffer: buf, length: len)
                     }
@@ -258,7 +262,7 @@ extension Ares {
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
 extension Ares {
-    struct QueryReplyHandler {
+    class QueryReplyHandler {
         private let _handler: (CInt, UnsafeMutablePointer<CUnsignedChar>?, CInt) -> Void
 
         init<Parser: AresQueryReplyParser>(parser: Parser, _ continuation: CheckedContinuation<Parser.Reply, Error>) {
@@ -276,11 +280,6 @@ extension Ares {
             }
         }
 
-        init(pointer: UnsafeMutableRawPointer) {
-            let handlerPointer = pointer.assumingMemoryBound(to: Self.self)
-            self = handlerPointer.pointee
-        }
-
         func handle(status: CInt, buffer: UnsafeMutablePointer<CUnsignedChar>?, length: CInt) {
             self._handler(status, buffer, length)
         }
diff --git a/Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift b/Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift
index 80f8849..c9a6575 100644
--- a/Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift
+++ b/Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift
@@ -115,18 +115,25 @@ struct DNSSD {
                 byteCount: MemoryLayout<QueryReplyHandler>.stride,
                 alignment: MemoryLayout<QueryReplyHandler>.alignment
             )
-            // The handler might be called multiple times so don't deallocate inside `callback`
-            defer { handlerPointer.deallocate() }
 
             handlerPointer.initializeMemory(as: QueryReplyHandler.self, repeating: handler, count: 1)
 
+            // The handler might be called multiple times so don't deallocate inside `callback`
+            defer {
+                let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
+                pointer.deinitialize(count: 1)
+                pointer.deallocate()
+            }
+
             // This is called once per record received
             let callback: DNSServiceQueryRecordReply = { _, _, _, errorCode, _, _, _, rdlen, rdata, _, context in
                 guard let handlerPointer = context else {
                     preconditionFailure("'context' is nil. This is a bug.")
                 }
 
-                let handler = QueryReplyHandler(pointer: handlerPointer)
+                let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
+                let handler = pointer.pointee
+
                 // This parses a record then adds it to the stream
                 handler.handleRecord(errorCode: errorCode, data: rdata, length: rdlen)
             }
@@ -171,7 +178,7 @@ struct DNSSD {
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
 extension DNSSD {
-    struct QueryReplyHandler {
+    class QueryReplyHandler {
         private let _handleRecord: (DNSServiceErrorType, UnsafeRawPointer?, UInt16) -> Void
 
         init<Handler: DNSSDQueryReplyHandler>(handler: Handler, _ continuation: AsyncThrowingStream<Handler.Record, Error>.Continuation) {
@@ -189,11 +196,6 @@ extension DNSSD {
             }
         }
 
-        init(pointer: UnsafeMutableRawPointer) {
-            let handlerPointer = pointer.assumingMemoryBound(to: Self.self)
-            self = handlerPointer.pointee
-        }
-
         func handleRecord(errorCode: DNSServiceErrorType, data: UnsafeRawPointer?, length: UInt16) {
             self._handleRecord(errorCode, data, length)
         }