Skip to content

Commit cd4d42b

Browse files
ethanndicksonethanndickson
committed
fix: support additional http headers on agent
1 parent 8c15192 commit cd4d42b

File tree

3 files changed

+106
-38
lines changed

3 files changed

+106
-38
lines changed

‎cli/agent.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
5050
slogJSONPath string
5151
slogStackdriverPath string
5252
blockFileTransfer bool
53+
agentHeaderCommand string
54+
agentHeader []string
5355
)
5456
cmd := &serpent.Command{
5557
Use: "agent",
@@ -176,6 +178,14 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
176178
// with large payloads can take a bit. e.g. startup scripts
177179
// may take a while to insert.
178180
client.SDK.HTTPClient.Timeout = 30 * time.Second
181+
// Attach header transport so we process --agent-header and
182+
// --agent-header-command flags
183+
headerTransport, err := headerTransport(ctx, r.agentURL, agentHeader, agentHeaderCommand)
184+
if err != nil {
185+
return xerrors.Errorf("configure header transport: %w", err)
186+
}
187+
headerTransport.Transport = client.SDK.HTTPClient.Transport
188+
client.SDK.HTTPClient.Transport = headerTransport
179189

180190
// Enable pprof handler
181191
// This prevents the pprof import from being accidentally deleted.
@@ -361,6 +371,18 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
361371
Value: serpent.StringOf(&pprofAddress),
362372
Description: "The address to serve pprof.",
363373
},
374+
{
375+
Flag: "agent-header-command",
376+
Env: "CODER_AGENT_HEADER_COMMAND",
377+
Value: serpent.StringOf(&agentHeaderCommand),
378+
Description: "An external command that outputs additional HTTP headers added to all requests. The command must output each header as `key=value` on its own line.",
379+
},
380+
{
381+
Flag: "agent-header",
382+
Env: "CODER_AGENT_HEADER",
383+
Value: serpent.StringArrayOf(&agentHeader),
384+
Description: "Additional HTTP headers added to all requests. Provide as " + `key=value` + ". Can be specified multiple times.",
385+
},
364386
{
365387
Flag: "no-reap",
366388

‎cli/agent_test.go

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@ package cli_test
33
import (
44
"context"
55
"fmt"
6+
"net/http"
7+
"net/http/httptest"
68
"os"
79
"path/filepath"
810
"runtime"
911
"strings"
12+
"sync/atomic"
1013
"testing"
1114

1215
"github.com/google/uuid"
@@ -229,6 +232,43 @@ func TestWorkspaceAgent(t *testing.T) {
229232
require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystems[0])
230233
require.Equal(t, codersdk.AgentSubsystemExectrace, resources[0].Agents[0].Subsystems[1])
231234
})
235+
t.Run("Header", func(t *testing.T) {
236+
t.Parallel()
237+
238+
var url string
239+
var called int64
240+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
241+
assert.Equal(t, "wow", r.Header.Get("X-Testing"))
242+
assert.Equal(t, "Ethan was Here!", r.Header.Get("Cool-Header"))
243+
assert.Equal(t, "very-wow-"+url, r.Header.Get("X-Process-Testing"))
244+
assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
245+
atomic.AddInt64(&called, 1)
246+
w.WriteHeader(http.StatusGone)
247+
}))
248+
defer srv.Close()
249+
url = srv.URL
250+
coderURLEnv := "$CODER_URL"
251+
if runtime.GOOS == "windows" {
252+
coderURLEnv = "%CODER_URL%"
253+
}
254+
255+
logDir := t.TempDir()
256+
inv, _ := clitest.New(t,
257+
"agent",
258+
"--auth", "token",
259+
"--agent-token", "fake-token",
260+
"--agent-url", srv.URL,
261+
"--log-dir", logDir,
262+
"--agent-header", "X-Testing=wow",
263+
"--agent-header", "Cool-Header=Ethan was Here!",
264+
"--agent-header-command", "printf X-Process-Testing=very-wow-"+coderURLEnv+"'\\r\\n'X-Process-Testing2=more-wow",
265+
)
266+
267+
clitest.Start(t, inv)
268+
require.Eventually(t, func() bool {
269+
return atomic.LoadInt64(&called) > 0
270+
}, testutil.WaitShort, testutil.IntervalFast)
271+
})
232272
}
233273

