opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: refactor instance identity to be a SessionTokenProvider (#19566)

Browse Source 
Refactors Agent instance identity to be a SessionTokenProvider.

Refactors the CLI to create Agent clients via a centralized function, rather than add-hoc via individual command handlers and their flags.

This allows commands besides `coder agent`, but which still use the agent identity, to support instance identity authentication.

Fixes #19111 by unifying all API requests to go thru the SessionTokenProvider for auth credentials.
This commit is contained in:
Spike Curtis
2025-09-03 10:38:42 +04:00
committed by GitHub
parent ee35ad3a57
commit 1354d84eb4
35 changed files with 640 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/agent.go
+4 -14
View File
@@ -74,7 +74,6 @@ type Options struct {
LogDir string
TempDir string
ScriptDataDir string
ExchangeToken func(ctx context.Context) (string, error)
Client Client
ReconnectingPTYTimeout time.Duration
EnvironmentVariables map[string]string
@@ -99,6 +98,7 @@ type Client interface {
proto.DRPCAgentClient26, tailnetproto.DRPCTailnetClient26, error,
)
tailnet.DERPMapRewriter
agentsdk.RefreshableSessionTokenProvider
}
type Agent interface {
@@ -131,11 +131,6 @@ func New(options Options) Agent {
}
options.ScriptDataDir = options.TempDir
}
if options.ExchangeToken == nil {
options.ExchangeToken = func(_ context.Context) (string, error) {
return "", nil
}
}
if options.ReportMetadataInterval == 0 {
options.ReportMetadataInterval = time.Second
}
@@ -172,7 +167,6 @@ func New(options Options) Agent {
coordDisconnected: make(chan struct{}),
environmentVariables: options.EnvironmentVariables,
client: options.Client,
exchangeToken: options.ExchangeToken,
filesystem: options.Filesystem,
logDir: options.LogDir,
tempDir: options.TempDir,
@@ -203,7 +197,6 @@ func New(options Options) Agent {
// coordinator during shut down.
close(a.coordDisconnected)
a.announcementBanners.Store(new([]codersdk.BannerConfig))
a.sessionToken.Store(new(string))
a.init()
return a
}
@@ -212,7 +205,6 @@ type agent struct {
clock quartz.Clock
logger slog.Logger
client Client
exchangeToken func(ctx context.Context) (string, error)
tailnetListenPort uint16
filesystem afero.Fs
logDir string
@@ -254,7 +246,6 @@ type agent struct {
scriptRunner *agentscripts.Runner
announcementBanners atomic.Pointer[[]codersdk.BannerConfig] // announcementBanners is atomic because it is periodically updated.
announcementBannersRefreshInterval time.Duration
sessionToken atomic.Pointer[string]
sshServer *agentssh.Server
sshMaxTimeout time.Duration
blockFileTransfer bool
@@ -916,11 +907,10 @@ func (a *agent) run() (retErr error) {
// This allows the agent to refresh its token if necessary.
// For instance identity this is required, since the instance
// may not have re-provisioned, but a new agent ID was created.
sessionToken, err := a.exchangeToken(a.hardCtx)
err := a.client.RefreshToken(a.hardCtx)
if err != nil {
return xerrors.Errorf("exchange token: %w", err)
return xerrors.Errorf("refresh token: %w", err)
}
a.sessionToken.Store(&sessionToken)
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs
aAPI, tAPI, err := a.client.ConnectRPC26(a.hardCtx)
@@ -1359,7 +1349,7 @@ func (a *agent) updateCommandEnv(current []string) (updated []string, err error)
"CODER_WORKSPACE_OWNER_NAME": manifest.OwnerName,
// Specific Coder subcommands require the agent token exposed!
"CODER_AGENT_TOKEN": *a.sessionToken.Load(),
"CODER_AGENT_TOKEN": a.client.GetSessionToken(),
// Git on Windows resolves with UNIX-style paths.
// If using backslashes, it's unable to find the executable.
agent/agent_test.go
+12 -22
View File
@@ -22,7 +22,6 @@ import (
"slices"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
@@ -2926,11 +2925,11 @@ func TestAgent_Speedtest(t *testing.T) {
func TestAgent_Reconnect(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
logger := testutil.Logger(t)
// After the agent is disconnected from a coordinator, it's supposed
// to reconnect!
coordinator := tailnet.NewCoordinator(logger)
defer coordinator.Close()
fCoordinator := tailnettest.NewFakeCoordinator()
agentID := uuid.New()
statsCh := make(chan *proto.Stats, 50)
@@ -2942,27 +2941,24 @@ func TestAgent_Reconnect(t *testing.T) {
DERPMap: derpMap,
},
statsCh,
coordinator,
fCoordinator,
)
defer client.Close()
initialized := atomic.Int32{}
closer := agent.New(agent.Options{
ExchangeToken: func(ctx context.Context) (string, error) {
initialized.Add(1)
return "", nil
},
Client: client,
Logger: logger.Named("agent"),
})
defer closer.Close()
require.Eventually(t, func() bool {
return coordinator.Node(agentID) != nil
}, testutil.WaitShort, testutil.IntervalFast)
client.LastWorkspaceAgent()
require.Eventually(t, func() bool {
return initialized.Load() == 2
}, testutil.WaitShort, testutil.IntervalFast)
call1 := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
require.Equal(t, client.GetNumRefreshTokenCalls(), 1)
close(call1.Resps) // hang up
// expect reconnect
testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
// Check that the agent refreshes the token when it reconnects.
require.Equal(t, client.GetNumRefreshTokenCalls(), 2)
closer.Close()
}
func TestAgent_WriteVSCodeConfigs(t *testing.T) {
@@ -2984,9 +2980,6 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
defer client.Close()
filesystem := afero.NewMemMapFs()
closer := agent.New(agent.Options{
ExchangeToken: func(ctx context.Context) (string, error) {
return "", nil
},
Client: client,
Logger: logger.Named("agent"),
Filesystem: filesystem,
@@ -3015,9 +3008,6 @@ func TestAgent_DebugServer(t *testing.T) {
conn, _, _, _, agnt := setupAgent(t, agentsdk.Manifest{
DERPMap: derpMap,
}, 0, func(c *agenttest.Client, o *agent.Options) {
o.ExchangeToken = func(context.Context) (string, error) {
return "token", nil
}
o.LogDir = logDir
})
agent/agenttest/agent.go
+1 -9
View File
@@ -1,7 +1,6 @@
package agenttest
import (
"context"
"net/url"
"testing"
@@ -31,18 +30,11 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent
}
if o.Client == nil {
agentClient := agentsdk.New(coderURL)
agentClient.SetSessionToken(agentToken)
agentClient := agentsdk.New(coderURL, agentsdk.WithFixedToken(agentToken))
agentClient.SDK.SetLogger(log)
o.Client = agentClient
}
if o.ExchangeToken == nil {
o.ExchangeToken = func(_ context.Context) (string, error) {
return agentToken, nil
}
}
if o.LogDir == "" {
o.LogDir = t.TempDir()
}
agent/agenttest/client.go
+30 -4
View File
@@ -3,6 +3,7 @@ package agenttest
import (
"context"
"io"
"net/http"
"slices"
"sync"
"sync/atomic"
@@ -28,6 +29,7 @@ import (
"github.com/coder/coder/v2/tailnet"
"github.com/coder/coder/v2/tailnet/proto"
"github.com/coder/coder/v2/testutil"
"github.com/coder/websocket"
)
const statsInterval = 500 * time.Millisecond
@@ -86,10 +88,34 @@ type Client struct {
fakeAgentAPI *FakeAgentAPI
LastWorkspaceAgent func()
mu sync.Mutex // Protects following.
logs []agentsdk.Log
derpMapUpdates chan *tailcfg.DERPMap
derpMapOnce sync.Once
mu sync.Mutex // Protects following.
logs []agentsdk.Log
derpMapUpdates chan *tailcfg.DERPMap
derpMapOnce sync.Once
refreshTokenCalls int
}
func (*Client) AsRequestOption() codersdk.RequestOption {
return func(_ *http.Request) {}
}
func (*Client) SetDialOption(*websocket.DialOptions) {}
func (*Client) GetSessionToken() string {
return "agenttest-token"
}
func (c *Client) RefreshToken(context.Context) error {
c.mu.Lock()
defer c.mu.Unlock()
c.refreshTokenCalls++
return nil
}
func (c *Client) GetNumRefreshTokenCalls() int {
c.mu.Lock()
defer c.mu.Unlock()
return c.refreshTokenCalls
}
func (*Client) RewriteDERPMap(*tailcfg.DERPMap) {}