Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitcd4d42b

Browse files
committed
fix: support additional http headers on agent
1 parent8c15192 commitcd4d42b

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
slogJSONPathstring
5151
slogStackdriverPathstring
5252
blockFileTransferbool
53+
agentHeaderCommandstring
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+
iferr!=nil {
185+
returnxerrors.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+
varurlstring
239+
varcalledint64
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+
defersrv.Close()
249+
url=srv.URL
250+
coderURLEnv:="$CODER_URL"
251+
ifruntime.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+
returnatomic.LoadInt64(&called)>0
270+
},testutil.WaitShort,testutil.IntervalFast)
271+
})
232272
}
233273

234274
funcmatchAgentWithVersion(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-
ifr.headerCommand!="" {
559-
shell:="sh"
560-
caller:="-c"
561-
ifruntime.GOOS=="windows" {
562-
shell="cmd.exe"
563-
caller="/c"
564-
}
565-
varoutBuf 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-
iferr!=nil {
573-
returnnil,xerrors.Errorf("failed to run %v: %w",cmd.Args,err)
574-
}
575-
scanner:=bufio.NewScanner(&outBuf)
576-
forscanner.Scan() {
577-
headers=append(headers,scanner.Text())
578-
}
579-
iferr:=scanner.Err();err!=nil {
580-
returnnil,xerrors.Errorf("scan %v: %w",cmd.Args,err)
581-
}
582-
}
583-
for_,header:=rangeheaders {
584-
parts:=strings.SplitN(header,"=",2)
585-
iflen(parts)<2 {
586-
returnnil,xerrors.Errorf("split header %q had less than two parts",header)
587-
}
588-
transport.Header.Add(parts[0],parts[1])
589-
}
590-
returntransport,nil
553+
returnheaderTransport(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 (rroundTripper)RoundTrip(req*http.Request) (*http.Response,error) {
12741237
returnr(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+
funcheaderTransport(ctx context.Context,serverURL*url.URL,header []string,headerCommandstring) (*codersdk.HeaderTransport,error) {
1243+
transport:=&codersdk.HeaderTransport{
1244+
Transport:http.DefaultTransport,
1245+
Header: http.Header{},
1246+
}
1247+
headers:=header
1248+
ifheaderCommand!="" {
1249+
shell:="sh"
1250+
caller:="-c"
1251+
ifruntime.GOOS=="windows" {
1252+
shell="cmd.exe"
1253+
caller="/c"
1254+
}
1255+
varoutBuf 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+
iferr!=nil {
1263+
returnnil,xerrors.Errorf("failed to run %v: %w",cmd.Args,err)
1264+
}
1265+
scanner:=bufio.NewScanner(&outBuf)
1266+
forscanner.Scan() {
1267+
headers=append(headers,scanner.Text())
1268+
}
1269+
iferr:=scanner.Err();err!=nil {
1270+
returnnil,xerrors.Errorf("scan %v: %w",cmd.Args,err)
1271+
}
1272+
}
1273+
for_,header:=rangeheaders {
1274+
parts:=strings.SplitN(header,"=",2)
1275+
iflen(parts)<2 {
1276+
returnnil,xerrors.Errorf("split header %q had less than two parts",header)
1277+
}
1278+
transport.Header.Add(parts[0],parts[1])
1279+
}
1280+
returntransport,nil
1281+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp