From 16f19c0c645014a41456097f1848018efca41d7e Mon Sep 17 00:00:00 2001
From: Cory Benfield <lukasa@apple.com>
Date: Tue, 26 Nov 2024 16:19:33 +0000
Subject: [PATCH] Prevent crash in Happy Eyeballs Resolver (#3003)

Motivation:

In vanishingly rare situations it is possible for the AAAA results to
come in on the same tick as the resolution delay timer completes. In
those cases, depending on the ordering of the tasks, we can get
situations where the resolution delay timer completion causes a crash.

Modifications:

Tolerate receiving the resolution delay timer after resolution
completes.

Result:

Fewer crashes
---
 Sources/NIOPosix/HappyEyeballs.swift        |  7 ++-
 Tests/NIOPosixTests/HappyEyeballsTest.swift | 47 +++++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/Sources/NIOPosix/HappyEyeballs.swift b/Sources/NIOPosix/HappyEyeballs.swift
index 015a8b4af6..25071ec32a 100644
--- a/Sources/NIOPosix/HappyEyeballs.swift
+++ b/Sources/NIOPosix/HappyEyeballs.swift
@@ -464,7 +464,12 @@ internal final class HappyEyeballsConnector<ChannelBuilderResult> {
         // notifications, and can also get late scheduled task callbacks. We want to just quietly
         // ignore these, as our transition into the complete state should have already sent
         // cleanup messages to all of these things.
-        case (.complete, .resolverACompleted),
+        //
+        // We can also get the resolutionDelayElapsed after allResolved, as it's possible that
+        // callback was already dequeued in the same tick as the cancellation. That's also fine:
+        // the resolution delay isn't interesting.
+        case (.allResolved, .resolutionDelayElapsed),
+            (.complete, .resolverACompleted),
             (.complete, .resolverAAAACompleted),
             (.complete, .connectSuccess),
             (.complete, .connectFailed),
diff --git a/Tests/NIOPosixTests/HappyEyeballsTest.swift b/Tests/NIOPosixTests/HappyEyeballsTest.swift
index 66f0dec09d..29f35820f8 100644
--- a/Tests/NIOPosixTests/HappyEyeballsTest.swift
+++ b/Tests/NIOPosixTests/HappyEyeballsTest.swift
@@ -1308,4 +1308,51 @@ public final class HappyEyeballsTest: XCTestCase {
 
         XCTAssertNoThrow(try client.close().wait())
     }
+
+    func testResolutionTimeoutAndResolutionInSameTick() throws {
+        var channels: [Channel] = []
+        let (eyeballer, resolver, loop) = buildEyeballer(host: "example.com", port: 80) {
+            let channelFuture = defaultChannelBuilder(loop: $0, family: $1)
+            channelFuture.whenSuccess { channel in
+                try! channel.pipeline.addHandler(ConnectionDelayer(), name: CONNECT_DELAYER, position: .first).wait()
+                channels.append(channel)
+            }
+            return channelFuture
+        }
+        let targetFuture = eyeballer.resolveAndConnect().flatMapThrowing { (channel) -> String? in
+            let target = channel.connectTarget()
+            _ = try (channel as! EmbeddedChannel).finish()
+            return target
+        }
+        loop.run()
+
+        // Then, queue a task to resolve the v6 promise after 50ms.
+        // Why 50ms? This is the same time as the resolution delay.
+        let promise = resolver.v6Promise
+        loop.scheduleTask(in: .milliseconds(50)) {
+            promise.fail(DummyError())
+        }
+
+        // Kick off the IPv4 resolution. This triggers the timer for the resolution delay.
+        resolver.v4Promise.succeed(SINGLE_IPv4_RESULT)
+        loop.run()
+
+        // Advance time 50ms.
+        loop.advanceTime(by: .milliseconds(50))
+
+        // Then complete the connection future.
+        XCTAssertEqual(channels.count, 1)
+        channels.first!.succeedConnection()
+
+        // Should be done.
+        let target = try targetFuture.wait()
+        XCTAssertEqual(target!, "10.0.0.1")
+
+        // We should have had queries for AAAA and A.
+        let expectedQueries: [DummyResolver.Event] = [
+            .aaaa(host: "example.com", port: 80),
+            .a(host: "example.com", port: 80),
+        ]
+        XCTAssertEqual(resolver.events, expectedQueries)
+    }
 }