From 5ecdafa2de7529d9c58ed20f500d610a83ad3d34 Mon Sep 17 00:00:00 2001
From: WendelHime <6754291+WendelHime@users.noreply.github.com>
Date: Fri, 26 Jul 2024 12:11:38 -0300
Subject: [PATCH] fix: hashing IP addresses and setting SNI to masquerades

---
 direct.go      | 12 +++---------
 direct_test.go |  2 +-
 masquerade.go  | 12 +++++++++++-
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/direct.go b/direct.go
index 0eac8f2..4357bec 100644
--- a/direct.go
+++ b/direct.go
@@ -421,17 +421,11 @@ func (d *direct) dialServerWith(m *Masquerade) (net.Conn, error) {
 	addr := m.IpAddress
 	var sendServerNameExtension bool
 
-	// looking for provider and using SNI if enabled
-	provider := d.findProviderFromMasquerade(m)
-	if provider != nil && provider.SNIConfig != nil && provider.SNIConfig.UseArbitrarySNIs {
+	if m.SNI != "" {
 		sendServerNameExtension = true
 
-		// selecting a random SNI
-		randomSNIIndex := rand.IntN(len(provider.SNIConfig.ArbitrarySNIs))
-		sniDomain := provider.SNIConfig.ArbitrarySNIs[randomSNIIndex]
-
-		op.Set("arbitrary_sni", sniDomain)
-		tlsConfig.ServerName = sniDomain
+		op.Set("arbitrary_sni", m.SNI)
+		tlsConfig.ServerName = m.SNI
 		tlsConfig.InsecureSkipVerify = true
 		tlsConfig.VerifyPeerCertificate = func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
 			log.Tracef("verifying peer certificate for masquerade domain %s", m.Domain)
diff --git a/direct_test.go b/direct_test.go
index 657d9c7..619307a 100644
--- a/direct_test.go
+++ b/direct_test.go
@@ -145,7 +145,7 @@ func TestLoadCandidates(t *testing.T) {
 	actual := make(map[Masquerade]bool)
 	count := 0
 	for _, m := range d.masquerades {
-		actual[Masquerade{m.Domain, m.IpAddress}] = true
+		actual[Masquerade{Domain: m.Domain, IpAddress: m.IpAddress}] = true
 		count++
 	}
 
diff --git a/masquerade.go b/masquerade.go
index 98cc541..2dceda6 100644
--- a/masquerade.go
+++ b/masquerade.go
@@ -2,6 +2,7 @@ package fronted
 
 import (
 	"fmt"
+	"hash/crc32"
 	"net"
 	"net/http"
 	"sort"
@@ -33,6 +34,9 @@ type Masquerade struct {
 	// IpAddress: pre-resolved ip address to use instead of Domain (if
 	// available)
 	IpAddress string
+
+	// SNI: the SNI to use for this masquerade
+	SNI string
 }
 
 type masquerade struct {
@@ -109,8 +113,14 @@ func NewProvider(hosts map[string]string, testURL string, masquerades []*Masquer
 	for k, v := range hosts {
 		d.HostAliases[strings.ToLower(k)] = v
 	}
+
 	for _, m := range masquerades {
-		d.Masquerades = append(d.Masquerades, &Masquerade{Domain: m.Domain, IpAddress: m.IpAddress})
+		var sni string
+		if d.SNIConfig != nil && d.SNIConfig.UseArbitrarySNIs {
+			crc32Hash := int(crc32.ChecksumIEEE([]byte(m.IpAddress)))
+			sni = d.SNIConfig.ArbitrarySNIs[crc32Hash%len(d.SNIConfig.ArbitrarySNIs)]
+		}
+		d.Masquerades = append(d.Masquerades, &Masquerade{Domain: m.Domain, IpAddress: m.IpAddress, SNI: sni})
 	}
 	d.PassthroughPatterns = append(d.PassthroughPatterns, passthrough...)
 	return d