234274
func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {

‎cli/root.go

Lines changed: 44 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -550,44 +550,7 @@ func (r *RootCmd) InitClient(client *codersdk.Client) serpent.MiddlewareFunc {
550550
// HeaderTransport creates a new transport that executes `--header-command`
551551
// if it is set to add headers for all outbound requests.
552552
func (r *RootCmd) HeaderTransport(ctx context.Context, serverURL *url.URL) (*codersdk.HeaderTransport, error) {
553-
transport := &codersdk.HeaderTransport{
554-
Transport: http.DefaultTransport,
555-
Header: http.Header{},
556-
}
557-
headers := r.header
558-
if r.headerCommand != "" {
559-
shell := "sh"
560-
caller := "-c"
561-
if runtime.GOOS == "windows" {
562-
shell = "cmd.exe"
563-
caller = "/c"
564-
}
565-
var outBuf bytes.Buffer
566-
// #nosec
567-
cmd := exec.CommandContext(ctx, shell, caller, r.headerCommand)
568-
cmd.Env = append(os.Environ(), "CODER_URL="+serverURL.String())
569-
cmd.Stdout = &outBuf
570-
cmd.Stderr = io.Discard
571-
err := cmd.Run()
572-
if err != nil {
573-
return nil, xerrors.Errorf("failed to run %v: %w", cmd.Args, err)
574-
}
575-
scanner := bufio.NewScanner(&outBuf)
576-
for scanner.Scan() {
577-
headers = append(headers, scanner.Text())
578-
}
579-
if err := scanner.Err(); err != nil {
580-
return nil, xerrors.Errorf("scan %v: %w", cmd.Args, err)
581-
}
582-
}
583-
for _, header := range headers {
584-
parts := strings.SplitN(header, "=", 2)
585-
if len(parts) < 2 {
586-
return nil, xerrors.Errorf("split header %q had less than two parts", header)
587-
}
588-
transport.Header.Add(parts[0], parts[1])
589-
}
590-
return transport, nil
553+
return headerTransport(ctx, serverURL, r.header, r.headerCommand)
591554
}
592555

593556
func (r *RootCmd) configureClient(ctx context.Context, client *codersdk.Client, serverURL *url.URL, inv *serpent.Invocation) error {
@@ -1273,3 +1236,46 @@ type roundTripper func(req *http.Request) (*http.Response, error)
12731236
func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
12741237
return r(req)
12751238
}
1239+
1240+
// HeaderTransport creates a new transport that executes `--header-command`
1241+
// if it is set to add headers for all outbound requests.
1242+
func headerTransport(ctx context.Context, serverURL *url.URL, header []string, headerCommand string) (*codersdk.HeaderTransport, error) {
1243+
transport := &codersdk.HeaderTransport{
1244+
Transport: http.DefaultTransport,
1245+
Header: http.Header{},
1246+
}
1247+
headers := header
1248+
if headerCommand != "" {
1249+
shell := "sh"
1250+
caller := "-c"
1251+
if runtime.GOOS == "windows" {
1252+
shell = "cmd.exe"
1253+
caller = "/c"
1254+
}
1255+
var outBuf bytes.Buffer
1256+
// #nosec
1257+
cmd := exec.CommandContext(ctx, shell, caller, headerCommand)
1258+
cmd.Env = append(os.Environ(), "CODER_URL="+serverURL.String())
1259+
cmd.Stdout = &outBuf
1260+
cmd.Stderr = io.Discard
1261+
err := cmd.Run()
1262+
if err != nil {
1263+
return nil, xerrors.Errorf("failed to run %v: %w", cmd.Args, err)
1264+
}
1265+
scanner := bufio.NewScanner(&outBuf)
1266+
for scanner.Scan() {
1267+
headers = append(headers, scanner.Text())
1268+
}
1269+
if err := scanner.Err(); err != nil {
1270+
return nil, xerrors.Errorf("scan %v: %w", cmd.Args, err)
1271+
}
1272+
}
1273+
for _, header := range headers {
1274+
parts := strings.SplitN(header, "=", 2)
1275+
if len(parts) < 2 {
1276+
return nil, xerrors.Errorf("split header %q had less than two parts", header)
1277+
}
1278+
transport.Header.Add(parts[0], parts[1])
1279+
}
1280+
return transport, nil
1281+
}

0 commit comments

Comments
 (0)