Skip to content

Commit 7b1affe

Browse files
spikecurtisspikecurtis
committed
chore: refactor instance identity to be a SessionTokenProvider
1 parent 7789799 commit 7b1affe

31 files changed

+471
-376
lines changed

‎agent/agent.go

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@ type Options struct {
7474
LogDir string
7575
TempDir string
7676
ScriptDataDir string
77-
ExchangeToken func(ctx context.Context) (string, error)
7877
Client Client
7978
ReconnectingPTYTimeout time.Duration
8079
EnvironmentVariables map[string]string
@@ -99,6 +98,7 @@ type Client interface {
9998
proto.DRPCAgentClient26, tailnetproto.DRPCTailnetClient26, error,
10099
)
101100
tailnet.DERPMapRewriter
101+
agentsdk.RefreshableSessionTokenProvider
102102
}
103103

104104
type Agent interface {
@@ -131,11 +131,6 @@ func New(options Options) Agent {
131131
}
132132
options.ScriptDataDir = options.TempDir
133133
}
134-
if options.ExchangeToken == nil {
135-
options.ExchangeToken = func(_ context.Context) (string, error) {
136-
return "", nil
137-
}
138-
}
139134
if options.ReportMetadataInterval == 0 {
140135
options.ReportMetadataInterval = time.Second
141136
}
@@ -172,7 +167,6 @@ func New(options Options) Agent {
172167
coordDisconnected: make(chan struct{}),
173168
environmentVariables: options.EnvironmentVariables,
174169
client: options.Client,
175-
exchangeToken: options.ExchangeToken,
176170
filesystem: options.Filesystem,
177171
logDir: options.LogDir,
178172
tempDir: options.TempDir,
@@ -203,7 +197,6 @@ func New(options Options) Agent {
203197
// coordinator during shut down.
204198
close(a.coordDisconnected)
205199
a.announcementBanners.Store(new([]codersdk.BannerConfig))
206-
a.sessionToken.Store(new(string))
207200
a.init()
208201
return a
209202
}
@@ -212,7 +205,6 @@ type agent struct {
212205
clock quartz.Clock
213206
logger slog.Logger
214207
client Client
215-
exchangeToken func(ctx context.Context) (string, error)
216208
tailnetListenPort uint16
217209
filesystem afero.Fs
218210
logDir string
@@ -254,7 +246,6 @@ type agent struct {
254246
scriptRunner *agentscripts.Runner
255247
announcementBanners atomic.Pointer[[]codersdk.BannerConfig] // announcementBanners is atomic because it is periodically updated.
256248
announcementBannersRefreshInterval time.Duration
257-
sessionToken atomic.Pointer[string]
258249
sshServer *agentssh.Server
259250
sshMaxTimeout time.Duration
260251
blockFileTransfer bool
@@ -916,11 +907,10 @@ func (a *agent) run() (retErr error) {
916907
// This allows the agent to refresh its token if necessary.
917908
// For instance identity this is required, since the instance
918909
// may not have re-provisioned, but a new agent ID was created.
919-
sessionToken, err := a.exchangeToken(a.hardCtx)
910+
err := a.client.RefreshToken(a.hardCtx)
920911
if err != nil {
921-
return xerrors.Errorf("exchange token: %w", err)
912+
return xerrors.Errorf("refresh token: %w", err)
922913
}
923-
a.sessionToken.Store(&sessionToken)
924914

925915
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs
926916
aAPI, tAPI, err := a.client.ConnectRPC26(a.hardCtx)
@@ -1359,7 +1349,7 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error)
13591349
"CODER_WORKSPACE_OWNER_NAME": manifest.OwnerName,
13601350

13611351
// Specific Coder subcommands require the agent token exposed!
1362-
"CODER_AGENT_TOKEN": *a.sessionToken.Load(),
1352+
"CODER_AGENT_TOKEN": a.client.GetSessionToken(),
13631353

13641354
// Git on Windows resolves with UNIX-style paths.
13651355
// If using backslashes, it's unable to find the executable.

‎agent/agenttest/agent.go

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package agenttest
22

33
import (
4-
"context"
54
"net/url"
65
"testing"
76

@@ -31,18 +30,11 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent
3130
}
3231

3332
if o.Client == nil {
34-
agentClient := agentsdk.New(coderURL)
35-
agentClient.SetSessionToken(agentToken)
33+
agentClient := agentsdk.New(coderURL, agentsdk.UsingFixedToken(agentToken))
3634
agentClient.SDK.SetLogger(log)
3735
o.Client = agentClient
3836
}
3937

40-
if o.ExchangeToken == nil {
41-
o.ExchangeToken = func(_ context.Context) (string, error) {
42-
return agentToken, nil
43-
}
44-
}
45-
4638
if o.LogDir == "" {
4739
o.LogDir = t.TempDir()
4840
}

‎agent/agenttest/client.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package agenttest
33
import (
44
"context"
55
"io"
6+
"net/http"
67
"slices"
78
"sync"
89
"sync/atomic"
@@ -28,6 +29,7 @@ import (
2829
"github.com/coder/coder/v2/tailnet"
2930
"github.com/coder/coder/v2/tailnet/proto"
3031
"github.com/coder/coder/v2/testutil"
32+
"github.com/coder/websocket"
3133
)
3234

3335
const statsInterval = 500 * time.Millisecond
@@ -92,6 +94,20 @@ type Client struct {
9294
derpMapOnce sync.Once
9395
}
9496

97+
func (c *Client) AsRequestOption() codersdk.RequestOption {
98+
return func(request *http.Request) {}
99+
}
100+
101+
func (c *Client) SetDialOption(*websocket.DialOptions) {}
102+
103+
func (c *Client) GetSessionToken() string {
104+
return "agenttest-token"
105+
}
106+
107+
func (c *Client) RefreshToken(context.Context) error {
108+
return nil
109+
}
110+
95111
func (*Client) RewriteDERPMap(*tailcfg.DERPMap) {}
96112

97113
func (c *Client) Close() {

‎cli/agent.go

Lines changed: 3 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import (
1515
"strings"
1616
"time"
1717

18-
"cloud.google.com/go/compute/metadata"
1918
"golang.org/x/xerrors"
2019
"gopkg.in/natefinch/lumberjack.v2"
2120

@@ -40,7 +39,6 @@ import (
4039

4140
func (r *RootCmd) workspaceAgent() *serpent.Command {
4241
var (
43-
auth string
4442
logDir string
4543
scriptDataDir string
4644
pprofAddress string
@@ -177,11 +175,10 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
177175
version := buildinfo.Version()
178176
logger.Info(ctx, "agent is starting now",
179177
slog.F("url", r.agentURL),
180-
slog.F("auth", auth),
178+
slog.F("auth", r.agentAuth),
181179
slog.F("version", version),
182180
)
183-
184-
client := agentsdk.New(r.agentURL)
181+
client, err := r.createAgentClient(ctx)
185182
client.SDK.SetLogger(logger)
186183
// Set a reasonable timeout so requests can't hang forever!
187184
// The timeout needs to be reasonably long, because requests
@@ -214,68 +211,6 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
214211
ignorePorts[port] = "debug"
215212
}
216213

217-
// exchangeToken returns a session token.
218-
// This is abstracted to allow for the same looping condition
219-
// regardless of instance identity auth type.
220-
var exchangeToken func(context.Context) (agentsdk.AuthenticateResponse, error)
221-
switch auth {
222-
case "token":
223-
token, _ := inv.ParsedFlags().GetString(varAgentToken)
224-
if token == "" {
225-
tokenFile, _ := inv.ParsedFlags().GetString(varAgentTokenFile)
226-
if tokenFile != "" {
227-
tokenBytes, err := os.ReadFile(tokenFile)
228-
if err != nil {
229-
return xerrors.Errorf("read token file %q: %w", tokenFile, err)
230-
}
231-
token = strings.TrimSpace(string(tokenBytes))
232-
}
233-
}
234-
if token == "" {
235-
return xerrors.Errorf("CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE must be set for token auth")
236-
}
237-
client.SetSessionToken(token)
238-
case "google-instance-identity":
239-
// This is *only* done for testing to mock client authentication.
240-
// This will never be set in a production scenario.
241-
var gcpClient *metadata.Client
242-
gcpClientRaw := ctx.Value("gcp-client")
243-
if gcpClientRaw != nil {
244-
gcpClient, _ = gcpClientRaw.(*metadata.Client)
245-
}
246-
exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
247-
return client.AuthGoogleInstanceIdentity(ctx, "", gcpClient)
248-
}
249-
case "aws-instance-identity":
250-
// This is *only* done for testing to mock client authentication.
251-
// This will never be set in a production scenario.
252-
var awsClient *http.Client
253-
awsClientRaw := ctx.Value("aws-client")
254-
if awsClientRaw != nil {
255-
awsClient, _ = awsClientRaw.(*http.Client)
256-
if awsClient != nil {
257-
client.SDK.HTTPClient = awsClient
258-
}
259-
}
260-
exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
261-
return client.AuthAWSInstanceIdentity(ctx)
262-
}
263-
case "azure-instance-identity":
264-
// This is *only* done for testing to mock client authentication.
265-
// This will never be set in a production scenario.
266-
var azureClient *http.Client
267-
azureClientRaw := ctx.Value("azure-client")
268-
if azureClientRaw != nil {
269-
azureClient, _ = azureClientRaw.(*http.Client)
270-
if azureClient != nil {
271-
client.SDK.HTTPClient = azureClient
272-
}
273-
}
274-
exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
275-
return client.AuthAzureInstanceIdentity(ctx)
276-
}
277-
}
278-
279214
executablePath, err := os.Executable()
280215
if err != nil {
281216
return xerrors.Errorf("getting os executable: %w", err)
@@ -343,18 +278,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
343278
LogDir: logDir,
344279
ScriptDataDir: scriptDataDir,
345280
// #nosec G115 - Safe conversion as tailnet listen port is within uint16 range (0-65535)
346-
TailnetListenPort: uint16(tailnetListenPort),
347-
ExchangeToken: func(ctx context.Context) (string, error) {
348-
if exchangeToken == nil {
349-
return client.SDK.SessionToken(), nil
350-
}
351-
resp, err := exchangeToken(ctx)
352-
if err != nil {
353-
return "", err
354-
}
355-
client.SetSessionToken(resp.SessionToken)
356-
return resp.SessionToken, nil
357-
},
281+
TailnetListenPort: uint16(tailnetListenPort),
358282
EnvironmentVariables: environmentVariables,
359283
IgnorePorts: ignorePorts,
360284
SSHMaxTimeout: sshMaxTimeout,
@@ -400,13 +324,6 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
400324
}
401325

402326
cmd.Options = serpent.OptionSet{
403-
{
404-
Flag: "auth",
405-
Default: "token",
406-
Description: "Specify the authentication type to use for the agent.",
407-
Env: "CODER_AGENT_AUTH",
408-
Value: serpent.StringOf(&auth),
409-
},
410327
{
411328
Flag: "log-dir",
412329
Default: os.TempDir(),

‎cli/exp_mcp.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
148148
binPath = testBinaryName
149149
}
150150
configureClaudeEnv := map[string]string{}
151-
agentClient, err := r.createAgentClient()
151+
agentClient, err := r.createAgentClient(inv.Context())
152152
if err != nil {
153153
cliui.Warnf(inv.Stderr, "failed to create agent client: %s", err)
154154
} else {
@@ -494,7 +494,7 @@ func (r *RootCmd) mcpServer() *serpent.Command {
494494
}
495495

496496
// Try to create an agent client for status reporting. Not validated.
497-
agentClient, err := r.createAgentClient()
497+
agentClient, err := r.createAgentClient(inv.Context())
498498
if err == nil {
499499
cliui.Infof(inv.Stderr, "Agent URL : %s", agentClient.SDK.URL.String())
500500
srv.agentClient = agentClient

‎cli/externalauth.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ fi
7575
return xerrors.Errorf("agent token not found")
7676
}
7777

78-
client, err := r.tryCreateAgentClient()
78+
client, err := r.createAgentClient(ctx)
7979
if err != nil {
8080
return xerrors.Errorf("create agent client: %w", err)
8181
}

‎cli/gitaskpass.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func (r *RootCmd) gitAskpass() *serpent.Command {
3333
return xerrors.Errorf("parse host: %w", err)
3434
}
3535

36-
client, err := r.tryCreateAgentClient()
36+
client, err := r.createAgentClient(ctx)
3737
if err != nil {
3838
return xerrors.Errorf("create agent client: %w", err)
3939
}

‎cli/gitaskpass_test.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/coder/coder/v2/codersdk"
1717
"github.com/coder/coder/v2/codersdk/agentsdk"
1818
"github.com/coder/coder/v2/pty/ptytest"
19+
"github.com/coder/coder/v2/testutil"
1920
)
2021

2122
func TestGitAskpass(t *testing.T) {
@@ -65,6 +66,7 @@ func TestGitAskpass(t *testing.T) {
6566

6667
t.Run("Poll", func(t *testing.T) {
6768
t.Parallel()
69+
ctx := testutil.Context(t, testutil.WaitShort)
6870
resp := atomic.Pointer[agentsdk.ExternalAuthResponse]{}
6971
resp.Store(&agentsdk.ExternalAuthResponse{
7072
URL: "https://something.org",
@@ -86,6 +88,7 @@ func TestGitAskpass(t *testing.T) {
8688

8789
inv, _ := clitest.New(t, "--agent-url", url, "--no-open", "Username for 'https://github.com':")
8890
inv.Environ.Set("GIT_PREFIX", "/")
91+
inv.Environ.Set("CODER_AGENT_TOKEN", "fake-token")
8992
stdout := ptytest.New(t)
9093
inv.Stdout = stdout.Output()
9194
stderr := ptytest.New(t)
@@ -94,7 +97,7 @@ func TestGitAskpass(t *testing.T) {
9497
err := inv.Run()
9598
assert.NoError(t, err)
9699
}()
97-
<-poll
100+
testutil.RequireReceive(ctx, t, poll)
98101
stderr.ExpectMatch("Open the following URL to authenticate")
99102
resp.Store(&agentsdk.ExternalAuthResponse{
100103
Username: "username",

‎cli/gitssh.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func (r *RootCmd) gitssh() *serpent.Command {
3838
return err
3939
}
4040

41-
client, err := r.tryCreateAgentClient()
41+
client, err := r.createAgentClient(ctx)
4242
if err != nil {
4343
return xerrors.Errorf("create agent client: %w", err)
4444
}

‎cli/gitssh_test.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,7 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str
5454
}).WithAgent().Do()
5555

5656
// start workspace agent
57-
agentClient := agentsdk.New(client.URL)
58-
agentClient.SetSessionToken(r.AgentToken)
57+
agentClient := agentsdk.New(client.URL, agentsdk.UsingFixedToken(r.AgentToken))
5958
_ = agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
6059
o.Client = agentClient
6160
})

0 commit comments

Comments
 (0)