package agent

import (

"bytes"

"context"

"encoding/json"

"errors"

"fmt"

"hash/fnv"

"io"

"net"

"net/http"

"net/netip"

"os"

"os/user"

"path/filepath"

"slices"

"sort"

"strconv"

"strings"

"sync"

"time"

"github.com/go-chi/chi/v5"

"github.com/google/uuid"

"github.com/prometheus/client_golang/prometheus"

"github.com/prometheus/common/expfmt"

"github.com/spf13/afero"

"go.uber.org/atomic"

"golang.org/x/sync/errgroup"

"golang.org/x/xerrors"

"google.golang.org/protobuf/types/known/timestamppb"

"tailscale.com/net/speedtest"

"tailscale.com/tailcfg"

"tailscale.com/types/netlogtype"

"tailscale.com/util/clientmetric"

"cdr.dev/slog"

"github.com/coder/clistat"

"github.com/coder/coder/v2/agent/agentcontainers"

"github.com/coder/coder/v2/agent/agentexec"

"github.com/coder/coder/v2/agent/agentscripts"

"github.com/coder/coder/v2/agent/agentssh"

"github.com/coder/coder/v2/agent/proto"

"github.com/coder/coder/v2/agent/proto/resourcesmonitor"

"github.com/coder/coder/v2/agent/reconnectingpty"

"github.com/coder/coder/v2/buildinfo"

"github.com/coder/coder/v2/cli/gitauth"

"github.com/coder/coder/v2/coderd/database/dbtime"

"github.com/coder/coder/v2/codersdk"

"github.com/coder/coder/v2/codersdk/agentsdk"

"github.com/coder/coder/v2/codersdk/workspacesdk"

"github.com/coder/coder/v2/tailnet"

tailnetproto "github.com/coder/coder/v2/tailnet/proto"

"github.com/coder/quartz"

"github.com/coder/retry"

)

const (

ProtocolReconnectingPTY = "reconnecting-pty"

ProtocolSSH = "ssh"

ProtocolDial = "dial"

)

// EnvProcPrioMgmt determines whether we attempt to manage

// process CPU and OOM Killer priority.

const (

EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"

EnvProcOOMScore = "CODER_PROC_OOM_SCORE"

)

type Options struct {

Filesystem afero.Fs

LogDir string

TempDir string

ScriptDataDir string

ExchangeToken func(ctx context.Context) (string, error)

Client Client

ReconnectingPTYTimeout time.Duration

EnvironmentVariables map[string]string

Logger slog.Logger

IgnorePorts map[int]string

PortCacheDuration time.Duration

SSHMaxTimeout time.Duration

TailnetListenPort uint16

Subsystems []codersdk.AgentSubsystem

PrometheusRegistry *prometheus.Registry

ReportMetadataInterval time.Duration

ServiceBannerRefreshInterval time.Duration

BlockFileTransfer bool

Execer agentexec.Execer

SubAgent bool

ExperimentalDevcontainersEnabled bool

ContainerAPIOptions []agentcontainers.Option // Enable ExperimentalDevcontainersEnabled for these to be effective.

}

type Client interface {

ConnectRPC24(ctx context.Context) (

proto.DRPCAgentClient24, tailnetproto.DRPCTailnetClient24, error,

)

RewriteDERPMap(derpMap *tailcfg.DERPMap)

}

type Agent interface {

HTTPDebug() http.Handler

// TailnetConn may be nil.

TailnetConn() *tailnet.Conn

io.Closer

}

func New(options Options) Agent {

if options.Filesystem == nil {

options.Filesystem = afero.NewOsFs()

}

if options.TempDir == "" {

options.TempDir = os.TempDir()

}

if options.LogDir == "" {

if options.TempDir != os.TempDir() {

options.Logger.Debug(context.Background(), "log dir not set, using temp dir", slog.F("temp_dir", options.TempDir))

} else {

options.Logger.Debug(context.Background(), "using log dir", slog.F("log_dir", options.LogDir))

}

options.LogDir = options.TempDir

}

if options.ScriptDataDir == "" {

if options.TempDir != os.TempDir() {

options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir))

} else {

options.Logger.Debug(context.Background(), "using script data dir", slog.F("script_data_dir", options.ScriptDataDir))

}

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

}

