chore: address comments

This commit is contained in:
Susana Cardoso Ferreira
2025-12-29 10:37:21 +00:00
parent 11ea5b9224
commit 5b58481440
2 changed files with 79 additions and 1 deletions
+76
View File
@@ -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))
}
+3 -1
View File
@@ -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