perf(jarm): coalesce concurrent fingerprints

dsl.go

@@ -34,7 +34,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/projectdiscovery/govaluate"
 	"github.com/Mzack9999/gcache"
 	"github.com/asaskevich/govalidator"
 	"github.com/brianvoe/gofakeit/v7"
@@ -48,6 +47,7 @@ import (
 	"github.com/projectdiscovery/dsl/randomip"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/gostruct"
+	"github.com/projectdiscovery/govaluate"
 	"github.com/projectdiscovery/mapcidr"
 	"github.com/projectdiscovery/utils/conn/connpool"
 	jarm "github.com/projectdiscovery/utils/crypto/jarm"
@@ -1504,15 +1504,17 @@ func init() {
 			return nil, err
 		}
 		// pick the first available proxy from common env vars (case-insensitive)
-		proxy := firstNonEmptyEnv("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy")
-		if proxy != "" {
-			socks5Dialer, err := connpool.NewCreateSOCKS5Dialer(proxy)
-			if err != nil {
-				return nil, err
+		proxy := jarmProxyFromEnvironment()
+		return coalesceJARM(proxy+"\x00"+host, func() (string, error) {
+			if proxy != "" {
+				socks5Dialer, err := connpool.NewCreateSOCKS5Dialer(proxy)
+				if err != nil {
+					return "", err
+				}
+				return jarm.HashWithDialer(socks5Dialer, hostname, port, 10)
 			}
-			return jarm.HashWithDialer(socks5Dialer, hostname, port, 10)
-		}
-		return jarm.HashWithDialer(nil, hostname, port, 10)
+			return jarm.HashWithDialer(nil, hostname, port, 10)
+		})
 	}))
 
 	MustAddFunction(NewWithSingleSignature("count",

go.mod

@@ -53,6 +53,7 @@ require (
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/net v0.57.0 // indirect
+	golang.org/x/sync v0.22.0
 	golang.org/x/sys v0.47.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -111,6 +111,8 @@ golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE=
 golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
+golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

jarm.go

@@ -0,0 +1,27 @@
+package dsl
+
+import "golang.org/x/sync/singleflight"
+
+var jarmFingerprints singleflight.Group
+
+func jarmProxyFromEnvironment() string {
+	// Nuclei and Aurora expose their SOCKS5 scan route using SOCKS5_PROXY.
+	// Keep the previous HTTP-family variables as fallbacks for callers that
+	// already place a socks5:// URL there.
+	return firstNonEmptyEnv(
+		"SOCKS5_PROXY", "socks5_proxy",
+		"HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy",
+	)
+}
+
+// coalesceJARM suppresses only overlapping calculations. Results are not kept
+// after the active callers return, so separate scans cannot observe stale data.
+func coalesceJARM(key string, fingerprint func() (string, error)) (string, error) {
+	value, err, _ := jarmFingerprints.Do(key, func() (interface{}, error) {
+		return fingerprint()
+	})
+	if err != nil {
+		return "", err
+	}
+	return value.(string), nil
+}

jarm_test.go

@@ -0,0 +1,88 @@
+package dsl
+
+import (
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCoalesceJARMOnlyRunsOverlappingFingerprintOnce(t *testing.T) {
+	firstStarted := make(chan struct{})
+	releaseFirst := make(chan struct{})
+	firstDone := make(chan struct {
+		value string
+		err   error
+	}, 1)
+	var calls atomic.Int32
+
+	go func() {
+		value, err := coalesceJARM("proxy\x00example.test:443", func() (string, error) {
+			calls.Add(1)
+			close(firstStarted)
+			<-releaseFirst
+			return "fingerprint", nil
+		})
+		firstDone <- struct {
+			value string
+			err   error
+		}{value: value, err: err}
+	}()
+
+	<-firstStarted
+	time.AfterFunc(10*time.Millisecond, func() { close(releaseFirst) })
+	value, err := coalesceJARM("proxy\x00example.test:443", func() (string, error) {
+		calls.Add(1)
+		return "duplicate", nil
+	})
+
+	require.NoError(t, err)
+	require.Equal(t, "fingerprint", value)
+	first := <-firstDone
+	require.NoError(t, first.err)
+	require.Equal(t, "fingerprint", first.value)
+	require.EqualValues(t, 1, calls.Load())
+}
+
+func TestCoalesceJARMDoesNotCacheCompletedFingerprint(t *testing.T) {
+	var calls atomic.Int32
+	for range 2 {
+		value, err := coalesceJARM("direct\x00example.test:443", func() (string, error) {
+			calls.Add(1)
+			return "fingerprint", nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, "fingerprint", value)
+	}
+	require.EqualValues(t, 2, calls.Load())
+}
+
+func TestJARMProxyPrefersConfiguredSOCKS5Route(t *testing.T) {
+	t.Setenv("SOCKS5_PROXY", "socks5://scan-route.example:1080")
+	t.Setenv("HTTP_PROXY", "socks5://legacy-http.example:1080")
+	t.Setenv("HTTPS_PROXY", "socks5://legacy-https.example:1080")
+	if runtime.GOOS != "windows" {
+		// Windows environment variable names are case-insensitive, so setting
+		// these aliases would overwrite the uppercase variables above.
+		t.Setenv("socks5_proxy", "socks5://lowercase.example:1080")
+		t.Setenv("http_proxy", "socks5://legacy-http-lower.example:1080")
+		t.Setenv("https_proxy", "socks5://legacy-https-lower.example:1080")
+	}
+
+	require.Equal(t, "socks5://scan-route.example:1080", jarmProxyFromEnvironment())
+}
+
+func TestJARMProxyRetainsLegacyFallback(t *testing.T) {
+	t.Setenv("SOCKS5_PROXY", "")
+	t.Setenv("HTTP_PROXY", "socks5://legacy.example:1080")
+	t.Setenv("HTTPS_PROXY", "")
+	if runtime.GOOS != "windows" {
+		t.Setenv("socks5_proxy", "")
+		t.Setenv("http_proxy", "")
+		t.Setenv("https_proxy", "")
+	}
+
+	require.Equal(t, "socks5://legacy.example:1080", jarmProxyFromEnvironment())
+}