if options.ServiceBannerRefreshInterval == 0 {

options.ServiceBannerRefreshInterval = 2 * time.Minute

}

if options.PortCacheDuration == 0 {

options.PortCacheDuration = 1 * time.Second

}

prometheusRegistry := options.PrometheusRegistry

if prometheusRegistry == nil {

prometheusRegistry = prometheus.NewRegistry()

}

if options.Execer == nil {

options.Execer = agentexec.DefaultExecer

}

hardCtx, hardCancel := context.WithCancel(context.Background())

gracefulCtx, gracefulCancel := context.WithCancel(hardCtx)

a := &agent{

tailnetListenPort: options.TailnetListenPort,

reconnectingPTYTimeout: options.ReconnectingPTYTimeout,

logger: options.Logger,

gracefulCtx: gracefulCtx,

gracefulCancel: gracefulCancel,

hardCtx: hardCtx,

hardCancel: hardCancel,

coordDisconnected: make(chan struct{}),

environmentVariables: options.EnvironmentVariables,

client: options.Client,

exchangeToken: options.ExchangeToken,

filesystem: options.Filesystem,

logDir: options.LogDir,

tempDir: options.TempDir,

scriptDataDir: options.ScriptDataDir,

lifecycleUpdate: make(chan struct{}, 1),

lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),

lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},

reportConnectionsUpdate: make(chan struct{}, 1),

ignorePorts: options.IgnorePorts,

portCacheDuration: options.PortCacheDuration,

reportMetadataInterval: options.ReportMetadataInterval,

announcementBannersRefreshInterval: options.ServiceBannerRefreshInterval,

sshMaxTimeout: options.SSHMaxTimeout,

subsystems: options.Subsystems,

logSender: agentsdk.NewLogSender(options.Logger),

blockFileTransfer: options.BlockFileTransfer,

prometheusRegistry: prometheusRegistry,

metrics: newAgentMetrics(prometheusRegistry),

execer: options.Execer,

subAgent: options.SubAgent,

experimentalDevcontainersEnabled: options.ExperimentalDevcontainersEnabled,

containerAPIOptions: options.ContainerAPIOptions,

}

// Initially, we have a closed channel, reflecting the fact that we are not initially connected.

// Each time we connect we replace the channel (while holding the closeMutex) with a new one

// that gets closed on disconnection. This is used to wait for graceful disconnection from the

// coordinator during shut down.

close(a.coordDisconnected)

a.announcementBanners.Store(new([]codersdk.BannerConfig))

a.sessionToken.Store(new(string))

a.init()

return a

}

type agent struct {

logger slog.Logger

client Client

exchangeToken func(ctx context.Context) (string, error)

tailnetListenPort uint16

filesystem afero.Fs

logDir string

tempDir string

scriptDataDir string

// ignorePorts tells the api handler which ports to ignore when

// listing all listening ports. This is helpful to hide ports that

// are used by the agent, that the user does not care about.

ignorePorts map[int]string

portCacheDuration time.Duration

subsystems []codersdk.AgentSubsystem

reconnectingPTYTimeout time.Duration

reconnectingPTYServer *reconnectingpty.Server

// we track 2 contexts and associated cancel functions: "graceful" which is Done when it is time

// to start gracefully shutting down and "hard" which is Done when it is time to close

// everything down (regardless of whether graceful shutdown completed).

gracefulCtx context.Context

gracefulCancel context.CancelFunc

hardCtx context.Context

hardCancel context.CancelFunc

// closeMutex protects the following:

closeMutex sync.Mutex

closeWaitGroup sync.WaitGroup

coordDisconnected chan struct{}

closing bool

// note that once the network is set to non-nil, it is never modified, as with the statsReporter. So, routines

// that run after createOrUpdateNetwork and check the networkOK checkpoint do not need to hold the lock to use them.

network *tailnet.Conn

statsReporter *statsReporter

// end fields protected by closeMutex

environmentVariables map[string]string

manifest atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.

reportMetadataInterval time.Duration

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

lifecycleUpdate chan struct{}

lifecycleReported chan codersdk.WorkspaceAgentLifecycle

lifecycleMu sync.RWMutex // Protects following.

lifecycleStates []agentsdk.PostLifecycleRequest

lifecycleLastReportedIndex int // Keeps track of the last lifecycle state we successfully reported.

reportConnectionsUpdate chan struct{}

reportConnectionsMu sync.Mutex

reportConnections []*proto.ReportConnectionRequest

logSender *agentsdk.LogSender

prometheusRegistry *prometheus.Registry

// metrics are prometheus registered metrics that will be collected and

// labeled in Coder with the agent + workspace.

metrics *agentMetrics

execer agentexec.Execer

subAgent bool

experimentalDevcontainersEnabled bool

containerAPIOptions []agentcontainers.Option

containerAPI atomic.Pointer[agentcontainers.API] // Set by apiHandler.

}

