mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: reroute AI provider requests to aibridged (#21343)
## Description Implements request routing for the AI Bridge Proxy. After MITM decryption, requests to known AI providers (Anthropic, OpenAI) are rewritten to the corresponding aibridged endpoint, while requests to unknown hosts are passed through to their original destination. ## Changes * Add `CoderAccessURL` configuration option for specifying the Coder deployment URL * Add `handleRequest` to route decrypted requests based on target host * Route known AI providers (Anthropic and OpenAI) to AI Bridge specific endpoint. * Passthrough requests to unknown hosts directly to their original destination * Inject Coder session token (from https://github.com/coder/coder/pull/21342) as `Authorization: Bearer` header for aibridged * Add tests for routing and passthrough behavior Depends on: https://github.com/coder/coder/pull/21342 Closes: https://github.com/coder/internal/issues/1181
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -16,6 +17,14 @@ import (
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog"
|
||||
|
||||
"github.com/coder/aibridge"
|
||||
)
|
||||
|
||||
// Known AI provider hosts.
|
||||
const (
|
||||
HostAnthropic = "api.anthropic.com"
|
||||
HostOpenAI = "api.openai.com"
|
||||
)
|
||||
|
||||
// loadMitmOnce ensures the MITM certificate is loaded exactly once.
|
||||
@@ -31,17 +40,21 @@ var loadMitmOnce sync.Once
|
||||
// - decrypting requests using the configured CA certificate
|
||||
// - forwarding requests to aibridged for processing
|
||||
type Server struct {
|
||||
ctx context.Context
|
||||
logger slog.Logger
|
||||
proxy *goproxy.ProxyHttpServer
|
||||
httpServer *http.Server
|
||||
listener net.Listener
|
||||
ctx context.Context
|
||||
logger slog.Logger
|
||||
proxy *goproxy.ProxyHttpServer
|
||||
httpServer *http.Server
|
||||
listener net.Listener
|
||||
coderAccessURL *url.URL
|
||||
}
|
||||
|
||||
// Options configures the AI Bridge Proxy server.
|
||||
type Options struct {
|
||||
// ListenAddr is the address the proxy server will listen on.
|
||||
ListenAddr string
|
||||
// CoderAccessURL is the URL of the Coder deployment where aibridged is running.
|
||||
// Requests to supported AI providers are forwarded here.
|
||||
CoderAccessURL string
|
||||
// CertFile is the path to the CA certificate file used for MITM.
|
||||
CertFile string
|
||||
// KeyFile is the path to the CA private key file used for MITM.
|
||||
@@ -62,6 +75,14 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
|
||||
return nil, xerrors.New("cert file and key file are required")
|
||||
}
|
||||
|
||||
if strings.TrimSpace(opts.CoderAccessURL) == "" {
|
||||
return nil, xerrors.New("coder access URL is required")
|
||||
}
|
||||
coderAccessURL, err := url.Parse(opts.CoderAccessURL)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("invalid coder access URL %q: %w", opts.CoderAccessURL, err)
|
||||
}
|
||||
|
||||
// Load CA certificate for MITM
|
||||
if err := loadMitmCertificate(opts.CertFile, opts.KeyFile); err != nil {
|
||||
return nil, xerrors.Errorf("failed to load MITM certificate: %w", err)
|
||||
@@ -70,9 +91,10 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
|
||||
proxy := goproxy.NewProxyHttpServer()
|
||||
|
||||
srv := &Server{
|
||||
ctx: ctx,
|
||||
logger: logger,
|
||||
proxy: proxy,
|
||||
ctx: ctx,
|
||||
logger: logger,
|
||||
proxy: proxy,
|
||||
coderAccessURL: coderAccessURL,
|
||||
}
|
||||
|
||||
// Reject CONNECT requests to non-standard ports.
|
||||
@@ -83,12 +105,14 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
|
||||
proxy.OnRequest().HandleConnectFunc(srv.portMiddleware(allowedPorts))
|
||||
|
||||
// Extract Coder session token from proxy authentication to forward to aibridged.
|
||||
// Decrypt all HTTPS requests via MITM. Requests are forwarded to
|
||||
// the original destination without modification for now.
|
||||
// TODO(ssncferreira): Route requests to aibridged will be implemented upstack.
|
||||
// Related to https://github.com/coder/internal/issues/1181
|
||||
proxy.OnRequest().HandleConnectFunc(srv.authMiddleware)
|
||||
|
||||
// Handle decrypted requests: route to aibridged for known AI providers, or passthrough to original destination.
|
||||
// TODO(ssncferreira): Currently the proxy always behaves as MITM, but this should only happen for known
|
||||
// AI providers as all other requests should be tunneled. This will be implemented upstack.
|
||||
// Related to https://github.com/coder/internal/issues/1182
|
||||
proxy.OnRequest().DoFunc(srv.handleRequest)
|
||||
|
||||
// Create listener first so we can get the actual address.
|
||||
// This is useful in tests where port 0 is used to avoid conflicts.
|
||||
listener, err := net.Listen("tcp", opts.ListenAddr)
|
||||
@@ -257,3 +281,97 @@ func extractCoderTokenFromProxyAuth(proxyAuth string) string {
|
||||
|
||||
return credentials[1]
|
||||
}
|
||||
|
||||
// providerFromHost maps the request host to the aibridge provider name.
|
||||
// - Known AI providers return their provider name, used to route to the
|
||||
// corresponding aibridge endpoint.
|
||||
// - Unknown hosts return empty string and are passed through directly.
|
||||
//
|
||||
// TODO(ssncferreira): Provider list configurable via domain allowlists will be implemented upstack.
|
||||
//
|
||||
// Related to https://github.com/coder/internal/issues/1182.
|
||||
func providerFromURL(reqURL *url.URL) string {
|
||||
if reqURL == nil {
|
||||
return ""
|
||||
}
|
||||
switch strings.ToLower(reqURL.Hostname()) {
|
||||
case HostAnthropic:
|
||||
return aibridge.ProviderAnthropic
|
||||
case HostOpenAI:
|
||||
return aibridge.ProviderOpenAI
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// handleRequest intercepts HTTP requests after MITM decryption.
|
||||
// - Requests to known AI providers are rewritten to aibridged, with the Coder session token
|
||||
// (from ctx.UserData, set during CONNECT) injected in the Authorization header.
|
||||
// - Unknown hosts are passed through to the original upstream.
|
||||
func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
|
||||
originalPath := req.URL.Path
|
||||
|
||||
// Check if this request is for a supported AI provider.
|
||||
provider := providerFromURL(req.URL)
|
||||
if provider == "" {
|
||||
// TODO(ssncferreira): After implementing selective MITM, this case should never
|
||||
// happen since unknown hosts will be tunneled, not decrypted.
|
||||
// Related to https://github.com/coder/internal/issues/1182
|
||||
s.logger.Debug(s.ctx, "passthrough request to unknown host",
|
||||
slog.F("host", req.Host),
|
||||
slog.F("method", req.Method),
|
||||
slog.F("path", originalPath),
|
||||
)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// Get the Coder session token stored during CONNECT.
|
||||
coderToken, _ := ctx.UserData.(string)
|
||||
|
||||
// Reject unauthenticated requests to AI providers.
|
||||
if coderToken == "" {
|
||||
s.logger.Warn(s.ctx, "rejecting unauthenticated request to AI provider",
|
||||
slog.F("host", req.Host),
|
||||
slog.F("provider", provider),
|
||||
)
|
||||
resp := goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusProxyAuthRequired, "Proxy authentication required")
|
||||
// Describe to the client how to authenticate with the proxy.
|
||||
resp.Header.Set("Proxy-Authenticate", `Basic realm="Coder AI Bridge Proxy"`)
|
||||
return req, resp
|
||||
}
|
||||
|
||||
// Rewrite the request to point to aibridged.
|
||||
if s.coderAccessURL == nil || s.coderAccessURL.String() == "" {
|
||||
s.logger.Error(s.ctx, "coderAccessURL is not configured")
|
||||
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Proxy misconfigured")
|
||||
}
|
||||
|
||||
aiBridgeURL, err := url.JoinPath(s.coderAccessURL.String(), "api/v2/aibridge", provider, originalPath)
|
||||
if err != nil {
|
||||
s.logger.Error(s.ctx, "failed to build aibridged URL", slog.Error(err))
|
||||
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to build AI Bridge URL")
|
||||
}
|
||||
|
||||
aiBridgeParsedURL, err := url.Parse(aiBridgeURL)
|
||||
if err != nil {
|
||||
s.logger.Error(s.ctx, "failed to parse aibridged URL", slog.Error(err))
|
||||
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to parse AI Bridge URL")
|
||||
}
|
||||
|
||||
// Preserve query parameters from the original request.
|
||||
// Both URL and Host must be set for the request to be properly routed.
|
||||
aiBridgeParsedURL.RawQuery = req.URL.RawQuery
|
||||
req.URL = aiBridgeParsedURL
|
||||
req.Host = aiBridgeParsedURL.Host
|
||||
|
||||
// Set Authorization header for aibridged authentication.
|
||||
req.Header.Set("Authorization", "Bearer "+coderToken)
|
||||
|
||||
s.logger.Debug(s.ctx, "routing request to aibridged",
|
||||
slog.F("provider", provider),
|
||||
slog.F("original_path", originalPath),
|
||||
slog.F("aibridged_url", aiBridgeParsedURL.String()),
|
||||
)
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -112,12 +113,60 @@ func TestNew(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "",
|
||||
ListenAddr: "",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "listen address is required")
|
||||
})
|
||||
|
||||
t.Run("MissingCoderAccessURL", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
certFile, keyFile := getSharedTestCA(t)
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "listen address is required")
|
||||
require.Contains(t, err.Error(), "coder access URL is required")
|
||||
})
|
||||
|
||||
t.Run("EmptyCoderAccessURL", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
certFile, keyFile := getSharedTestCA(t)
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: " ",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "coder access URL is required")
|
||||
})
|
||||
|
||||
t.Run("InvalidCoderAccessURL", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
certFile, keyFile := getSharedTestCA(t)
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: "://invalid",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "invalid coder access URL")
|
||||
})
|
||||
|
||||
t.Run("MissingCertFile", func(t *testing.T) {
|
||||
@@ -126,8 +175,9 @@ func TestNew(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: ":0",
|
||||
KeyFile: "key.pem",
|
||||
ListenAddr: ":0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
KeyFile: "key.pem",
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "cert file and key file are required")
|
||||
@@ -139,8 +189,9 @@ func TestNew(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: ":0",
|
||||
CertFile: "cert.pem",
|
||||
ListenAddr: ":0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: "cert.pem",
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "cert file and key file are required")
|
||||
@@ -152,9 +203,10 @@ func TestNew(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
_, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: ":0",
|
||||
CertFile: "/nonexistent/cert.pem",
|
||||
KeyFile: "/nonexistent/key.pem",
|
||||
ListenAddr: ":0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: "/nonexistent/cert.pem",
|
||||
KeyFile: "/nonexistent/key.pem",
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "failed to load MITM certificate")
|
||||
@@ -167,9 +219,10 @@ func TestNew(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, srv)
|
||||
@@ -186,9 +239,10 @@ func TestClose(t *testing.T) {
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -228,7 +282,7 @@ func TestProxy_PortValidation(t *testing.T) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("hello from target"))
|
||||
}))
|
||||
defer targetServer.Close()
|
||||
t.Cleanup(func() { targetServer.Close() })
|
||||
|
||||
targetURL, err := url.Parse(targetServer.URL)
|
||||
require.NoError(t, err)
|
||||
@@ -248,10 +302,11 @@ func TestProxy_PortValidation(t *testing.T) {
|
||||
|
||||
// Start the proxy server on a random port to avoid conflicts when running tests in parallel.
|
||||
srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
AllowedPorts: allowedPorts,
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
AllowedPorts: allowedPorts,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = srv.Close() })
|
||||
@@ -352,7 +407,7 @@ func TestProxy_Authentication(t *testing.T) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("hello from target"))
|
||||
}))
|
||||
defer targetServer.Close()
|
||||
t.Cleanup(func() { targetServer.Close() })
|
||||
|
||||
targetURL, err := url.Parse(targetServer.URL)
|
||||
require.NoError(t, err)
|
||||
@@ -363,10 +418,11 @@ func TestProxy_Authentication(t *testing.T) {
|
||||
// Start the proxy server on a random port to avoid conflicts when running tests in parallel.
|
||||
// The actual port is accessible via srv.Addr().
|
||||
srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
AllowedPorts: []string{targetURL.Port()},
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: "http://localhost:3000",
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
AllowedPorts: []string{targetURL.Port()},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = srv.Close() })
|
||||
@@ -430,3 +486,181 @@ func TestProxy_Authentication(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxy_MITM(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
targetHost string
|
||||
targetPort string // optional, if empty uses default HTTPS port (443)
|
||||
targetPath string
|
||||
expectedPath string
|
||||
passthrough bool
|
||||
}{
|
||||
{
|
||||
name: "AnthropicMessages",
|
||||
targetHost: "api.anthropic.com",
|
||||
targetPath: "/v1/messages",
|
||||
expectedPath: "/api/v2/aibridge/anthropic/v1/messages",
|
||||
},
|
||||
{
|
||||
name: "AnthropicNonDefaultPort",
|
||||
targetHost: "api.anthropic.com",
|
||||
targetPort: "8443",
|
||||
targetPath: "/v1/messages",
|
||||
expectedPath: "/api/v2/aibridge/anthropic/v1/messages",
|
||||
},
|
||||
{
|
||||
name: "OpenAIChatCompletions",
|
||||
targetHost: "api.openai.com",
|
||||
targetPath: "/v1/chat/completions",
|
||||
expectedPath: "/api/v2/aibridge/openai/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "OpenAINonDefaultPort",
|
||||
targetHost: "api.openai.com",
|
||||
targetPort: "8443",
|
||||
targetPath: "/v1/chat/completions",
|
||||
expectedPath: "/api/v2/aibridge/openai/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "UnknownHostPassthrough",
|
||||
targetPath: "/some/path",
|
||||
passthrough: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Track what aibridged receives.
|
||||
var receivedPath string
|
||||
var receivedAuth string
|
||||
|
||||
// Create a mock aibridged server.
|
||||
aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedPath = r.URL.Path
|
||||
receivedAuth = r.Header.Get("Authorization")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("hello from aibridged"))
|
||||
}))
|
||||
t.Cleanup(func() { aibridgedServer.Close() })
|
||||
|
||||
// Create a mock target server for passthrough tests.
|
||||
targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("hello from passthrough"))
|
||||
}))
|
||||
t.Cleanup(func() { targetServer.Close() })
|
||||
|
||||
certFile, keyFile := getSharedTestCA(t)
|
||||
logger := slogtest.Make(t, nil)
|
||||
|
||||
// Configure allowed ports based on test case.
|
||||
// AI provider tests connect to the specified port, or 443 if not specified.
|
||||
// Passthrough tests connect directly to the local target server's random port.
|
||||
var allowedPorts []string
|
||||
switch {
|
||||
case tt.passthrough:
|
||||
parsedTargetURL, err := url.Parse(targetServer.URL)
|
||||
require.NoError(t, err)
|
||||
allowedPorts = []string{parsedTargetURL.Port()}
|
||||
case tt.targetPort != "":
|
||||
allowedPorts = []string{tt.targetPort}
|
||||
default:
|
||||
allowedPorts = []string{"443"}
|
||||
}
|
||||
|
||||
// Start the proxy server pointing to our mock aibridged.
|
||||
srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
|
||||
ListenAddr: "127.0.0.1:0",
|
||||
CoderAccessURL: aibridgedServer.URL,
|
||||
CertFile: certFile,
|
||||
KeyFile: keyFile,
|
||||
AllowedPorts: allowedPorts,
|
||||
})
|
||||
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.
|
||||
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,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Build the target URL:
|
||||
// - For passthrough, target the local mock TLS server.
|
||||
// - For AI providers, use their real hostnames to trigger routing.
|
||||
// Non-default ports are included explicitly; default port (443) is omitted.
|
||||
var targetURL string
|
||||
switch {
|
||||
case tt.passthrough:
|
||||
targetURL, err = url.JoinPath(targetServer.URL, tt.targetPath)
|
||||
require.NoError(t, err)
|
||||
case tt.targetPort != "":
|
||||
targetURL, err = url.JoinPath("https://"+tt.targetHost+":"+tt.targetPort, tt.targetPath)
|
||||
require.NoError(t, err)
|
||||
default:
|
||||
targetURL, err = url.JoinPath("https://"+tt.targetHost, tt.targetPath)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Make a request through the proxy to the target URL.
|
||||
req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, targetURL, strings.NewReader(`{}`))
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := client.Do(req)
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
if tt.passthrough {
|
||||
// Verify request went to target server, not aibridged.
|
||||
require.Equal(t, "hello from passthrough", string(body))
|
||||
require.Empty(t, receivedPath, "aibridged should not receive passthrough requests")
|
||||
require.Empty(t, receivedAuth, "aibridged should not receive passthrough requests")
|
||||
} else {
|
||||
// Verify the request was routed to aibridged correctly.
|
||||
require.Equal(t, "hello from aibridged", string(body))
|
||||
require.Equal(t, tt.expectedPath, receivedPath)
|
||||
require.Equal(t, "Bearer test-session-token", receivedAuth)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,9 +18,10 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API) (*aibridgeproxyd.Server, error
|
||||
logger := coderAPI.Logger.Named("aibridgeproxyd")
|
||||
|
||||
srv, err := aibridgeproxyd.New(ctx, logger, aibridgeproxyd.Options{
|
||||
ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(),
|
||||
CertFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.CertFile.String(),
|
||||
KeyFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.KeyFile.String(),
|
||||
ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(),
|
||||
CoderAccessURL: coderAPI.AccessURL.String(),
|
||||
CertFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.CertFile.String(),
|
||||
KeyFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.KeyFile.String(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("failed to start in-memory aibridgeproxy daemon: %w", err)
|
||||
|
||||
Reference in New Issue
Block a user