mirror of
https://github.com/safedep/pmg.git
synced 2026-08-03 07:24:09 +02:00
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>
This commit is contained in:
@@ -13,6 +13,7 @@ import (
|
|||||||
"github.com/safedep/pmg/internal/audit"
|
"github.com/safedep/pmg/internal/audit"
|
||||||
"github.com/safedep/pmg/proxy"
|
"github.com/safedep/pmg/proxy"
|
||||||
gobreaker "github.com/sony/gobreaker/v2"
|
gobreaker "github.com/sony/gobreaker/v2"
|
||||||
|
"golang.org/x/sync/singleflight"
|
||||||
"google.golang.org/grpc/codes"
|
"google.golang.org/grpc/codes"
|
||||||
"google.golang.org/grpc/status"
|
"google.golang.org/grpc/status"
|
||||||
)
|
)
|
||||||
@@ -26,6 +27,13 @@ type baseRegistryInterceptor struct {
|
|||||||
confirmationChan chan *ConfirmationRequest
|
confirmationChan chan *ConfirmationRequest
|
||||||
circuitBreaker *gobreaker.CircuitBreaker[*analyzer.PackageVersionAnalysisResult]
|
circuitBreaker *gobreaker.CircuitBreaker[*analyzer.PackageVersionAnalysisResult]
|
||||||
execContext InterceptorContext
|
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] {
|
func newAnalyzerCircuitBreaker(name string) *gobreaker.CircuitBreaker[*analyzer.PackageVersionAnalysisResult] {
|
||||||
@@ -110,6 +118,14 @@ func (b *baseRegistryInterceptor) analyzePackage(
|
|||||||
|
|
||||||
log.Debugf("[%s] Analyzing package %s@%s", ctx.RequestID, packageName, packageVersion)
|
log.Debugf("[%s] Analyzing package %s@%s", ctx.RequestID, packageName, packageVersion)
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
result, err := b.circuitBreaker.Execute(func() (*analyzer.PackageVersionAnalysisResult, error) {
|
result, err := b.circuitBreaker.Execute(func() (*analyzer.PackageVersionAnalysisResult, error) {
|
||||||
analysisCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
analysisCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@@ -131,10 +147,22 @@ func (b *baseRegistryInterceptor) analyzePackage(
|
|||||||
return res, err
|
return res, err
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("analyzer failed: %w", err)
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
b.cache.Set(ecosystem.String(), packageName, packageVersion, result)
|
b.cache.Set(ecosystem.String(), packageName, packageVersion, result)
|
||||||
|
return result, nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("analyzer failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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)
|
log.Debugf("[%s] Analysis complete for %s@%s: action=%d", ctx.RequestID, packageName, packageVersion, result.Action)
|
||||||
|
|
||||||
|
|||||||
@@ -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")
|
||||||
|
}
|
||||||
+88
-3
@@ -9,6 +9,7 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
"os"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
@@ -23,6 +24,15 @@ import (
|
|||||||
// tunnel connections, so this must be large enough for a full bulk install.
|
// tunnel connections, so this must be large enough for a full bulk install.
|
||||||
const defaultServerReadWriteTimeout = 30 * time.Minute
|
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
|
// ProxyServer manages the proxy lifecycle
|
||||||
type ProxyServer interface {
|
type ProxyServer interface {
|
||||||
// Start begins listening on the configured address
|
// Start begins listening on the configured address
|
||||||
@@ -68,6 +78,10 @@ type ProxyConfig struct {
|
|||||||
//
|
//
|
||||||
// If zero, defaults to 30 minutes.
|
// If zero, defaults to 30 minutes.
|
||||||
ServerReadWriteTimeout time.Duration
|
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
|
// DefaultProxyConfig returns a configuration with sensible defaults
|
||||||
@@ -78,6 +92,7 @@ func DefaultProxyConfig() *ProxyConfig {
|
|||||||
ConnectTimeout: 30 * time.Second,
|
ConnectTimeout: 30 * time.Second,
|
||||||
RequestTimeout: 5 * time.Minute,
|
RequestTimeout: 5 * time.Minute,
|
||||||
ServerReadWriteTimeout: defaultServerReadWriteTimeout,
|
ServerReadWriteTimeout: defaultServerReadWriteTimeout,
|
||||||
|
UpstreamRetries: defaultUpstreamRetries,
|
||||||
Interceptors: []Interceptor{},
|
Interceptors: []Interceptor{},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -86,6 +101,7 @@ type proxyServer struct {
|
|||||||
config *ProxyConfig
|
config *ProxyConfig
|
||||||
proxy *goproxy.ProxyHttpServer
|
proxy *goproxy.ProxyHttpServer
|
||||||
server *http.Server
|
server *http.Server
|
||||||
|
roundTripper goproxy.RoundTripper
|
||||||
|
|
||||||
listener net.Listener
|
listener net.Listener
|
||||||
interceptors map[string]Interceptor
|
interceptors map[string]Interceptor
|
||||||
@@ -120,9 +136,13 @@ func NewProxyServer(config *ProxyConfig) (ProxyServer, error) {
|
|||||||
proxy.Logger = &goproxyLoggerWrapper{}
|
proxy.Logger = &goproxyLoggerWrapper{}
|
||||||
proxy.Tr = newUpstreamTransport(config)
|
proxy.Tr = newUpstreamTransport(config)
|
||||||
|
|
||||||
// Set verbose to true for verbose logging.
|
// goproxy emits several log lines per request when Verbose is set. During a
|
||||||
// Logging is handled by our own logger which has log level controls.
|
// large install (5000+ packages) that is a substantial amount of per-request
|
||||||
proxy.Verbose = true
|
// 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
|
// Configure connection timeout for upstream connections during CONNECT requests
|
||||||
proxy.ConnectDial = func(network, addr string) (net.Conn, error) {
|
proxy.ConnectDial = func(network, addr string) (net.Conn, error) {
|
||||||
@@ -139,6 +159,10 @@ func NewProxyServer(config *ProxyConfig) (ProxyServer, error) {
|
|||||||
interceptors: make(map[string]Interceptor),
|
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 {
|
for _, interceptor := range config.Interceptors {
|
||||||
if err := ps.AddInterceptor(interceptor); err != nil {
|
if err := ps.AddInterceptor(interceptor); err != nil {
|
||||||
return nil, fmt.Errorf("failed to add interceptor %s: %w", interceptor.Name(), err)
|
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() {
|
func (ps *proxyServer) registerHandlers() {
|
||||||
ps.proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
|
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.
|
// 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.
|
// We detect and fix this before processing.
|
||||||
normalizeRequestURL(req)
|
normalizeRequestURL(req)
|
||||||
|
|
||||||
|
ctx.RoundTripper = ps.roundTripper
|
||||||
|
|
||||||
reqCtx, err := newRequestContext(req)
|
reqCtx, err := newRequestContext(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Errorf("Failed to create request context: %v", err)
|
log.Errorf("Failed to create request context: %v", err)
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user