func (a *agent) TailnetConn() *tailnet.Conn {

a.closeMutex.Lock()

defer a.closeMutex.Unlock()

return a.network

}

func (a *agent) init() {

// pass the "hard" context because we explicitly close the SSH server as part of graceful shutdown.

sshSrv, err := agentssh.NewServer(a.hardCtx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.execer, &agentssh.Config{

MaxTimeout: a.sshMaxTimeout,

MOTDFile: func() string { return a.manifest.Load().MOTDFile },

AnnouncementBanners: func() *[]codersdk.BannerConfig { return a.announcementBanners.Load() },

UpdateEnv: a.updateCommandEnv,

WorkingDirectory: func() string { return a.manifest.Load().Directory },

BlockFileTransfer: a.blockFileTransfer,

ReportConnection: func(id uuid.UUID, magicType agentssh.MagicSessionType, ip string) func(code int, reason string) {

var connectionType proto.Connection_Type

switch magicType {

case agentssh.MagicSessionTypeSSH:

connectionType = proto.Connection_SSH

case agentssh.MagicSessionTypeVSCode:

connectionType = proto.Connection_VSCODE

case agentssh.MagicSessionTypeJetBrains:

connectionType = proto.Connection_JETBRAINS

case agentssh.MagicSessionTypeUnknown:

connectionType = proto.Connection_TYPE_UNSPECIFIED

default:

a.logger.Error(a.hardCtx, "unhandled magic session type when reporting connection", slog.F("magic_type", magicType))

connectionType = proto.Connection_TYPE_UNSPECIFIED

}

return a.reportConnection(id, connectionType, ip)

},

ExperimentalDevContainersEnabled: a.experimentalDevcontainersEnabled,

})

if err != nil {

panic(err)

}

a.sshServer = sshSrv

a.scriptRunner = agentscripts.New(agentscripts.Options{

LogDir: a.logDir,

DataDirBase: a.scriptDataDir,

Logger: a.logger,

SSHServer: sshSrv,

Filesystem: a.filesystem,

GetScriptLogger: func(logSourceID uuid.UUID) agentscripts.ScriptLogger {

return a.logSender.GetScriptLogger(logSourceID)

},

})

// Register runner metrics. If the prom registry is nil, the metrics

// will not report anywhere.

a.scriptRunner.RegisterMetrics(a.prometheusRegistry)

a.reconnectingPTYServer = reconnectingpty.NewServer(

a.logger.Named("reconnecting-pty"),

a.sshServer,

func(id uuid.UUID, ip string) func(code int, reason string) {

return a.reportConnection(id, proto.Connection_RECONNECTING_PTY, ip)

},

a.metrics.connectionsTotal, a.metrics.reconnectingPTYErrors,

a.reconnectingPTYTimeout,

func(s *reconnectingpty.Server) {

s.ExperimentalDevcontainersEnabled = a.experimentalDevcontainersEnabled

},

)

go a.runLoop()

}

// runLoop attempts to start the agent in a retry loop.

