fix(proxy): harden MITM proxy reliability and scale for bulk installs (#314)

* fix(proxy): harden MITM proxy reliability and scale for bulk installs

Deep-dive analysis of dropped connections during large installs (5000+
packages with concurrent downloads) surfaced three issues, each verified
with a reproduction test before fixing.

1. Transient upstream errors tore down whole keep-alive tunnels.
   goproxy returns false (closing the entire MITM client tunnel) when a
   single upstream round-trip errors. Under load, CDN-fronted registries
   (e.g. Cloudflare for registry.npmjs.org) intermittently reset
   connections, so one transient reset dropped a pooled keep-alive socket
   and surfaced to the package manager as ECONNRESET / "socket hang up".
   Fix: route upstream round-trips through a resilient round tripper that
   retries idempotent, body-less requests with bounded linear backoff,
   absorbing transient resets and keeping the tunnel alive. A reproduction
   test shows the tunnel count drop from 3 to 1 across a transient failure.

2. Head-of-line amplification against the external analysis service.
   Concurrent requests for the same package version each issued their own
   gRPC call. Fix: de-duplicate in-flight analyses with singleflight so a
   burst of identical requests collapses into one upstream call.

3. Per-request goproxy verbose logging on the hot path.
   proxy.Verbose was always on, formatting several log lines per request
   even when discarded below debug level. Fix: enable goproxy verbose
   logging only when PMG runs at debug level.

Note: the hypothesis that the http.Server Read/WriteTimeout leaks onto
hijacked CONNECT tunnels was investigated and disproven (Go clears the
deadlines on Hijack); the behavioral guard tests for long-lived
connections and slow transfers are retained.

* fix: Type assertion error handling

---------

Co-authored-by: Claude <noreply@anthropic.com>

proxy/interceptors/base_registry.go

@@ -13,6 +13,7 @@ import (
 	"github.com/safedep/pmg/internal/audit"
 	"github.com/safedep/pmg/proxy"
 	gobreaker "github.com/sony/gobreaker/v2"
+	"golang.org/x/sync/singleflight"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -26,6 +27,13 @@ type baseRegistryInterceptor struct {
 	confirmationChan chan *ConfirmationRequest
 	circuitBreaker   *gobreaker.CircuitBreaker[*analyzer.PackageVersionAnalysisResult]
 	execContext      InterceptorContext
+
+	// inflight collapses concurrent analyses of the same package version into a
+	// single upstream call. During a large install the same transitive
+	// dependency is frequently requested across several connections at once.
+	// Without de-duplication each would issue its own gRPC call to the external
+	// analysis service, amplifying load and latency (head-of-line blocking).
+	inflight singleflight.Group
 }
 
 func newAnalyzerCircuitBreaker(name string) *gobreaker.CircuitBreaker[*analyzer.PackageVersionAnalysisResult] {
@@ -110,31 +118,51 @@ func (b *baseRegistryInterceptor) analyzePackage(
 
 	log.Debugf("[%s] Analyzing package %s@%s", ctx.RequestID, packageName, packageVersion)
 
-	result, err := b.circuitBreaker.Execute(func() (*analyzer.PackageVersionAnalysisResult, error) {
-		analysisCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
-
-		res, err := b.analyzer.Analyze(analysisCtx, pkgVersion)
-		if err != nil {
-			// NotFound means the package is not in the analysis DB — this is expected
-			// and should not count as a circuit breaker failure.
-			// Since gRPC v1.75.0, status.FromError unwraps error chains via errors.As.
-			if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
-				log.Debugf("[%s] Package %s@%s not found in analysis DB, allowing", ctx.RequestID, packageName, packageVersion)
-				return &analyzer.PackageVersionAnalysisResult{
-					PackageVersion: pkgVersion,
-					Action:         analyzer.ActionAllow,
-				}, nil
-			}
+	key := ecosystem.String() + ":" + packageName + ":" + packageVersion
+	resultAny, err, _ := b.inflight.Do(key, func() (interface{}, error) {
+		// Re-check the cache: a previous in-flight analysis for this key may
+		// have populated it after our own cache miss above.
+		if cached, ok := b.cache.Get(ecosystem.String(), packageName, packageVersion); ok {
+			return cached, nil
 		}
 
-		return res, err
+		result, err := b.circuitBreaker.Execute(func() (*analyzer.PackageVersionAnalysisResult, error) {
+			analysisCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
+			res, err := b.analyzer.Analyze(analysisCtx, pkgVersion)
+			if err != nil {
+				// NotFound means the package is not in the analysis DB — this is expected
+				// and should not count as a circuit breaker failure.
+				// Since gRPC v1.75.0, status.FromError unwraps error chains via errors.As.
+				if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
+					log.Debugf("[%s] Package %s@%s not found in analysis DB, allowing", ctx.RequestID, packageName, packageVersion)
+					return &analyzer.PackageVersionAnalysisResult{
+						PackageVersion: pkgVersion,
+						Action:         analyzer.ActionAllow,
+					}, nil
+				}
+			}
+
+			return res, err
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		b.cache.Set(ecosystem.String(), packageName, packageVersion, result)
+		return result, nil
 	})
 	if err != nil {
 		return nil, fmt.Errorf("analyzer failed: %w", err)
 	}
 
-	b.cache.Set(ecosystem.String(), packageName, packageVersion, result)
+	// Fail loudly rather than panic in the proxy hot path if a future change to
+	// the singleflight closure ever returns a different type or a nil result.
+	result, ok := resultAny.(*analyzer.PackageVersionAnalysisResult)
+	if !ok || result == nil {
+		return nil, fmt.Errorf("analyzer returned unexpected result type %T for %s@%s", resultAny, packageName, packageVersion)
+	}
 
 	log.Debugf("[%s] Analysis complete for %s@%s: action=%d", ctx.RequestID, packageName, packageVersion, result.Action)
 

proxy/interceptors/inflight_test.go

@@ -0,0 +1,67 @@
+package interceptors
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
+	"github.com/safedep/pmg/analyzer"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// slowAtomicAnalyzer counts calls atomically and blocks briefly so concurrent
+// callers overlap inside analyzePackage.
+type slowAtomicAnalyzer struct {
+	calls atomic.Int64
+	delay time.Duration
+}
+
+func (s *slowAtomicAnalyzer) Name() string { return "slow-atomic" }
+
+func (s *slowAtomicAnalyzer) Analyze(_ context.Context, pv *packagev1.PackageVersion) (*analyzer.PackageVersionAnalysisResult, error) {
+	s.calls.Add(1)
+	time.Sleep(s.delay)
+	return &analyzer.PackageVersionAnalysisResult{PackageVersion: pv, Action: analyzer.ActionAllow}, nil
+}
+
+// TestAnalyzePackageDeduplicatesConcurrentCalls verifies that a burst of
+// concurrent analyses for the same package version collapses into a single
+// upstream analyzer call (head-of-line mitigation), while distinct packages
+// are analyzed independently.
+func TestAnalyzePackageDeduplicatesConcurrentCalls(t *testing.T) {
+	mock := &slowAtomicAnalyzer{delay: 50 * time.Millisecond}
+	base := newTestBaseInterceptor(mock)
+	ctx := newTestRequestContext()
+
+	const concurrency = 50
+
+	var wg sync.WaitGroup
+	for i := 0; i < concurrency; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			result, err := base.analyzePackage(ctx, packagev1.Ecosystem_ECOSYSTEM_NPM, "same-pkg", "1.0.0")
+			assert.NoError(t, err)
+			require.NotNil(t, result)
+			assert.Equal(t, analyzer.ActionAllow, result.Action)
+		}()
+	}
+	wg.Wait()
+
+	assert.EqualValues(t, 1, mock.calls.Load(),
+		"concurrent analyses of the same package version should collapse into one upstream call")
+
+	// A subsequent call for the same package is served from cache (no new call).
+	_, err := base.analyzePackage(ctx, packagev1.Ecosystem_ECOSYSTEM_NPM, "same-pkg", "1.0.0")
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, mock.calls.Load(), "result should be cached after the first analysis")
+
+	// A different package version triggers its own analysis.
+	_, err = base.analyzePackage(ctx, packagev1.Ecosystem_ECOSYSTEM_NPM, "other-pkg", "2.0.0")
+	require.NoError(t, err)
+	assert.EqualValues(t, 2, mock.calls.Load(), "distinct package versions are analyzed independently")
+}

proxy/proxy.go

@@ -9,6 +9,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"os"
 	"strings"
 	"sync"
 	"time"
@@ -23,6 +24,15 @@ import (
 // tunnel connections, so this must be large enough for a full bulk install.
 const defaultServerReadWriteTimeout = 30 * time.Minute
 
+// defaultUpstreamRetries is the number of times an idempotent upstream
+// request is retried when the round-trip fails before any response is
+// received. Registries fronted by CDNs (e.g. Cloudflare for
+// registry.npmjs.org) intermittently reset connections under the connection
+// burst of a large "npm install". Retrying transparently here prevents a
+// single transient reset from tearing down the whole keep-alive MITM tunnel
+// (which surfaces to the client as a "socket hang up"/ECONNRESET).
+const defaultUpstreamRetries = 2
+
 // ProxyServer manages the proxy lifecycle
 type ProxyServer interface {
 	// Start begins listening on the configured address
@@ -68,6 +78,10 @@ type ProxyConfig struct {
 	//
 	// If zero, defaults to 30 minutes.
 	ServerReadWriteTimeout time.Duration
+
+	// UpstreamRetries bounds retries of idempotent upstream requests on
+	// transient round-trip failures. Zero disables retries.
+	UpstreamRetries int
 }
 
 // DefaultProxyConfig returns a configuration with sensible defaults
@@ -78,14 +92,16 @@ func DefaultProxyConfig() *ProxyConfig {
 		ConnectTimeout:         30 * time.Second,
 		RequestTimeout:         5 * time.Minute,
 		ServerReadWriteTimeout: defaultServerReadWriteTimeout,
+		UpstreamRetries:        defaultUpstreamRetries,
 		Interceptors:           []Interceptor{},
 	}
 }
 
 type proxyServer struct {
-	config *ProxyConfig
-	proxy  *goproxy.ProxyHttpServer
-	server *http.Server
+	config       *ProxyConfig
+	proxy        *goproxy.ProxyHttpServer
+	server       *http.Server
+	roundTripper goproxy.RoundTripper
 
 	listener     net.Listener
 	interceptors map[string]Interceptor
@@ -120,9 +136,13 @@ func NewProxyServer(config *ProxyConfig) (ProxyServer, error) {
 	proxy.Logger = &goproxyLoggerWrapper{}
 	proxy.Tr = newUpstreamTransport(config)
 
-	// Set verbose to true for verbose logging.
-	// Logging is handled by our own logger which has log level controls.
-	proxy.Verbose = true
+	// goproxy emits several log lines per request when Verbose is set. During a
+	// large install (5000+ packages) that is a substantial amount of per-request
+	// formatting and allocation on the hot path, even though dry/log discards
+	// the lines below the debug level. Only enable goproxy's verbose logging
+	// when PMG itself is running at debug level so the cost is paid only when
+	// the output is actually wanted.
+	proxy.Verbose = strings.EqualFold(os.Getenv("APP_LOG_LEVEL"), "debug")
 
 	// Configure connection timeout for upstream connections during CONNECT requests
 	proxy.ConnectDial = func(network, addr string) (net.Conn, error) {
@@ -139,6 +159,10 @@ func NewProxyServer(config *ProxyConfig) (ProxyServer, error) {
 		interceptors: make(map[string]Interceptor),
 	}
 
+	ps.roundTripper = goproxy.RoundTripperFunc(func(req *http.Request, _ *goproxy.ProxyCtx) (*http.Response, error) {
+		return ps.upstreamRoundTrip(req)
+	})
+
 	for _, interceptor := range config.Interceptors {
 		if err := ps.AddInterceptor(interceptor); err != nil {
 			return nil, fmt.Errorf("failed to add interceptor %s: %w", interceptor.Name(), err)
@@ -406,6 +430,65 @@ func (ps *proxyServer) configureMITM() {
 	}))
 }
 
+// upstreamRoundTrip executes the upstream round-trip with bounded retries for
+// idempotent requests. goproxy tears down the entire client MITM tunnel when a
+// single round-trip returns an error (see handleHttps in goproxy), so a lone
+// transient upstream reset would otherwise drop a pooled keep-alive connection
+// and surface as a "socket hang up"/ECONNRESET to the package manager. Retrying
+// here absorbs those transient failures and keeps the tunnel alive.
+//
+// Only requests that can be safely replayed are retried: idempotent methods
+// with no request body, and only while the client request context is live.
+func (ps *proxyServer) upstreamRoundTrip(req *http.Request) (*http.Response, error) {
+	maxRetries := ps.config.UpstreamRetries
+	if maxRetries < 0 {
+		maxRetries = 0
+	}
+
+	var resp *http.Response
+	var err error
+
+	for attempt := 0; ; attempt++ {
+		resp, err = ps.proxy.Tr.RoundTrip(req)
+		if err == nil {
+			return resp, nil
+		}
+
+		if attempt >= maxRetries || !isReplayableRequest(req) || req.Context().Err() != nil {
+			return resp, err
+		}
+
+		// Linear backoff capped at 500ms to avoid adding meaningful latency
+		// while still spacing out retries against a struggling upstream.
+		backoff := time.Duration(attempt+1) * 50 * time.Millisecond
+		if backoff > 500*time.Millisecond {
+			backoff = 500 * time.Millisecond
+		}
+
+		log.Debugf("Retrying upstream %s %s after transient error (attempt %d/%d): %v",
+			req.Method, req.URL.Host, attempt+1, maxRetries, err)
+
+		select {
+		case <-time.After(backoff):
+		case <-req.Context().Done():
+			return resp, err
+		}
+	}
+}
+
+// isReplayableRequest reports whether a request can be safely retried after an
+// upstream failure. Only idempotent methods without a body qualify, which
+// covers registry metadata reads and tarball downloads.
+func isReplayableRequest(req *http.Request) bool {
+	switch req.Method {
+	case http.MethodGet, http.MethodHead, http.MethodOptions:
+	default:
+		return false
+	}
+
+	return req.Body == nil || req.Body == http.NoBody
+}
+
 func (ps *proxyServer) registerHandlers() {
 	ps.proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
 		// Fix malformed URLs produced by goproxy's MITM URL reconstruction.
@@ -415,6 +498,8 @@ func (ps *proxyServer) registerHandlers() {
 		// We detect and fix this before processing.
 		normalizeRequestURL(req)
 
+		ctx.RoundTripper = ps.roundTripper
+
 		reqCtx, err := newRequestContext(req)
 		if err != nil {
 			log.Errorf("Failed to create request context: %v", err)

proxy/scale_load_test.go

@@ -0,0 +1,106 @@
+package proxy
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"golang.org/x/net/http2"
+)
+
+// TestLoadHTTP2UpstreamConcurrent hammers the MITM proxy with many concurrent
+// workers that reuse keep-alive connections (mimicking npm/pip socket pools),
+// against an HTTP/2 upstream with a bounded MaxConcurrentStreams (mimicking
+// registries fronted by Cloudflare). It reports failure counts and types so we
+// can observe whether the proxy drops connections under load.
+func TestLoadHTTP2UpstreamConcurrent(t *testing.T) {
+	if testing.Short() {
+		t.Skip("load test")
+	}
+
+	const (
+		workers          = 80
+		reqsPerWorker    = 40
+		payload          = 32 * 1024
+		upstreamLatency  = 5 * time.Millisecond
+		analyzeLatency   = 10 * time.Millisecond
+		maxConcurStreams = 50
+	)
+
+	var upstreamConns atomic.Int64
+
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(upstreamLatency)
+		_, _ = w.Write(make([]byte, payload))
+	})
+
+	upstream := httptest.NewUnstartedServer(handler)
+	upstream.EnableHTTP2 = true
+	upstream.Config.ConnState = func(_ net.Conn, state http.ConnState) {
+		if state == http.StateNew {
+			upstreamConns.Add(1)
+		}
+	}
+	// Bound concurrent streams per h2 connection like a real CDN.
+	if err := http2.ConfigureServer(upstream.Config, &http2.Server{MaxConcurrentStreams: maxConcurStreams}); err != nil {
+		t.Fatalf("configure h2: %v", err)
+	}
+	upstream.StartTLS()
+	defer upstream.Close()
+
+	host := mustHost(t, upstream.URL)
+	_, client := buildReproProxy(t, host, 30*time.Minute, analyzeLatency)
+	// Reuse connections aggressively across workers.
+	client.Transport.(*http.Transport).MaxIdleConnsPerHost = workers
+	client.Transport.(*http.Transport).MaxConnsPerHost = workers
+
+	var (
+		wg        sync.WaitGroup
+		fail      atomic.Int64
+		short     atomic.Int64
+		ok        atomic.Int64
+		firstErr  atomic.Value
+		startTime = time.Now()
+	)
+
+	for w := 0; w < workers; w++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			for i := 0; i < reqsPerWorker; i++ {
+				resp, err := client.Get(fmt.Sprintf("%s/pkg-%d-%d.tgz", upstream.URL, id, i))
+				if err != nil {
+					fail.Add(1)
+					firstErr.CompareAndSwap(nil, err.Error())
+					continue
+				}
+				n, err := io.Copy(io.Discard, resp.Body)
+				_ = resp.Body.Close()
+				if err != nil {
+					fail.Add(1)
+					firstErr.CompareAndSwap(nil, err.Error())
+					continue
+				}
+				if n != payload {
+					short.Add(1)
+					continue
+				}
+				ok.Add(1)
+			}
+		}(w)
+	}
+	wg.Wait()
+
+	total := workers * reqsPerWorker
+	t.Logf("total=%d ok=%d fail=%d short=%d elapsed=%s upstreamConns=%d",
+		total, ok.Load(), fail.Load(), short.Load(), time.Since(startTime), upstreamConns.Load())
+	if fe := firstErr.Load(); fe != nil {
+		t.Logf("first error: %v", fe)
+	}
+}

proxy/scale_teardown_test.go

@@ -0,0 +1,140 @@
+package proxy
+
+import (
+	"crypto/tls"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
+	"testing"
+
+	"github.com/safedep/pmg/proxy/certmanager"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// countingInterceptor records how many CONNECTs are MITM'd, so we can detect
+// whether the proxy tears down a keep-alive tunnel (forcing the client to open
+// a fresh CONNECT) when an individual upstream request fails.
+type countingInterceptor struct {
+	host     string
+	connects atomic.Int64
+}
+
+func (c *countingInterceptor) Name() string { return "counting" }
+
+func (c *countingInterceptor) ShouldIntercept(ctx *RequestContext) bool {
+	return ctx.Hostname == c.host
+}
+
+func (c *countingInterceptor) ShouldMITM(ctx *RequestContext) bool {
+	if ctx.Method == "CONNECT" {
+		c.connects.Add(1)
+	}
+	return true
+}
+
+func (c *countingInterceptor) HandleRequest(ctx *RequestContext) (*InterceptorResponse, error) {
+	return &InterceptorResponse{Action: ActionAllow}, nil
+}
+
+// TestTunnelSurvivesTransientUpstreamError verifies that a transient upstream
+// failure (server resets the connection before responding) is absorbed by the
+// proxy's upstream retry instead of tearing down the entire keep-alive MITM
+// tunnel. Without the retry, goproxy returns the error which closes the tunnel,
+// forcing the client to open a brand-new CONNECT tunnel — the mechanism behind
+// "random connection drops" / "socket hang up" under load.
+func TestTunnelSurvivesTransientUpstreamError(t *testing.T) {
+	// Number of upstream attempts that should fail before succeeding.
+	var failuresRemaining atomic.Int64
+
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if failuresRemaining.Add(-1) >= 0 {
+			// Reset the connection before sending any response, simulating a
+			// transient upstream failure (CDN rate-limit / RST).
+			hj, ok := w.(http.Hijacker)
+			require.True(t, ok)
+			conn, _, _ := hj.Hijack()
+			_ = conn.Close()
+			return
+		}
+		_, _ = w.Write([]byte("ok"))
+	}))
+	defer upstream.Close()
+
+	host := mustHost(t, upstream.URL)
+
+	ca, err := certmanager.GenerateCA(certmanager.DefaultCertManagerConfig())
+	require.NoError(t, err)
+	cm, err := certmanager.NewCertificateManagerWithCA(ca, certmanager.DefaultCertManagerConfig())
+	require.NoError(t, err)
+
+	ci := &countingInterceptor{host: host}
+
+	cfg := DefaultProxyConfig()
+	cfg.CertManager = cm
+	cfg.Interceptors = []Interceptor{ci}
+
+	server, err := NewProxyServer(cfg)
+	require.NoError(t, err)
+	ps := server.(*proxyServer)
+	ps.proxy.Tr.TLSClientConfig.InsecureSkipVerify = true
+	require.NoError(t, ps.Start())
+	t.Cleanup(func() { _ = ps.Stop(t.Context()) })
+
+	client := newProxyClient(t, ps.Address())
+	// Disable transparent retry so we observe the raw failure, but keep
+	// keep-alive so the tunnel is reused.
+	tr := client.Transport.(*http.Transport)
+	tr.DisableKeepAlives = false
+
+	doGet := func() (int, error) {
+		resp, err := client.Get(upstream.URL + "/pkg")
+		if err != nil {
+			return 0, err
+		}
+		defer func() { _ = resp.Body.Close() }()
+		_, _ = io.Copy(io.Discard, resp.Body)
+		return resp.StatusCode, nil
+	}
+
+	// 1) Warm up: establishes exactly one CONNECT tunnel.
+	code, err := doGet()
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, code)
+	require.EqualValues(t, 1, ci.connects.Load(), "expected a single CONNECT tunnel after warmup")
+
+	// 2) Inject a single transient upstream failure on the reused tunnel. The
+	// proxy should retry upstream and recover transparently, keeping the
+	// existing tunnel alive.
+	failuresRemaining.Store(1)
+	code, err = doGet()
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusOK, code, "request should recover via upstream retry")
+
+	// 3) A few more requests should keep reusing the same tunnel.
+	for i := 0; i < 3; i++ {
+		code, err := doGet()
+		require.NoError(t, err)
+		require.Equal(t, http.StatusOK, code)
+	}
+
+	connects := ci.connects.Load()
+	t.Logf("total CONNECT tunnels opened: %d (ideal: 1)", connects)
+
+	assert.EqualValues(t, 1, connects,
+		"a transient upstream error should not tear down the keep-alive tunnel")
+}
+
+func newProxyClient(t *testing.T, addr string) *http.Client {
+	t.Helper()
+	proxyURL, err := url.Parse("http://" + addr)
+	require.NoError(t, err)
+	return &http.Client{
+		Transport: &http.Transport{
+			Proxy:           http.ProxyURL(proxyURL),
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+}

proxy/scale_test.go

@@ -0,0 +1,218 @@
+package proxy
+
+import (
+	"crypto/tls"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/safedep/pmg/proxy/certmanager"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// reproInterceptor MITMs a single host and optionally simulates analyzer latency.
+type reproInterceptor struct {
+	host  string
+	delay time.Duration
+}
+
+func (r *reproInterceptor) Name() string { return "repro" }
+
+func (r *reproInterceptor) ShouldIntercept(ctx *RequestContext) bool {
+	return ctx.Hostname == r.host
+}
+
+func (r *reproInterceptor) ShouldMITM(ctx *RequestContext) bool { return true }
+
+func (r *reproInterceptor) HandleRequest(ctx *RequestContext) (*InterceptorResponse, error) {
+	if r.delay > 0 {
+		time.Sleep(r.delay)
+	}
+	return &InterceptorResponse{Action: ActionAllow}, nil
+}
+
+func newReproCertManager(t *testing.T) certmanager.CertificateManager {
+	t.Helper()
+	ca, err := certmanager.GenerateCA(certmanager.DefaultCertManagerConfig())
+	require.NoError(t, err)
+	cm, err := certmanager.NewCertificateManagerWithCA(ca, certmanager.DefaultCertManagerConfig())
+	require.NoError(t, err)
+	return cm
+}
+
+// buildReproProxy wires a MITM proxy in front of the given upstream host with
+// the supplied server read/write timeout and analyzer delay. It returns a
+// client configured to route through the proxy.
+func buildReproProxy(t *testing.T, upstreamHost string, serverRWTimeout, analyzeDelay time.Duration) (*proxyServer, *http.Client) {
+	t.Helper()
+
+	cm := newReproCertManager(t)
+
+	cfg := DefaultProxyConfig()
+	cfg.CertManager = cm
+	cfg.ServerReadWriteTimeout = serverRWTimeout
+	cfg.Interceptors = []Interceptor{&reproInterceptor{host: upstreamHost, delay: analyzeDelay}}
+
+	server, err := NewProxyServer(cfg)
+	require.NoError(t, err)
+
+	ps := server.(*proxyServer)
+
+	// Trust the upstream test server's self-signed cert (test-only).
+	ps.proxy.Tr.TLSClientConfig.InsecureSkipVerify = true
+
+	require.NoError(t, ps.Start())
+	t.Cleanup(func() { _ = ps.Stop(t.Context()) })
+
+	proxyURL, err := url.Parse("http://" + ps.Address())
+	require.NoError(t, err)
+
+	client := &http.Client{
+		Transport: &http.Transport{
+			Proxy:           http.ProxyURL(proxyURL),
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+
+	return ps, client
+}
+
+// TestActiveTransferSurvivesServerTimeout guards that the http.Server
+// ReadTimeout/WriteTimeout does NOT abort an in-flight MITM response whose body
+// streams for longer than the timeout. Go clears the server deadlines when the
+// CONNECT is hijacked, so a slow tarball stream must still complete. This test
+// pins that behavior so a future regression (or a config that re-applies a
+// deadline to hijacked tunnels) is caught.
+func TestActiveTransferSurvivesServerTimeout(t *testing.T) {
+	const total = 512 * 1024
+
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		flusher, ok := w.(http.Flusher)
+		require.True(t, ok)
+
+		w.Header().Set("Content-Length", "")
+		// Send a first chunk immediately, then stall longer than the server
+		// timeout before sending the rest. This models a slow tarball stream.
+		first := make([]byte, 1024)
+		_, _ = w.Write(first)
+		flusher.Flush()
+
+		time.Sleep(2 * time.Second)
+
+		rest := make([]byte, total-len(first))
+		_, _ = w.Write(rest)
+		flusher.Flush()
+	}))
+	defer upstream.Close()
+
+	host := mustHost(t, upstream.URL)
+
+	// Server timeout shorter than the upstream body stall (2s) to force the
+	// leaked deadline to fire during the transfer.
+	_, client := buildReproProxy(t, host, 1*time.Second, 0)
+
+	resp, err := client.Get(upstream.URL + "/some-package-1.0.0.tgz")
+	require.NoError(t, err)
+	defer func() { _ = resp.Body.Close() }()
+
+	n, err := io.Copy(io.Discard, resp.Body)
+
+	assert.NoError(t, err, "active transfer should not be dropped by the server deadline")
+	assert.EqualValues(t, total, n, "client should receive the full response body")
+}
+
+// TestLongLivedConnectionSurvivesServerTimeout guards that a keep-alive MITM
+// connection reused across a window longer than ServerReadWriteTimeout keeps
+// working (the hijacked tunnel must not inherit the server read deadline).
+func TestLongLivedConnectionSurvivesServerTimeout(t *testing.T) {
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write([]byte("ok"))
+	}))
+	defer upstream.Close()
+
+	host := mustHost(t, upstream.URL)
+
+	// Short server timeout; we then reuse the same connection past it.
+	_, client := buildReproProxy(t, host, 1*time.Second, 0)
+
+	// First request establishes the CONNECT tunnel + keep-alive MITM conn.
+	doGet := func() error {
+		resp, err := client.Get(upstream.URL + "/pkg")
+		if err != nil {
+			return err
+		}
+		defer func() { _ = resp.Body.Close() }()
+		_, err = io.Copy(io.Discard, resp.Body)
+		return err
+	}
+
+	require.NoError(t, doGet())
+
+	// Idle past the server read/write timeout, then reuse the connection.
+	time.Sleep(1500 * time.Millisecond)
+
+	assert.NoError(t, doGet(), "reused keep-alive connection should survive past the server timeout")
+}
+
+// TestConcurrentDownloads validates that many concurrent downloads through the
+// MITM proxy with simulated analyzer latency all complete without dropped
+// connections. This is the baseline scale/correctness guard.
+func TestConcurrentDownloads(t *testing.T) {
+	const (
+		payload     = 64 * 1024
+		concurrency = 200
+	)
+
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write(make([]byte, payload))
+	}))
+	defer upstream.Close()
+
+	host := mustHost(t, upstream.URL)
+
+	_, client := buildReproProxy(t, host, 30*time.Minute, 20*time.Millisecond)
+
+	var (
+		wg       sync.WaitGroup
+		failures atomic.Int64
+		shortRd  atomic.Int64
+	)
+
+	for i := 0; i < concurrency; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			resp, err := client.Get(upstream.URL + "/pkg.tgz")
+			if err != nil {
+				failures.Add(1)
+				return
+			}
+			defer func() { _ = resp.Body.Close() }()
+			n, err := io.Copy(io.Discard, resp.Body)
+			if err != nil {
+				failures.Add(1)
+				return
+			}
+			if n != payload {
+				shortRd.Add(1)
+			}
+		}()
+	}
+	wg.Wait()
+
+	assert.EqualValues(t, 0, failures.Load(), "no requests should fail under concurrency")
+	assert.EqualValues(t, 0, shortRd.Load(), "no truncated responses under concurrency")
+}
+
+func mustHost(t *testing.T, rawURL string) string {
+	t.Helper()
+	u, err := url.Parse(rawURL)
+	require.NoError(t, err)
+	return u.Hostname()
+}