From b522c9471a6a030758e46f8c04156250bfb0ef4a Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 29 Dec 2025 16:16:31 +0000 Subject: [PATCH] feat: add certificate caching for AI Bridge Proxy (#21344) ## Description Implements in-memory certificate caching for the AI Bridge MITM proxy. Certificate generation is expensive (RSA key generation + signing), so caching avoids repeated generation for the same hostname. ## Changes * Add `certCache` struct implementing `goproxy.CertStorage` with thread-safe double-check locking * Wire certificate cache into the proxy via `proxy.CertStore` * Add unit tests for cache behavior (hit, miss, errors, concurrency) * Add integration test to verify caching works end-to-end through the proxy Closes https://github.com/coder/internal/issues/1183 --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 14 ++ .../aibridgeproxyd/aibridgeproxyd_test.go | 85 +++++++++++ enterprise/aibridgeproxyd/certcache.go | 71 +++++++++ enterprise/aibridgeproxyd/certcache_test.go | 144 ++++++++++++++++++ 4 files changed, 314 insertions(+) create mode 100644 enterprise/aibridgeproxyd/certcache.go create mode 100644 enterprise/aibridgeproxyd/certcache_test.go diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index f57a353e64..6c4d6259d3 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -62,6 +62,9 @@ type Options struct { // AllowedPorts is the list of ports allowed for CONNECT requests. // Defaults to ["80", "443"] if empty. AllowedPorts []string + // CertStore is an optional certificate cache for MITM. If nil, a default + // cache is created. Exposed for testing. + CertStore goproxy.CertStorage } func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) { @@ -90,6 +93,17 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) proxy := goproxy.NewProxyHttpServer() + // Cache generated leaf certificates to avoid expensive RSA key generation + // and signing on every request to the same hostname. + // TODO(ssncferreira): Currently certs are cached for all MITM'd hosts, but once + // host filtering is implemented, only AI provider certs will be cached. + // Related to https://github.com/coder/internal/issues/1182 + if opts.CertStore != nil { + proxy.CertStore = opts.CertStore + } else { + proxy.CertStore = NewCertCache() + } + srv := &Server{ ctx: ctx, logger: logger, diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index c4ca3011e5..ec38ad93ae 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -254,6 +254,91 @@ func TestClose(t *testing.T) { require.NoError(t, err) } +func TestProxy_CertCaching(t *testing.T) { + t.Parallel() + + // Create a mock HTTPS server that will be the target of our proxied request. + targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer targetServer.Close() + + targetURL, err := url.Parse(targetServer.URL) + require.NoError(t, err) + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // Create a cert cache so we can inspect it after the request. + certCache := aibridgeproxyd.NewCertCache() + + // Start the proxy server with the certificate cache. + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + CertStore: certCache, + AllowedPorts: []string{targetURL.Port()}, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Close() }) + + proxyAddr := srv.Addr() + require.NotEmpty(t, proxyAddr) + + // Wait for the proxy server to be ready. + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", proxyAddr) + if err != nil { + return false + } + _ = conn.Close() + return true + }, testutil.WaitShort, testutil.IntervalFast) + + // Load the CA certificate so the client trusts the proxy's MITM certificate. + certPEM, err := os.ReadFile(certFile) + require.NoError(t, err) + certPool := x509.NewCertPool() + certPool.AppendCertsFromPEM(certPEM) + + // Create an HTTP client configured to use the proxy. + proxyURL, err := url.Parse("http://" + proxyAddr) + require.NoError(t, err) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + ProxyConnectHeader: http.Header{ + "Proxy-Authorization": []string{makeProxyAuthHeader("test-session-token")}, + }, + TLSClientConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: certPool, + }, + }, + } + + // Make a request through the proxy to the target server. + // This triggers MITM and caches the generated certificate. + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetServer.URL, nil) + require.NoError(t, err) + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + // Fetch with a generator that tracks calls: if the certificate was cached + // during the request above, the generator should not be called. + genCalls := 0 + _, err = certCache.Fetch(targetURL.Hostname(), func() (*tls.Certificate, error) { + genCalls++ + return &tls.Certificate{}, nil + }) + require.NoError(t, err) + require.Equal(t, 0, genCalls, "certificate should have been cached during request") +} + func TestProxy_PortValidation(t *testing.T) { t.Parallel() diff --git a/enterprise/aibridgeproxyd/certcache.go b/enterprise/aibridgeproxyd/certcache.go new file mode 100644 index 0000000000..9cdeb5228b --- /dev/null +++ b/enterprise/aibridgeproxyd/certcache.go @@ -0,0 +1,71 @@ +package aibridgeproxyd + +import ( + "crypto/tls" + "sync" + + "golang.org/x/xerrors" + "tailscale.com/util/singleflight" +) + +// CertCache implements goproxy.CertStorage to cache generated leaf certificates +// in memory. Certificate generation is expensive (RSA key generation + signing), +// so caching avoids repeated generation for the same hostname during MITM. +type CertCache struct { + mu sync.RWMutex + certs map[string]*tls.Certificate + singleFlight singleflight.Group[string, *tls.Certificate] +} + +// NewCertCache creates a new certificate cache that maps hostnames to their +// generated TLS certificates. +func NewCertCache() *CertCache { + return &CertCache{ + certs: make(map[string]*tls.Certificate), + } +} + +// Fetch retrieves a cached certificate for the given hostname, or generates +// and caches a new one using the provided generator function. +// +// Uses singleflight to ensure concurrent requests for the same hostname share +// a single in-flight generation rather than waiting on a mutex. This means only +// one goroutine generates the certificate while others wait on the result directly. +func (c *CertCache) Fetch(hostname string, genFunc func() (*tls.Certificate, error)) (*tls.Certificate, error) { + // Cache hit: check cache with read lock. + c.mu.RLock() + cert, ok := c.certs[hostname] + c.mu.RUnlock() + if ok { + return cert, nil + } + + // Cache miss: use singleflight to ensure only one goroutine generates + // the certificate for a given hostname, even under concurrent requests. + cert, err, _ := c.singleFlight.Do(hostname, func() (*tls.Certificate, error) { + // Double-check cache inside singleflight in case another call + // already populated it. + c.mu.RLock() + if cert, ok := c.certs[hostname]; ok { + c.mu.RUnlock() + return cert, nil + } + c.mu.RUnlock() + + cert, err := genFunc() + if err != nil { + return nil, err + } + if cert == nil { + return nil, xerrors.New("generator function returned nil certificate") + } + + c.mu.Lock() + c.certs[hostname] = cert + c.mu.Unlock() + + return cert, nil + }) + + return cert, err +} diff --git a/enterprise/aibridgeproxyd/certcache_test.go b/enterprise/aibridgeproxyd/certcache_test.go new file mode 100644 index 0000000000..75bbdf1b2a --- /dev/null +++ b/enterprise/aibridgeproxyd/certcache_test.go @@ -0,0 +1,144 @@ +package aibridgeproxyd_test + +import ( + "crypto/tls" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/enterprise/aibridgeproxyd" +) + +func TestCertCache_Fetch(t *testing.T) { + t.Parallel() + + t.Run("CacheMiss", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + expectedCert := &tls.Certificate{} + genCalls := 0 + + cert, err := cache.Fetch("example.com", func() (*tls.Certificate, error) { + genCalls++ + return expectedCert, nil + }) + + require.NoError(t, err) + require.Same(t, expectedCert, cert) + require.Equal(t, 1, genCalls) + }) + + t.Run("CacheHit", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + expectedCert := &tls.Certificate{} + genCalls := 0 + + gen := func() (*tls.Certificate, error) { + genCalls++ + return expectedCert, nil + } + + // First call: cache miss + cert1, err := cache.Fetch("example.com", gen) + require.NoError(t, err) + require.Same(t, expectedCert, cert1) + require.Equal(t, 1, genCalls) + + // Second call: cache hit, generator should not be called + cert2, err := cache.Fetch("example.com", gen) + require.NoError(t, err) + require.Same(t, expectedCert, cert2) + require.Equal(t, 1, genCalls) + }) + + t.Run("DifferentHostnames", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + cert1 := &tls.Certificate{} + cert2 := &tls.Certificate{} + + result1, err := cache.Fetch("example1.com", func() (*tls.Certificate, error) { + return cert1, nil + }) + require.NoError(t, err) + require.Same(t, cert1, result1) + + result2, err := cache.Fetch("example2.com", func() (*tls.Certificate, error) { + return cert2, nil + }) + require.NoError(t, err) + require.Same(t, cert2, result2) + + // Verify different hostnames have different certificates. + require.NotSame(t, result1, result2) + }) + + t.Run("GeneratorError", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + expectedErr := xerrors.New("generation failed") + + cert, err := cache.Fetch("example.com", func() (*tls.Certificate, error) { + return nil, expectedErr + }) + + require.ErrorIs(t, err, expectedErr) + require.Nil(t, cert) + }) + + t.Run("GeneratorReturnsNil", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + + cert, err := cache.Fetch("example.com", func() (*tls.Certificate, error) { + //nolint:nilnil // Intentionally testing this edge case + return nil, nil + }) + + require.ErrorContains(t, err, "generator function returned nil certificate") + require.Nil(t, cert) + }) + + t.Run("ConcurrentFetchSameHostname", func(t *testing.T) { + t.Parallel() + + cache := aibridgeproxyd.NewCertCache() + expectedCert := &tls.Certificate{} + var genCalls atomic.Int32 + + const numGoroutines = 10 + var wg sync.WaitGroup + wg.Add(numGoroutines) + + var fetchErrors atomic.Int32 + + // Spawn multiple goroutines that all request the same hostname concurrently. + for range numGoroutines { + go func() { + defer wg.Done() + cert, err := cache.Fetch("example.com", func() (*tls.Certificate, error) { + genCalls.Add(1) + return expectedCert, nil + }) + if err != nil || cert != expectedCert { + fetchErrors.Add(1) + } + }() + } + wg.Wait() + + require.Equal(t, int32(0), fetchErrors.Load()) + + // Generator should only be called once due to double-check locking. + require.Equal(t, int32(1), genCalls.Load()) + }) +}