Skip to content

Commit 724efff

Browse files
mafredrimafredri
committed
feat: add agentapi endpoint to report connections for audit
This change adds a new `ReportConnection` endpoint to the `agentapi` and bumps the protocol version. This allows the agent to report connection events, for example when the user connects to the workspace via SSH or VS Code. Updates #15139
1 parent de017dc commit 724efff

File tree

13 files changed

+1126
-587
lines changed

13 files changed

+1126
-587
lines changed

‎agent/agent.go

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,8 @@ type Options struct {
8585
}
8686

8787
type Client interface {
88-
ConnectRPC23(ctx context.Context) (
89-
proto.DRPCAgentClient23, tailnetproto.DRPCTailnetClient23, error,
88+
ConnectRPC24(ctx context.Context) (
89+
proto.DRPCAgentClient24, tailnetproto.DRPCTailnetClient24, error,
9090
)
9191
RewriteDERPMap(derpMap *tailcfg.DERPMap)
9292
}
@@ -362,7 +362,6 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
362362
// Important: if the command times out, we may see a misleading error like
363363
// "exit status 1", so it's important to include the context error.
364364
err = errors.Join(err, ctx.Err())
365-
366365
if err != nil {
367366
result.Error = fmt.Sprintf("run cmd: %+v", err)
368367
}
@@ -399,7 +398,7 @@ func (t *trySingleflight) Do(key string, fn func()) {
399398
fn()
400399
}
401400

402-
func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
401+
func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
403402
tickerDone := make(chan struct{})
404403
collectDone := make(chan struct{})
405404
ctx, cancel := context.WithCancel(ctx)
@@ -615,7 +614,7 @@ func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient23
615614

616615
// reportLifecycle reports the current lifecycle state once. All state
617616
// changes are reported in order.
618-
func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
617+
func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
619618
for {
620619
select {
621620
case <-a.lifecycleUpdate:
@@ -697,7 +696,7 @@ func (a *agent) setLifecycle(state codersdk.WorkspaceAgentLifecycle) {
697696
// fetchServiceBannerLoop fetches the service banner on an interval. It will
698697
// not be fetched immediately; the expectation is that it is primed elsewhere
699698
// (and must be done before the session actually starts).
700-
func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
699+
func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
701700
ticker := time.NewTicker(a.announcementBannersRefreshInterval)
702701
defer ticker.Stop()
703702
for {
@@ -733,7 +732,7 @@ func (a *agent) run() (retErr error) {
733732
a.sessionToken.Store(&sessionToken)
734733

735734
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs
736-
aAPI, tAPI, err := a.client.ConnectRPC23(a.hardCtx)
735+
aAPI, tAPI, err := a.client.ConnectRPC24(a.hardCtx)
737736
if err != nil {
738737
return err
739738
}
@@ -750,7 +749,7 @@ func (a *agent) run() (retErr error) {
750749
connMan := newAPIConnRoutineManager(a.gracefulCtx, a.hardCtx, a.logger, aAPI, tAPI)
751750

752751
connMan.startAgentAPI("init notification banners", gracefulShutdownBehaviorStop,
753-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
752+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
754753
bannersProto, err := aAPI.GetAnnouncementBanners(ctx, &proto.GetAnnouncementBannersRequest{})
755754
if err != nil {
756755
return xerrors.Errorf("fetch service banner: %w", err)
@@ -767,7 +766,7 @@ func (a *agent) run() (retErr error) {
767766
// sending logs gets gracefulShutdownBehaviorRemain because we want to send logs generated by
768767
// shutdown scripts.
769768
connMan.startAgentAPI("send logs", gracefulShutdownBehaviorRemain,
770-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
769+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
771770
err := a.logSender.SendLoop(ctx, aAPI)
772771
if xerrors.Is(err, agentsdk.LogLimitExceededError) {
773772
// we don't want this error to tear down the API connection and propagate to the
@@ -807,7 +806,7 @@ func (a *agent) run() (retErr error) {
807806
connMan.startAgentAPI("handle manifest", gracefulShutdownBehaviorStop, a.handleManifest(manifestOK))
808807

809808
connMan.startAgentAPI("app health reporter", gracefulShutdownBehaviorStop,
810-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
809+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
811810
if err := manifestOK.wait(ctx); err != nil {
812811
return xerrors.Errorf("no manifest: %w", err)
813812
}
@@ -840,7 +839,7 @@ func (a *agent) run() (retErr error) {
840839

841840
connMan.startAgentAPI("fetch service banner loop", gracefulShutdownBehaviorStop, a.fetchServiceBannerLoop)
842841

843-
connMan.startAgentAPI("stats report loop", gracefulShutdownBehaviorStop, func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
842+
connMan.startAgentAPI("stats report loop", gracefulShutdownBehaviorStop, func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
844843
if err := networkOK.wait(ctx); err != nil {
845844
return xerrors.Errorf("no network: %w", err)
846845
}
@@ -851,8 +850,8 @@ func (a *agent) run() (retErr error) {
851850
}
852851

853852
// handleManifest returns a function that fetches and processes the manifest
854-
func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
855-
return func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
853+
func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
854+
return func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
856855
var (
857856
sentResult = false
858857
err error
@@ -961,8 +960,8 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context,
961960

962961
// createOrUpdateNetwork waits for the manifest to be set using manifestOK, then creates or updates
963962
// the tailnet using the information in the manifest
964-
func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(context.Context, proto.DRPCAgentClient23) error {
965-
return func(ctx context.Context, _ proto.DRPCAgentClient23) (retErr error) {
963+
func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(context.Context, proto.DRPCAgentClient24) error {
964+
return func(ctx context.Context, _ proto.DRPCAgentClient24) (retErr error) {
966965
if err := manifestOK.wait(ctx); err != nil {
967966
return xerrors.Errorf("no manifest: %w", err)
968967
}
@@ -1683,7 +1682,7 @@ const (
16831682

16841683
type apiConnRoutineManager struct {
16851684
logger slog.Logger
1686-
aAPI proto.DRPCAgentClient23
1685+
aAPI proto.DRPCAgentClient24
16871686
tAPI tailnetproto.DRPCTailnetClient23
16881687
eg *errgroup.Group
16891688
stopCtx context.Context
@@ -1692,7 +1691,7 @@ type apiConnRoutineManager struct {
16921691

16931692
func newAPIConnRoutineManager(
16941693
gracefulCtx, hardCtx context.Context, logger slog.Logger,
1695-
aAPI proto.DRPCAgentClient23, tAPI tailnetproto.DRPCTailnetClient23,
1694+
aAPI proto.DRPCAgentClient24, tAPI tailnetproto.DRPCTailnetClient23,
16961695
) *apiConnRoutineManager {
16971696
// routines that remain in operation during graceful shutdown use the remainCtx. They'll still
16981697
// exit if the errgroup hits an error, which usually means a problem with the conn.
@@ -1725,7 +1724,7 @@ func newAPIConnRoutineManager(
17251724
// but for Tailnet.
17261725
func (a *apiConnRoutineManager) startAgentAPI(
17271726
name string, behavior gracefulShutdownBehavior,
1728-
f func(context.Context, proto.DRPCAgentClient23) error,
1727+
f func(context.Context, proto.DRPCAgentClient24) error,
17291728
) {
17301729
logger := a.logger.With(slog.F("name", name))
17311730
var ctx context.Context

‎agent/agenttest/client.go

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"golang.org/x/exp/slices"
1616
"golang.org/x/xerrors"
1717
"google.golang.org/protobuf/types/known/durationpb"
18+
"google.golang.org/protobuf/types/known/emptypb"
1819
"storj.io/drpc/drpcmux"
1920
"storj.io/drpc/drpcserver"
2021
"tailscale.com/tailcfg"
@@ -96,8 +97,8 @@ func (c *Client) Close() {
9697
c.derpMapOnce.Do(func() { close(c.derpMapUpdates) })
9798
}
9899

99-
func (c *Client) ConnectRPC23(ctx context.Context) (
100-
agentproto.DRPCAgentClient23, proto.DRPCTailnetClient23, error,
100+
func (c *Client) ConnectRPC24(ctx context.Context) (
101+
agentproto.DRPCAgentClient24, proto.DRPCTailnetClient24, error,
101102
) {
102103
conn, lis := drpcsdk.MemTransportPipe()
103104
c.LastWorkspaceAgent = func() {
@@ -170,6 +171,7 @@ type FakeAgentAPI struct {
170171
lifecycleStates []codersdk.WorkspaceAgentLifecycle
171172
metadata map[string]agentsdk.Metadata
172173
timings []*agentproto.Timing
174+
connections []*agentproto.Connection
173175

174176
getAnnouncementBannersFunc func() ([]codersdk.BannerConfig, error)
175177
}
@@ -309,12 +311,20 @@ func (f *FakeAgentAPI) BatchCreateLogs(ctx context.Context, req *agentproto.Batc
309311

310312
func (f *FakeAgentAPI) ScriptCompleted(_ context.Context, req *agentproto.WorkspaceAgentScriptCompletedRequest) (*agentproto.WorkspaceAgentScriptCompletedResponse, error) {
311313
f.Lock()
312-
f.timings = append(f.timings, req.Timing)
314+
f.timings = append(f.timings, req.GetTiming())
313315
f.Unlock()
314316

315317
return &agentproto.WorkspaceAgentScriptCompletedResponse{}, nil
316318
}
317319

320+
func (f *FakeAgentAPI) ReportConnection(_ context.Context, req *agentproto.ReportConnectionRequest) (*emptypb.Empty, error) {
321+
f.Lock()
322+
f.connections = append(f.connections, req.GetConnection())
323+
f.Unlock()
324+
325+
return &emptypb.Empty{}, nil
326+
}
327+
318328
func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest, statsCh chan *agentproto.Stats) *FakeAgentAPI {
319329
return &FakeAgentAPI{
320330
t: t,

0 commit comments

Comments
 (0)