From c5095b12643de60afaa3eb0aff7a541605846a73 Mon Sep 17 00:00:00 2001
From: Rod Hynes <rod-hynes@users.noreply.github.com>
Date: Thu, 25 Aug 2022 10:55:33 -0400
Subject: [PATCH] Fix DNS resolver nil pointer dereference

---
 psiphon/common/resolver/resolver.go      |  2 +-
 psiphon/common/resolver/resolver_test.go | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/psiphon/common/resolver/resolver.go b/psiphon/common/resolver/resolver.go
index 8b415fe65..60b1d339e 100644
--- a/psiphon/common/resolver/resolver.go
+++ b/psiphon/common/resolver/resolver.go
@@ -865,7 +865,7 @@ func (r *Resolver) ResolveIP(
 	if result != nil &&
 		resolveCtx.Err() == nil &&
 		atomic.LoadInt64(&inFlight) > 0 &&
-		(atomic.LoadInt32(&awaitA) != 0 || atomic.LoadInt32(&awaitAAAA) != 0) ||
+		(atomic.LoadInt32(&awaitA) != 0 || atomic.LoadInt32(&awaitAAAA) != 0) &&
 		params.AwaitTimeout > 0 {
 
 		resetTimer(params.AwaitTimeout)
diff --git a/psiphon/common/resolver/resolver_test.go b/psiphon/common/resolver/resolver_test.go
index a3645e6bc..9ff123a57 100644
--- a/psiphon/common/resolver/resolver_test.go
+++ b/psiphon/common/resolver/resolver_test.go
@@ -625,6 +625,27 @@ func runTestResolver() error {
 		return errors.TraceNew("unexpected success")
 	}
 
+	// Test: cancel context while resolving
+
+	// This test exercises the additional answers and await cases in
+	// ResolveIP. The test is timing dependent, and so imperfect, but this
+	// configuration can reproduce panics in those cases before bugs were
+	// fixed, where DNS responses need to be received just as the context is
+	// cancelled.
+
+	networkConfig.GetDNSServers = func() []string { return []string{okServer.getAddr()} }
+	networkID = "networkID-6"
+
+	for i := 0; i < 100; i++ {
+		resolver.cache.Flush()
+
+		ctx, cancelFunc := context.WithTimeout(
+			context.Background(), time.Duration(i%10*20)*time.Microsecond)
+		defer cancelFunc()
+
+		_, _ = resolver.ResolveIP(ctx, networkID, params, exampleDomain)
+	}
+
 	return nil
 }