From 6f8af37568d474dc36d7f6ec6bb09433cde64c74 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Tue, 8 Mar 2022 21:58:49 -0500
Subject: [PATCH] fix: a concurrency issue in fakedns (#1666)

In rare cases different domains asking for dns will return the same IP, Add a mutex.
---
 app/dns/fakedns/fake.go         |  8 +++++++-
 app/dns/fakedns/fakedns_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/app/dns/fakedns/fake.go b/app/dns/fakedns/fake.go
index 124de4abc..b5ca89670 100644
--- a/app/dns/fakedns/fake.go
+++ b/app/dns/fakedns/fake.go
@@ -8,6 +8,7 @@ import (
 	"math"
 	"math/big"
 	gonet "net"
+	"sync"
 
 	"github.com/v2fly/v2ray-core/v5/common"
 	"github.com/v2fly/v2ray-core/v5/common/cache"
@@ -18,6 +19,7 @@ import (
 type Holder struct {
 	domainToIP cache.Lru
 	nextIP     *big.Int
+	mu         *sync.Mutex
 
 	ipRange *gonet.IPNet
 
@@ -54,6 +56,7 @@ func (fkdns *Holder) Close() error {
 	fkdns.domainToIP = nil
 	fkdns.nextIP = nil
 	fkdns.ipRange = nil
+	fkdns.mu = nil
 	return nil
 }
 
@@ -72,7 +75,7 @@ func NewFakeDNSHolder() (*Holder, error) {
 }
 
 func NewFakeDNSHolderConfigOnly(conf *FakeDnsPool) (*Holder, error) {
-	return &Holder{nil, nil, nil, conf}, nil
+	return &Holder{nil, nil, nil, nil, conf}, nil
 }
 
 func (fkdns *Holder) initializeFromConfig() error {
@@ -102,11 +105,14 @@ func (fkdns *Holder) initialize(ipPoolCidr string, lruSize int) error {
 	fkdns.domainToIP = cache.NewLru(lruSize)
 	fkdns.ipRange = ipRange
 	fkdns.nextIP = currentIP
+	fkdns.mu = new(sync.Mutex)
 	return nil
 }
 
 // GetFakeIPForDomain checks and generate a fake IP for a domain name
 func (fkdns *Holder) GetFakeIPForDomain(domain string) []net.Address {
+	fkdns.mu.Lock()
+	defer fkdns.mu.Unlock()
 	if v, ok := fkdns.domainToIP.Get(domain); ok {
 		return []net.Address{v.(net.Address)}
 	}
diff --git a/app/dns/fakedns/fakedns_test.go b/app/dns/fakedns/fakedns_test.go
index 48defb83f..7b37f4a84 100644
--- a/app/dns/fakedns/fakedns_test.go
+++ b/app/dns/fakedns/fakedns_test.go
@@ -2,10 +2,13 @@ package fakedns
 
 import (
 	gonet "net"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 
+	"golang.org/x/sync/errgroup"
+
 	"github.com/v2fly/v2ray-core/v5/common"
 	"github.com/v2fly/v2ray-core/v5/common/net"
 	"github.com/v2fly/v2ray-core/v5/common/uuid"
@@ -71,6 +74,31 @@ func TestFakeDnsHolderCreateMappingManySingleDomain(t *testing.T) {
 	assert.Equal(t, "198.18.0.0", addr2[0].IP().String())
 }
 
+func TestGetFakeIPForDomainConcurrently(t *testing.T) {
+	fkdns, err := NewFakeDNSHolder()
+	common.Must(err)
+
+	total := 200
+	addr := make([][]net.Address, total+1)
+	var errg errgroup.Group
+	for i := 0; i < total; i++ {
+		errg.Go(testGetFakeIP(i, addr, fkdns))
+	}
+	errg.Wait()
+	for i := 0; i < total; i++ {
+		for j := i + 1; j < total; j++ {
+			assert.NotEqual(t, addr[i][0].IP().String(), addr[j][0].IP().String())
+		}
+	}
+}
+
+func testGetFakeIP(index int, addr [][]net.Address, fkdns *Holder) func() error {
+	return func() error {
+		addr[index] = fkdns.GetFakeIPForDomain("fakednstest" + strconv.Itoa(index) + ".example.com")
+		return nil
+	}
+}
+
 func TestFakeDnsHolderCreateMappingAndRollOver(t *testing.T) {
 	fkdns, err := NewFakeDNSHolderConfigOnly(&FakeDnsPool{
 		IpPool:  "240.0.0.0/12",