// Coder may be offline temporarily, a connection issue

// may be happening, but regardless after the intermittent

// failure, you'll want the agent to reconnect.

func (a *agent) runLoop() {

// need to keep retrying up to the hardCtx so that we can send graceful shutdown-related

// messages.

ctx := a.hardCtx

for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {

a.logger.Info(ctx, "connecting to coderd")

err := a.run()

if err == nil {

continue

}

if ctx.Err() != nil {

// Context canceled errors may come from websocket pings, so we

// don't want to use `errors.Is(err, context.Canceled)` here.

a.logger.Warn(ctx, "runLoop exited with error", slog.Error(ctx.Err()))

return

}

if a.isClosed() {

a.logger.Warn(ctx, "runLoop exited because agent is closed")

return

}

if errors.Is(err, io.EOF) {

a.logger.Info(ctx, "disconnected from coderd")

continue

}

a.logger.Warn(ctx, "run exited with error", slog.Error(err))

}

}

func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentMetadataDescription, now time.Time) *codersdk.WorkspaceAgentMetadataResult {

var out bytes.Buffer

result := &codersdk.WorkspaceAgentMetadataResult{

// CollectedAt is set here for testing purposes and overrode by

// coderd to the time of server receipt to solve clock skew.

//

// In the future, the server may accept the timestamp from the agent

// if it can guarantee the clocks are synchronized.

CollectedAt: now,

}

cmdPty, err := a.sshServer.CreateCommand(ctx, md.Script, nil, nil)

if err != nil {

result.Error = fmt.Sprintf("create cmd: %+v", err)

return result

}

cmd := cmdPty.AsExec()

cmd.Stdout = &out

cmd.Stderr = &out

cmd.Stdin = io.LimitReader(nil, 0)

// We split up Start and Wait instead of calling Run so that we can return a more precise error.

err = cmd.Start()

if err != nil {

result.Error = fmt.Sprintf("start cmd: %+v", err)

return result

}

// This error isn't mutually exclusive with useful output.

err = cmd.Wait()

const bufLimit = 10 << 10

if out.Len() > bufLimit {

err = errors.Join(

err,

xerrors.Errorf("output truncated from %v to %v bytes", out.Len(), bufLimit),

)

out.Truncate(bufLimit)

}

// Important: if the command times out, we may see a misleading error like

// "exit status 1", so it's important to include the context error.

err = errors.Join(err, ctx.Err())

if err != nil {

result.Error = fmt.Sprintf("run cmd: %+v", err)

}

result.Value = out.String()

return result

}

type metadataResultAndKey struct {

result *codersdk.WorkspaceAgentMetadataResult

key string

}

type trySingleflight struct {

mu sync.Mutex

m map[string]struct{}

}

func (t *trySingleflight) Do(key string, fn func()) {

t.mu.Lock()

_, ok := t.m[key]

if ok {

t.mu.Unlock()

return

}

t.m[key] = struct{}{}

t.mu.Unlock()

defer func() {

t.mu.Lock()

delete(t.m, key)

t.mu.Unlock()

}()

fn()

}

func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

tickerDone := make(chan struct{})

collectDone := make(chan struct{})

ctx, cancel := context.WithCancel(ctx)

defer func() {

cancel()

<-collectDone

<-tickerDone

}()

var (

logger = a.logger.Named("metadata")

report = make(chan struct{}, 1)

collect = make(chan struct{}, 1)

metadataResults = make(chan metadataResultAndKey, 1)

)

// Set up collect and report as a single ticker with two channels,

// this is to allow collection and reporting to be triggered

// independently of each other.

go func() {

t := time.NewTicker(a.reportMetadataInterval)

defer func() {

t.Stop()

close(report)

close(collect)

close(tickerDone)

}()

wake := func(c chan<- struct{}) {

select {

case c <- struct{}{}:

default:

}

}

wake(collect) // Start immediately.

for {

select {

case <-ctx.Done():

return

case <-t.C:

wake(report)

wake(collect)

}

}

}()

go func() {

defer close(collectDone)

var (

// We use a custom singleflight that immediately returns if there is already

// a goroutine running for a given key. This is to prevent a build-up of

// goroutines waiting on Do when the script takes many multiples of

// baseInterval to run.

flight = trySingleflight{m: map[string]struct{}{}}

lastCollectedAtMu sync.RWMutex

lastCollectedAts = make(map[string]time.Time)

)

for {

select {

case <-ctx.Done():

return

case <-collect:

}

manifest := a.manifest.Load()

if manifest == nil {

continue

}

// If the manifest changes (e.g. on agent reconnect) we need to

// purge old cache values to prevent lastCollectedAt from growing

// boundlessly.

lastCollectedAtMu.Lock()

for key := range lastCollectedAts {

if slices.IndexFunc(manifest.Metadata, func(md codersdk.WorkspaceAgentMetadataDescription) bool {

return md.Key == key

}) < 0 {

logger.Debug(ctx, "deleting lastCollected key, missing from manifest",

slog.F("key", key),

)

delete(lastCollectedAts, key)

}

}

lastCollectedAtMu.Unlock()

// Spawn a goroutine for each metadata collection, and use a

// channel to synchronize the results and avoid both messy

// mutex logic and overloading the API.

for _, md := range manifest.Metadata {

md := md

// We send the result to the channel in the goroutine to avoid

// sending the same result multiple times. So, we don't care about

// the return values.

go flight.Do(md.Key, func() {

ctx := slog.With(ctx, slog.F("key", md.Key))

lastCollectedAtMu.RLock()

collectedAt, ok := lastCollectedAts[md.Key]

lastCollectedAtMu.RUnlock()

if ok {

// If the interval is zero, we assume the user just wants

// a single collection at startup, not a spinning loop.

if md.Interval == 0 {

return

}

intervalUnit := time.Second

// reportMetadataInterval is only less than a second in tests,

// so adjust the interval unit for them.

if a.reportMetadataInterval < time.Second {

intervalUnit = 100 * time.Millisecond

}

// The last collected value isn't quite stale yet, so we skip it.

if collectedAt.Add(time.Duration(md.Interval) * intervalUnit).After(time.Now()) {

return

}

}

timeout := md.Timeout

if timeout == 0 {

if md.Interval != 0 {

timeout = md.Interval

} else if interval := int64(a.reportMetadataInterval.Seconds()); interval != 0 {

// Fallback to the report interval

timeout = interval * 3

} else {

// If the interval is still 0 (possible if the interval

// is less than a second), default to 5. This was

// randomly picked.

timeout = 5

}

}

ctxTimeout := time.Duration(timeout) * time.Second

ctx, cancel := context.WithTimeout(ctx, ctxTimeout)

defer cancel()

now := time.Now()

select {

case <-ctx.Done():

logger.Warn(ctx, "metadata collection timed out", slog.F("timeout", ctxTimeout))

case metadataResults <- metadataResultAndKey{

key: md.Key,

result: a.collectMetadata(ctx, md, now),

}:

lastCollectedAtMu.Lock()

lastCollectedAts[md.Key] = now

lastCollectedAtMu.Unlock()

}

})

}

}

}()

// Gather metadata updates and report them once every interval. If a

// previous report is in flight, wait for it to complete before

// sending a new one. If the network conditions are bad, we won't

// benefit from canceling the previous send and starting a new one.

var (

updatedMetadata = make(map[string]*codersdk.WorkspaceAgentMetadataResult)

reportTimeout = 30 * time.Second

reportError = make(chan error, 1)

reportInFlight = false

)

for {

select {

case <-ctx.Done():

return ctx.Err()

case mr := <-metadataResults:

// This can overwrite unsent values, but that's fine because

// we're only interested about up-to-date values.

updatedMetadata[mr.key] = mr.result

continue

case err := <-reportError:

logMsg := "batch update metadata complete"

if err != nil {

a.logger.Debug(ctx, logMsg, slog.Error(err))

return xerrors.Errorf("failed to report metadata: %w", err)

}

a.logger.Debug(ctx, logMsg)

reportInFlight = false

case <-report:

if len(updatedMetadata) == 0 {

continue

}

if reportInFlight {

// If there's already a report in flight, don't send

// another one, wait for next tick instead.

a.logger.Debug(ctx, "skipped metadata report tick because report is in flight")

continue

}

metadata := make([]*proto.Metadata, 0, len(updatedMetadata))

for key, result := range updatedMetadata {

pr := agentsdk.ProtoFromMetadataResult(*result)

metadata = append(metadata, &proto.Metadata{

Key: key,

Result: pr,

})

delete(updatedMetadata, key)

}

reportInFlight = true

go func() {

a.logger.Debug(ctx, "batch updating metadata")

ctx, cancel := context.WithTimeout(ctx, reportTimeout)

defer cancel()

_, err := aAPI.BatchUpdateMetadata(ctx, &proto.BatchUpdateMetadataRequest{Metadata: metadata})

reportError <- err

}()

}

}

}

// reportLifecycle reports the current lifecycle state once. All state

// changes are reported in order.

func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

for {

select {

case <-a.lifecycleUpdate:

case <-ctx.Done():

return ctx.Err()

}

for {

a.lifecycleMu.RLock()

lastIndex := len(a.lifecycleStates) - 1

report := a.lifecycleStates[a.lifecycleLastReportedIndex]

if len(a.lifecycleStates) > a.lifecycleLastReportedIndex+1 {

report = a.lifecycleStates[a.lifecycleLastReportedIndex+1]

}

a.lifecycleMu.RUnlock()

if lastIndex == a.lifecycleLastReportedIndex {

break

}

l, err := agentsdk.ProtoFromLifecycle(report)

if err != nil {

a.logger.Critical(ctx, "failed to convert lifecycle state", slog.F("report", report))

// Skip this report; there is no point retrying. Maybe we can successfully convert the next one?

a.lifecycleLastReportedIndex++

continue

}

payload := &proto.UpdateLifecycleRequest{Lifecycle: l}

logger := a.logger.With(slog.F("payload", payload))

logger.Debug(ctx, "reporting lifecycle state")

_, err = aAPI.UpdateLifecycle(ctx, payload)

if err != nil {

return xerrors.Errorf("failed to update lifecycle: %w", err)

}

logger.Debug(ctx, "successfully reported lifecycle state")

a.lifecycleLastReportedIndex++

select {

case a.lifecycleReported <- report.State:

case <-a.lifecycleReported:

a.lifecycleReported <- report.State

}

if a.lifecycleLastReportedIndex < lastIndex {

// Keep reporting until we've sent all messages, we can't

// rely on the channel triggering us before the backlog is

// consumed.

continue

}

break

}

}

}

// setLifecycle sets the lifecycle state and notifies the lifecycle loop.

// The state is only updated if it's a valid state transition.

func (a *agent) setLifecycle(state codersdk.WorkspaceAgentLifecycle) {

report := agentsdk.PostLifecycleRequest{

State: state,

ChangedAt: dbtime.Now(),

}

a.lifecycleMu.Lock()

lastReport := a.lifecycleStates[len(a.lifecycleStates)-1]

if slices.Index(codersdk.WorkspaceAgentLifecycleOrder, lastReport.State) >= slices.Index(codersdk.WorkspaceAgentLifecycleOrder, report.State) {

a.logger.Warn(context.Background(), "attempted to set lifecycle state to a previous state", slog.F("last", lastReport), slog.F("current", report))

a.lifecycleMu.Unlock()

return

}

a.lifecycleStates = append(a.lifecycleStates, report)

a.logger.Debug(context.Background(), "set lifecycle state", slog.F("current", report), slog.F("last", lastReport))

a.lifecycleMu.Unlock()

select {

case a.lifecycleUpdate <- struct{}{}:

default:

}

}

// reportConnectionsLoop reports connections to the agent for auditing.

func (a *agent) reportConnectionsLoop(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

for {

select {

case <-a.reportConnectionsUpdate:

case <-ctx.Done():

return ctx.Err()

}

for {

a.reportConnectionsMu.Lock()

if len(a.reportConnections) == 0 {

a.reportConnectionsMu.Unlock()

break

}

payload := a.reportConnections[0]

// Release lock while we send the payload, this is safe

// since we only append to the slice.

a.reportConnectionsMu.Unlock()

logger := a.logger.With(slog.F("payload", payload))

logger.Debug(ctx, "reporting connection")

_, err := aAPI.ReportConnection(ctx, payload)

if err != nil {

return xerrors.Errorf("failed to report connection: %w", err)

}

logger.Debug(ctx, "successfully reported connection")

// Remove the payload we sent.

a.reportConnectionsMu.Lock()

a.reportConnections[0] = nil // Release the pointer from the underlying array.

a.reportConnections = a.reportConnections[1:]

a.reportConnectionsMu.Unlock()

}

}

}

const (

// reportConnectionBufferLimit limits the number of connection reports we

// buffer to avoid growing the buffer indefinitely. This should not happen

// unless the agent has lost connection to coderd for a long time or if

// the agent is being spammed with connections.

//

// If we assume ~150 byte per connection report, this would be around 300KB

// of memory which seems acceptable. We could reduce this if necessary by

// not using the proto struct directly.

reportConnectionBufferLimit = 2048

)

func (a *agent) reportConnection(id uuid.UUID, connectionType proto.Connection_Type, ip string) (disconnected func(code int, reason string)) {

// Remove the port from the IP because ports are not supported in coderd.

if host, _, err := net.SplitHostPort(ip); err != nil {

a.logger.Error(a.hardCtx, "split host and port for connection report failed", slog.F("ip", ip), slog.Error(err))

} else {

// Best effort.

ip = host

}

a.reportConnectionsMu.Lock()

defer a.reportConnectionsMu.Unlock()

if len(a.reportConnections) >= reportConnectionBufferLimit {

a.logger.Warn(a.hardCtx, "connection report buffer limit reached, dropping connect",

slog.F("limit", reportConnectionBufferLimit),

slog.F("connection_id", id),

slog.F("connection_type", connectionType),

slog.F("ip", ip),

)

} else {

a.reportConnections = append(a.reportConnections, &proto.ReportConnectionRequest{

Connection: &proto.Connection{

Id: id[:],

Action: proto.Connection_CONNECT,

Type: connectionType,

Timestamp: timestamppb.New(time.Now()),

Ip: ip,

StatusCode: 0,

Reason: nil,

},

})

select {

case a.reportConnectionsUpdate <- struct{}{}:

default:

}

}

return func(code int, reason string) {

a.reportConnectionsMu.Lock()

defer a.reportConnectionsMu.Unlock()

if len(a.reportConnections) >= reportConnectionBufferLimit {

a.logger.Warn(a.hardCtx, "connection report buffer limit reached, dropping disconnect",

slog.F("limit", reportConnectionBufferLimit),

slog.F("connection_id", id),

slog.F("connection_type", connectionType),

slog.F("ip", ip),

)

return

}

a.reportConnections = append(a.reportConnections, &proto.ReportConnectionRequest{

Connection: &proto.Connection{

Id: id[:],

Action: proto.Connection_DISCONNECT,

Type: connectionType,

Timestamp: timestamppb.New(time.Now()),

Ip: ip,

StatusCode: int32(code), //nolint:gosec

Reason: &reason,

},

})

select {

case a.reportConnectionsUpdate <- struct{}{}:

default:

}

}

}

// fetchServiceBannerLoop fetches the service banner on an interval. It will

// not be fetched immediately; the expectation is that it is primed elsewhere

// (and must be done before the session actually starts).

func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

ticker := time.NewTicker(a.announcementBannersRefreshInterval)

defer ticker.Stop()

