Skip to content

Commit 761ba0a

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 ffc8fd6 commit 761ba0a

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
@@ -87,8 +87,8 @@ type Options struct {
8787
}
8888

8989
type Client interface {
90-
ConnectRPC23(ctx context.Context) (
91-
proto.DRPCAgentClient23, tailnetproto.DRPCTailnetClient23, error,
90+
ConnectRPC24(ctx context.Context) (
91+
proto.DRPCAgentClient24, tailnetproto.DRPCTailnetClient24, error,
9292
)
9393
RewriteDERPMap(derpMap *tailcfg.DERPMap)
9494
}
@@ -369,7 +369,6 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
369369
// Important: if the command times out, we may see a misleading error like
370370
// "exit status 1", so it's important to include the context error.
371371
err = errors.Join(err, ctx.Err())
372-
373372
if err != nil {
374373
result.Error = fmt.Sprintf("run cmd: %+v", err)
375374
}
@@ -406,7 +405,7 @@ func (t *trySingleflight) Do(key string, fn func()) {
406405
fn()
407406
}
408407

409-
func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
408+
func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
410409
tickerDone := make(chan struct{})
411410
collectDone := make(chan struct{})
412411
ctx, cancel := context.WithCancel(ctx)
@@ -622,7 +621,7 @@ func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient23
622621

623622
// reportLifecycle reports the current lifecycle state once. All state
624623
// changes are reported in order.
625-
func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
624+
func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
626625
for {
627626
select {
628627
case <-a.lifecycleUpdate:
@@ -704,7 +703,7 @@ func (a *agent) setLifecycle(state codersdk.WorkspaceAgentLifecycle) {
704703
// fetchServiceBannerLoop fetches the service banner on an interval. It will
705704
// not be fetched immediately; the expectation is that it is primed elsewhere
706705
// (and must be done before the session actually starts).
707-
func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
706+
func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
708707
ticker := time.NewTicker(a.announcementBannersRefreshInterval)
709708
defer ticker.Stop()
710709
for {
@@ -740,7 +739,7 @@ func (a *agent) run() (retErr error) {
740739
a.sessionToken.Store(&sessionToken)
741740

742741
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs
743-
aAPI, tAPI, err := a.client.ConnectRPC23(a.hardCtx)
742+
aAPI, tAPI, err := a.client.ConnectRPC24(a.hardCtx)
744743
if err != nil {
745744
return err
746745
}
@@ -757,7 +756,7 @@ func (a *agent) run() (retErr error) {
757756
connMan := newAPIConnRoutineManager(a.gracefulCtx, a.hardCtx, a.logger, aAPI, tAPI)
758757

759758
connMan.startAgentAPI("init notification banners", gracefulShutdownBehaviorStop,
760-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
759+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
761760
bannersProto, err := aAPI.GetAnnouncementBanners(ctx, &proto.GetAnnouncementBannersRequest{})
762761
if err != nil {
763762
return xerrors.Errorf("fetch service banner: %w", err)
@@ -774,7 +773,7 @@ func (a *agent) run() (retErr error) {
774773
// sending logs gets gracefulShutdownBehaviorRemain because we want to send logs generated by
775774
// shutdown scripts.
776775
connMan.startAgentAPI("send logs", gracefulShutdownBehaviorRemain,
777-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
776+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
778777
err := a.logSender.SendLoop(ctx, aAPI)
779778
if xerrors.Is(err, agentsdk.LogLimitExceededError) {
780779
// we don't want this error to tear down the API connection and propagate to the
@@ -814,7 +813,7 @@ func (a *agent) run() (retErr error) {
814813
connMan.startAgentAPI("handle manifest", gracefulShutdownBehaviorStop, a.handleManifest(manifestOK))
815814

816815
connMan.startAgentAPI("app health reporter", gracefulShutdownBehaviorStop,
817-
func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
816+
func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
818817
if err := manifestOK.wait(ctx); err != nil {
819818
return xerrors.Errorf("no manifest: %w", err)
820819
}
@@ -847,7 +846,7 @@ func (a *agent) run() (retErr error) {
847846

848847
connMan.startAgentAPI("fetch service banner loop", gracefulShutdownBehaviorStop, a.fetchServiceBannerLoop)
849848

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

860859
// handleManifest returns a function that fetches and processes the manifest
861-
func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
862-
return func(ctx context.Context, aAPI proto.DRPCAgentClient23) error {
860+
func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
861+
return func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {
863862
var (
864863
sentResult = false
865864
err error
@@ -968,8 +967,8 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context,
968967

969968
// createOrUpdateNetwork waits for the manifest to be set using manifestOK, then creates or updates
970969
// the tailnet using the information in the manifest
971-
func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(context.Context, proto.DRPCAgentClient23) error {
972-
return func(ctx context.Context, _ proto.DRPCAgentClient23) (retErr error) {
970+
func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(context.Context, proto.DRPCAgentClient24) error {
971+
return func(ctx context.Context, _ proto.DRPCAgentClient24) (retErr error) {
973972
if err := manifestOK.wait(ctx); err != nil {
974973
return xerrors.Errorf("no manifest: %w", err)
975974
}
@@ -1690,7 +1689,7 @@ const (
16901689

16911690
type apiConnRoutineManager struct {
16921691
logger slog.Logger
1693-
aAPI proto.DRPCAgentClient23
1692+
aAPI proto.DRPCAgentClient24
16941693
tAPI tailnetproto.DRPCTailnetClient23
16951694
eg *errgroup.Group
16961695
stopCtx context.Context
@@ -1699,7 +1698,7 @@ type apiConnRoutineManager struct {
16991698

17001699
func newAPIConnRoutineManager(
17011700
gracefulCtx, hardCtx context.Context, logger slog.Logger,
1702-
aAPI proto.DRPCAgentClient23, tAPI tailnetproto.DRPCTailnetClient23,
1701+
aAPI proto.DRPCAgentClient24, tAPI tailnetproto.DRPCTailnetClient23,
17031702
) *apiConnRoutineManager {
17041703
// routines that remain in operation during graceful shutdown use the remainCtx. They'll still
17051704
// exit if the errgroup hits an error, which usually means a problem with the conn.
@@ -1732,7 +1731,7 @@ func newAPIConnRoutineManager(
17321731
// but for Tailnet.
17331732
func (a *apiConnRoutineManager) startAgentAPI(
17341733
name string, behavior gracefulShutdownBehavior,
1735-
f func(context.Context, proto.DRPCAgentClient23) error,
1734+
f func(context.Context, proto.DRPCAgentClient24) error,
17361735
) {
17371736
logger := a.logger.With(slog.F("name", name))
17381737
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy