chore: address comments

ssncferreira · ssncferreira · commit feca199d79ff · 2025-12-17T12:56:04.000Z
diff --git a/enterprise/aiproxyd/aiproxyd_test.go b/enterprise/aiproxyd/aiproxyd_test.go
@@ -3,10 +3,16 @@ package aiproxyd_test
 import (
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
 	"encoding/pem"
+	"io"
 	"math/big"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
 	"os"
 	"path/filepath"
 	"testing"
@@ -17,6 +23,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 
 	"github.com/coder/coder/v2/enterprise/aiproxyd"
+	"github.com/coder/coder/v2/testutil"
 )
 
 // generateTestCA creates a temporary CA certificate and key for testing.
@@ -76,6 +83,7 @@ func TestNew(t *testing.T) {
 
 	t.Run("MissingCertFile", func(t *testing.T) {
 		t.Parallel()
+
 		logger := slogtest.Make(t, nil)
 
 		_, err := aiproxyd.New(t.Context(), logger, aiproxyd.Options{
@@ -88,6 +96,7 @@ func TestNew(t *testing.T) {
 
 	t.Run("MissingKeyFile", func(t *testing.T) {
 		t.Parallel()
+
 		logger := slogtest.Make(t, nil)
 
 		_, err := aiproxyd.New(t.Context(), logger, aiproxyd.Options{
@@ -100,6 +109,7 @@ func TestNew(t *testing.T) {
 
 	t.Run("InvalidCertFile", func(t *testing.T) {
 		t.Parallel()
+
 		logger := slogtest.Make(t, nil)
 
 		_, err := aiproxyd.New(t.Context(), logger, aiproxyd.Options{
@@ -150,3 +160,69 @@ func TestClose(t *testing.T) {
 	err = srv.Close()
 	require.NoError(t, err)
 }
+
+func TestProxy_MITM(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)
+		_, _ = w.Write([]byte("hello from target"))
+	}))
+	defer targetServer.Close()
+
+	certFile, keyFile := generateTestCA(t)
+	logger := slogtest.Make(t, nil)
+
+	// Start the proxy server.
+	srv, err := aiproxyd.New(t.Context(), logger, aiproxyd.Options{
+		ListenAddr: "127.0.0.1:8888",
+		CertFile:   certFile,
+		KeyFile:    keyFile,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = srv.Close() })
+
+	// Wait for the proxy server to be ready.
+	require.Eventually(t, func() bool {
+		conn, err := net.Dial("tcp", "127.0.0.1:8888")
+		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://127.0.0.1:8888")
+	require.NoError(t, err)
+
+	client := &http.Client{
+		Transport: &http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+			TLSClientConfig: &tls.Config{
+				MinVersion: tls.VersionTLS12,
+				RootCAs:    certPool,
+			},
+		},
+	}
+
+	// Make a request through the proxy to the target server.
+	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()
+
+	// Verify the response was successfully proxied.
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	require.Equal(t, "hello from target", string(body))
+}
diff --git a/enterprise/cli/aiproxyd.go b/enterprise/cli/aiproxyd.go
@@ -5,6 +5,8 @@ package cli
 import (
 	"context"
 
+	"golang.org/x/xerrors"
+
 	"github.com/coder/coder/v2/enterprise/aiproxyd"
 	"github.com/coder/coder/v2/enterprise/coderd"
 )
@@ -21,7 +23,7 @@ func newAIProxyDaemon(coderAPI *coderd.API) (*aiproxyd.Server, error) {
 		KeyFile:    coderAPI.DeploymentValues.AI.ProxyConfig.KeyFile.String(),
 	})
 	if err != nil {
-		return nil, err
+		return nil, xerrors.Errorf("failed to start in-memory aiproxy daemon: %w", err)
 	}
 
 	return srv, nil
