package detectors import ( "context" "fmt" "net" "net/http" "net/http/httptest" "sync" "sync/atomic" "testing" "time" "github.com/stretchr/testify/assert" "github.com/trufflesecurity/trufflehog/v3/pkg/pb/detector_typepb" ) func TestWithNoLocalIP(t *testing.T) { t.Run("Prevents dialing local IP", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") _, err := transport.DialContext(context.Background(), "tcp", "127.0.0.1:8080") assert.Error(t, err) assert.ErrorIs(t, err, ErrNoLocalIP) }) t.Run("Prevents dialing wildcard IP", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") _, err := transport.DialContext(context.Background(), "tcp", "0.0.0.0:8080") assert.Error(t, err) assert.ErrorIs(t, err, ErrNoLocalIP) }) t.Run("Prevents dialing IPv6 wildcard IP", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") _, err := transport.DialContext(context.Background(), "tcp", "[::]:8080") assert.Error(t, err) assert.ErrorIs(t, err, ErrNoLocalIP) }) t.Run("Allows dialing non-local host", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") conn, err := transport.DialContext(context.Background(), "tcp", "google.com:80") assert.NoError(t, err) assert.NotNil(t, conn) _ = conn.Close() }) t.Run("Allows dialing non-local IP", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") conn, err := transport.DialContext(context.Background(), "tcp", "1.1.1.1:80") assert.NoError(t, err) assert.NotNil(t, conn) _ = conn.Close() }) t.Run("Handles invalid address", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") _, err := transport.DialContext(context.Background(), "tcp", "invalid-address") assert.Error(t, err) }) t.Run("Handles non-existent hostname", func(t *testing.T) { client := &http.Client{} WithNoLocalIP()(client) transport, ok := client.Transport.(*http.Transport) assert.True(t, ok, "Expected transport to be *http.Transport") _, err := transport.DialContext(context.Background(), "tcp", "non-existent-host.local:80") assert.Error(t, err) }) } // TestDoWithDedup_Singleflight verifies that concurrent DoWithDedup calls sharing the // same detector type and credential are coalesced into one network call. Each request // the server receives returns a distinct body, so all goroutines should observe the // body from exactly one actual server-side request. func TestDoWithDedup_Singleflight(t *testing.T) { var requestCount int32 // The 20 ms sleep keeps the first request in-flight long enough for all // goroutines to call DoWithDedup before the result is ready. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { n := atomic.AddInt32(&requestCount, 1) time.Sleep(20 * time.Millisecond) _, _ = fmt.Fprintf(w, `{"request":%d}`, n) })) defer server.Close() client := NewClientWithDedup(server.Client()) const goroutines = 5 bodies := make([]string, goroutines) statuses := make([]int, goroutines) errs := make([]error, goroutines) var wg sync.WaitGroup for i := range goroutines { wg.Add(1) go func(i int) { defer wg.Done() req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, server.URL, http.NoBody) if err != nil { errs[i] = err return } resp, err := DoWithDedup(client, detector_typepb.DetectorType_Meraki, "test-credential", req) if err != nil { errs[i] = err return } defer func() { _ = resp.Body.Close() }() var buf [512]byte n, _ := resp.Body.Read(buf[:]) bodies[i] = string(buf[:n]) statuses[i] = resp.StatusCode }(i) } wg.Wait() for _, err := range errs { assert.NoError(t, err) } for _, s := range statuses { assert.Equal(t, http.StatusOK, s) } assert.Equal(t, int32(1), atomic.LoadInt32(&requestCount), "singleflight should coalesce all concurrent calls into one HTTP request") for i := 1; i < goroutines; i++ { assert.Equal(t, bodies[0], bodies[i]) } } // TestDoWithDedup_WaiterContextCancelled verifies that a waiter whose context is // cancelled bails out with an error while other waiters still receive the response, // and only one HTTP request is made. func TestDoWithDedup_WaiterContextCancelled(t *testing.T) { var requestCount int32 // inFlight is closed by the server once it starts handling the request, // giving us a reliable signal to cancel one waiter mid-flight. inFlight := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { atomic.AddInt32(&requestCount, 1) close(inFlight) time.Sleep(50 * time.Millisecond) w.WriteHeader(http.StatusOK) })) defer server.Close() client := NewClientWithDedup(server.Client()) ctx1, cancel1 := context.WithCancel(context.Background()) defer cancel1() type result struct { status int err error } results := make([]result, 3) var wg sync.WaitGroup for i, ctx := range []context.Context{ctx1, context.Background(), context.Background()} { wg.Add(1) go func(i int, ctx context.Context) { defer wg.Done() req, _ := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, http.NoBody) resp, err := DoWithDedup(client, detector_typepb.DetectorType_Meraki, "cred", req) if err != nil { results[i] = result{err: err} return } defer func() { _ = resp.Body.Close() }() results[i] = result{status: resp.StatusCode} }(i, ctx) } <-inFlight cancel1() wg.Wait() assert.ErrorIs(t, results[0].err, context.Canceled, "cancelled waiter should get context error") assert.NoError(t, results[1].err) assert.Equal(t, http.StatusOK, results[1].status) assert.NoError(t, results[2].err) assert.Equal(t, http.StatusOK, results[2].status) assert.Equal(t, int32(1), atomic.LoadInt32(&requestCount), "only one HTTP request should be made") } // TestDoWithDedup_FirstCallerContextCancelled verifies that cancelling the first // caller's context does not abort the shared in-flight HTTP call: the second caller // should still receive a valid response, and only one HTTP request is made. func TestDoWithDedup_FirstCallerContextCancelled(t *testing.T) { var requestCount int32 inFlight := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { atomic.AddInt32(&requestCount, 1) close(inFlight) time.Sleep(50 * time.Millisecond) w.WriteHeader(http.StatusOK) })) defer server.Close() client := NewClientWithDedup(server.Client()) ctx1, cancel1 := context.WithCancel(context.Background()) defer cancel1() var firstErr error var wg sync.WaitGroup wg.Add(1) go func() { defer wg.Done() req, _ := http.NewRequestWithContext(ctx1, http.MethodGet, server.URL, http.NoBody) resp, err := DoWithDedup(client, detector_typepb.DetectorType_Meraki, "cred", req) if err != nil { firstErr = err return } _ = resp.Body.Close() }() // Cancel the first caller once the server is processing, then immediately // start a second caller that should coalesce into the still-running call. <-inFlight cancel1() var secondStatus int var secondErr error wg.Add(1) go func() { defer wg.Done() req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, http.NoBody) resp, err := DoWithDedup(client, detector_typepb.DetectorType_Meraki, "cred", req) if err != nil { secondErr = err return } defer func() { _ = resp.Body.Close() }() secondStatus = resp.StatusCode }() wg.Wait() assert.ErrorIs(t, firstErr, context.Canceled, "first caller should get context error") assert.NoError(t, secondErr, "second caller should succeed despite first caller's cancellation") assert.Equal(t, http.StatusOK, secondStatus) assert.Equal(t, int32(1), atomic.LoadInt32(&requestCount), "only one HTTP request should be made") } // TestDoWithDedup_DeadlinePreserved verifies that the client timeout still applies // to the shared in-flight call after context.WithoutCancel strips cancellation. // A hanging server must not cause an indefinite leak. func TestDoWithDedup_DeadlinePreserved(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { select { case <-time.After(10 * time.Second): case <-r.Context().Done(): } })) defer server.Close() base := server.Client() base.Timeout = 75 * time.Millisecond client := NewClientWithDedup(base) req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, http.NoBody) start := time.Now() resp, err := DoWithDedup(client, detector_typepb.DetectorType_Meraki, "cred", req) if err == nil { defer func() { _ = resp.Body.Close() }() } elapsed := time.Since(start) assert.Error(t, err, "request to hanging server should fail") assert.Less(t, elapsed, time.Second, "timeout should be enforced by client deadline, not run indefinitely") } func TestIsLocalIP(t *testing.T) { testCases := []struct { name string ip net.IP expected bool }{ {"Loopback IPv4", net.ParseIP("127.0.0.1"), true}, {"Loopback IPv6", net.ParseIP("::1"), true}, {"Private IPv4", net.ParseIP("192.168.1.1"), true}, {"Private IPv6", net.ParseIP("fd00::1"), true}, {"Unspecified IPv4", net.ParseIP("0.0.0.0"), true}, {"Unspecified IPv6", net.ParseIP("::"), true}, {"Public IPv4", net.ParseIP("8.8.8.8"), false}, {"Public IPv6", net.ParseIP("2001:4860:4860::8888"), false}, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { result := isLocalIP(tc.ip) assert.Equal(t, tc.expected, result) }) } }