for {

select {

case <-ctx.Done():

return ctx.Err()

case <-ticker.C:

bannersProto, err := aAPI.GetAnnouncementBanners(ctx, &proto.GetAnnouncementBannersRequest{})

if err != nil {

if ctx.Err() != nil {

return ctx.Err()

}

a.logger.Error(ctx, "failed to update notification banners", slog.Error(err))

return err

}

banners := make([]codersdk.BannerConfig, 0, len(bannersProto.AnnouncementBanners))

for _, bannerProto := range bannersProto.AnnouncementBanners {

banners = append(banners, agentsdk.BannerConfigFromProto(bannerProto))

}

a.announcementBanners.Store(&banners)

}

}

}

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)

if err != nil {

return xerrors.Errorf("exchange 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.ConnectRPC24(a.hardCtx)

if err != nil {

return err

}

defer func() {

cErr := aAPI.DRPCConn().Close()

if cErr != nil {

a.logger.Debug(a.hardCtx, "error closing drpc connection", slog.Error(cErr))

}

}()

// A lot of routines need the agent API / tailnet API connection. We run them in their own

// goroutines in parallel, but errors in any routine will cause them all to exit so we can

// redial the coder server and retry.

connMan := newAPIConnRoutineManager(a.gracefulCtx, a.hardCtx, a.logger, aAPI, tAPI)

connMan.startAgentAPI("init notification banners", gracefulShutdownBehaviorStop,

func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

bannersProto, err := aAPI.GetAnnouncementBanners(ctx, &proto.GetAnnouncementBannersRequest{})

if err != nil {

return xerrors.Errorf("fetch service banner: %w", err)

}

banners := make([]codersdk.BannerConfig, 0, len(bannersProto.AnnouncementBanners))

for _, bannerProto := range bannersProto.AnnouncementBanners {

banners = append(banners, agentsdk.BannerConfigFromProto(bannerProto))

}

a.announcementBanners.Store(&banners)

return nil

},

)

// sending logs gets gracefulShutdownBehaviorRemain because we want to send logs generated by

// shutdown scripts.

connMan.startAgentAPI("send logs", gracefulShutdownBehaviorRemain,

func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

err := a.logSender.SendLoop(ctx, aAPI)

if xerrors.Is(err, agentsdk.ErrLogLimitExceeded) {

// we don't want this error to tear down the API connection and propagate to the

// other routines that use the API. The LogSender has already dropped a warning

// log, so just return nil here.

return nil

}

return err

})

// part of graceful shut down is reporting the final lifecycle states, e.g "ShuttingDown" so the

// lifecycle reporting has to be via gracefulShutdownBehaviorRemain

connMan.startAgentAPI("report lifecycle", gracefulShutdownBehaviorRemain, a.reportLifecycle)

// metadata reporting can cease as soon as we start gracefully shutting down

connMan.startAgentAPI("report metadata", gracefulShutdownBehaviorStop, a.reportMetadata)

// resources monitor can cease as soon as we start gracefully shutting down.

connMan.startAgentAPI("resources monitor", gracefulShutdownBehaviorStop, func(ctx context.Context, aAPI proto.DRPCAgentClient24) error {

logger := a.logger.Named("resources_monitor")

clk := quartz.NewReal()

config, err := aAPI.GetResourcesMonitoringConfiguration(ctx, &proto.GetResourcesMonitoringConfigurationRequest{})

if err != nil {

return xerrors.Errorf("failed to get resources monitoring configuration: %w", err)

}

statfetcher, err := clistat.New()

if err != nil {

return xerrors.Errorf("failed to create resources fetcher: %w", err)

}

resourcesFetcher, err := resourcesmonitor.NewFetcher(statfetcher)

if err != nil {

return xerrors.Errorf("new resource fetcher: %w", err)

}

resourcesmonitor := resourcesmonitor.NewResourcesMonitor(logger, clk, config, resourcesFetcher, aAPI)

return resourcesmonitor.Start(ctx)

})

// Connection reports are part of auditing, we should keep sending them via

// gracefulShutdownBehaviorRemain.

connMan.startAgentAPI("report connections", gracefulShutdownBehaviorRemain, a.reportConnectionsLoop)

// channels to sync goroutines below

// handle manifest

// |

// manifestOK

// | |

// | +----------------------+

// V |