Thanks to visit codestin.com
Credit goes to github.com

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

+17-18
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

+13-3
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)