From 695f57f7ffb776a284319c5733cc28f3ef029984 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Tue, 5 Dec 2023 02:13:42 -0800 Subject: [PATCH] fix: use header flags in wsproxy server (#10985) --- cli/exp_scaletest.go | 18 +++++----- cli/root.go | 53 ++++++++++++----------------- codersdk/client.go | 30 +++++++++++++++++ codersdk/workspaceagents.go | 7 ++-- enterprise/cli/proxyserver.go | 11 +++++- enterprise/cli/proxyserver_test.go | 54 ++++++++++++++++++++++++++++++ 6 files changed, 127 insertions(+), 46 deletions(-) create mode 100644 enterprise/cli/proxyserver_test.go diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index 663f195164..9c88272e95 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -419,9 +419,9 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd { } client.HTTPClient = &http.Client{ - Transport: &headerTransport{ - transport: http.DefaultTransport, - header: map[string][]string{ + Transport: &codersdk.HeaderTransport{ + Transport: http.DefaultTransport, + Header: map[string][]string{ codersdk.BypassRatelimitHeader: {"true"}, }, }, @@ -570,9 +570,9 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd { } client.HTTPClient = &http.Client{ - Transport: &headerTransport{ - transport: http.DefaultTransport, - header: map[string][]string{ + Transport: &codersdk.HeaderTransport{ + Transport: http.DefaultTransport, + Header: map[string][]string{ codersdk.BypassRatelimitHeader: {"true"}, }, }, @@ -896,9 +896,9 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd { // Bypass rate limiting client.HTTPClient = &http.Client{ - Transport: &headerTransport{ - transport: http.DefaultTransport, - header: map[string][]string{ + Transport: &codersdk.HeaderTransport{ + Transport: http.DefaultTransport, + Header: map[string][]string{ codersdk.BypassRatelimitHeader: {"true"}, }, }, diff --git a/cli/root.go b/cli/root.go index 31aa5fec47..fbabad6d18 100644 --- a/cli/root.go +++ b/cli/root.go @@ -471,11 +471,11 @@ type RootCmd struct { } func addTelemetryHeader(client *codersdk.Client, inv *clibase.Invocation) { - transport, ok := client.HTTPClient.Transport.(*headerTransport) + transport, ok := client.HTTPClient.Transport.(*codersdk.HeaderTransport) if !ok { - transport = &headerTransport{ - transport: client.HTTPClient.Transport, - header: http.Header{}, + transport = &codersdk.HeaderTransport{ + Transport: client.HTTPClient.Transport, + Header: http.Header{}, } client.HTTPClient.Transport = transport } @@ -509,7 +509,7 @@ func addTelemetryHeader(client *codersdk.Client, inv *clibase.Invocation) { return } - transport.header.Add(codersdk.CLITelemetryHeader, s) + transport.Header.Add(codersdk.CLITelemetryHeader, s) } // InitClient sets client to a new client. @@ -609,10 +609,10 @@ func (r *RootCmd) initClientInternal(client *codersdk.Client, allowTokenMissing } } -func (r *RootCmd) setClient(ctx context.Context, client *codersdk.Client, serverURL *url.URL) error { - transport := &headerTransport{ - transport: http.DefaultTransport, - header: http.Header{}, +func (r *RootCmd) HeaderTransport(ctx context.Context, serverURL *url.URL) (*codersdk.HeaderTransport, error) { + transport := &codersdk.HeaderTransport{ + Transport: http.DefaultTransport, + Header: http.Header{}, } headers := r.header if r.headerCommand != "" { @@ -630,23 +630,32 @@ func (r *RootCmd) setClient(ctx context.Context, client *codersdk.Client, server cmd.Stderr = io.Discard err := cmd.Run() if err != nil { - return xerrors.Errorf("failed to run %v: %w", cmd.Args, err) + return nil, xerrors.Errorf("failed to run %v: %w", cmd.Args, err) } scanner := bufio.NewScanner(&outBuf) for scanner.Scan() { headers = append(headers, scanner.Text()) } if err := scanner.Err(); err != nil { - return xerrors.Errorf("scan %v: %w", cmd.Args, err) + return nil, xerrors.Errorf("scan %v: %w", cmd.Args, err) } } for _, header := range headers { parts := strings.SplitN(header, "=", 2) if len(parts) < 2 { - return xerrors.Errorf("split header %q had less than two parts", header) + return nil, xerrors.Errorf("split header %q had less than two parts", header) } - transport.header.Add(parts[0], parts[1]) + transport.Header.Add(parts[0], parts[1]) } + return transport, nil +} + +func (r *RootCmd) setClient(ctx context.Context, client *codersdk.Client, serverURL *url.URL) error { + transport, err := r.HeaderTransport(ctx, serverURL) + if err != nil { + return xerrors.Errorf("create header transport: %w", err) + } + client.URL = serverURL client.HTTPClient = &http.Client{ Transport: transport, @@ -853,24 +862,6 @@ func (r *RootCmd) Verbosef(inv *clibase.Invocation, fmtStr string, args ...inter } } -type headerTransport struct { - transport http.RoundTripper - header http.Header -} - -func (h *headerTransport) Header() http.Header { - return h.header.Clone() -} - -func (h *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) { - for k, v := range h.header { - for _, vv := range v { - req.Header.Add(k, vv) - } - } - return h.transport.RoundTrip(req) -} - // DumpHandler provides a custom SIGQUIT and SIGTRAP handler that dumps the // stacktrace of all goroutines to stderr and a well-known file in the home // directory. This is useful for debugging deadlock issues that may occur in diff --git a/codersdk/client.go b/codersdk/client.go index 965f2653d9..b95af72ab3 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -530,3 +530,33 @@ func WithQueryParam(key, value string) RequestOption { r.URL.RawQuery = q.Encode() } } + +// HeaderTransport is a http.RoundTripper that adds some headers to all requests. +// @typescript-ignore HeaderTransport +type HeaderTransport struct { + Transport http.RoundTripper + Header http.Header +} + +var _ http.RoundTripper = &HeaderTransport{} + +func (h *HeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) { + for k, v := range h.Header { + for _, vv := range v { + req.Header.Add(k, vv) + } + } + if h.Transport == nil { + h.Transport = http.DefaultTransport + } + return h.Transport.RoundTrip(req) +} + +func (h *HeaderTransport) CloseIdleConnections() { + type closeIdler interface { + CloseIdleConnections() + } + if tr, ok := h.Transport.(closeIdler); ok { + tr.CloseIdleConnections() + } +} diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go index ac3f28aa28..d3cfabcb63 100644 --- a/codersdk/workspaceagents.go +++ b/codersdk/workspaceagents.go @@ -273,11 +273,8 @@ func (c *Client) DialWorkspaceAgent(dialCtx context.Context, agentID uuid.UUID, ip := tailnet.IP() var header http.Header - headerTransport, ok := c.HTTPClient.Transport.(interface { - Header() http.Header - }) - if ok { - header = headerTransport.Header() + if headerTransport, ok := c.HTTPClient.Transport.(*HeaderTransport); ok { + header = headerTransport.Header } conn, err := tailnet.NewConn(&tailnet.Options{ Addresses: []netip.Prefix{netip.PrefixFrom(ip, 128)}, diff --git a/enterprise/cli/proxyserver.go b/enterprise/cli/proxyserver.go index c245426b51..bf5db6d0f7 100644 --- a/enterprise/cli/proxyserver.go +++ b/enterprise/cli/proxyserver.go @@ -43,7 +43,7 @@ func (c *closers) Add(f func()) { *c = append(*c, f) } -func (*RootCmd) proxyServer() *clibase.Cmd { +func (r *RootCmd) proxyServer() *clibase.Cmd { var ( cfg = new(codersdk.DeploymentValues) // Filter options for only relevant ones. @@ -192,6 +192,15 @@ func (*RootCmd) proxyServer() *clibase.Cmd { defer httpClient.CloseIdleConnections() closers.Add(httpClient.CloseIdleConnections) + // Attach header transport so we process --header and + // --header-command flags + headerTransport, err := r.HeaderTransport(ctx, primaryAccessURL.Value()) + if err != nil { + return xerrors.Errorf("configure header transport: %w", err) + } + headerTransport.Transport = httpClient.Transport + httpClient.Transport = headerTransport + // A newline is added before for visibility in terminal output. cliui.Infof(inv.Stdout, "\nView the Web UI: %s", cfg.AccessURL.String()) diff --git a/enterprise/cli/proxyserver_test.go b/enterprise/cli/proxyserver_test.go new file mode 100644 index 0000000000..72f576008f --- /dev/null +++ b/enterprise/cli/proxyserver_test.go @@ -0,0 +1,54 @@ +package cli_test + +import ( + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/pty/ptytest" +) + +func Test_Headers(t *testing.T) { + t.Parallel() + + const ( + headerName1 = "X-Test-Header-1" + headerVal1 = "test-value-1" + headerName2 = "X-Test-Header-2" + headerVal2 = "test-value-2" + ) + + // We're not going to actually start a proxy, we're going to point it + // towards a fake server that returns an unexpected status code. This'll + // cause the proxy to exit with an error that we can check for. + var called int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&called, 1) + assert.Equal(t, headerVal1, r.Header.Get(headerName1)) + assert.Equal(t, headerVal2, r.Header.Get(headerName2)) + + w.WriteHeader(http.StatusTeapot) // lol + })) + defer srv.Close() + + inv, _ := newCLI(t, "wsproxy", "server", + "--primary-access-url", srv.URL, + "--proxy-session-token", "test-token", + "--access-url", "http://localhost:8080", + "--header", fmt.Sprintf("%s=%s", headerName1, headerVal1), + "--header-command", fmt.Sprintf("printf %s=%s", headerName2, headerVal2), + ) + pty := ptytest.New(t) + inv.Stdout = pty.Output() + err := inv.Run() + require.Error(t, err) + require.ErrorContains(t, err, "unexpected status code 418") + require.NoError(t, pty.Close()) + + assert.EqualValues(t, 1, atomic.LoadInt64(&called)) +}