package cliui

import (

"context"

"fmt"

"io"

"strconv"

"strings"

"time"

"github.com/google/uuid"

"golang.org/x/xerrors"

"tailscale.com/tailcfg"

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

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

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

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

)

var errAgentShuttingDown = xerrors.New("agent is shutting down")

// fetchAgentResult is used to pass agent fetch results through channels.

type fetchAgentResult struct {

agent codersdk.WorkspaceAgent

err error

}

type AgentOptions struct {

FetchInterval time.Duration

Fetch func(ctx context.Context, agentID uuid.UUID) (codersdk.WorkspaceAgent, error)

FetchLogs func(ctx context.Context, agentID uuid.UUID, after int64, follow bool) (<-chan []codersdk.WorkspaceAgentLog, io.Closer, error)

Wait bool // If true, wait for the agent to be ready (startup script).

DocsURL string

}

// agentWaiter encapsulates the state machine for waiting on a workspace agent.

type agentWaiter struct {

opts AgentOptions

sw *stageWriter

logSources map[uuid.UUID]codersdk.WorkspaceAgentLogSource

fetchAgent func(context.Context) (codersdk.WorkspaceAgent, error)

}

// Agent displays a spinning indicator that waits for a workspace agent to connect.

func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentOptions) error {

ctx, cancel := context.WithCancel(ctx)

defer cancel()

if opts.FetchInterval == 0 {

opts.FetchInterval = 500 * time.Millisecond

}

if opts.FetchLogs == nil {

opts.FetchLogs = func(_ context.Context, _ uuid.UUID, _ int64, _ bool) (<-chan []codersdk.WorkspaceAgentLog, io.Closer, error) {

c := make(chan []codersdk.WorkspaceAgentLog)

close(c)

return c, closeFunc(func() error { return nil }), nil

}

}

fetchedAgent := make(chan fetchAgentResult, 1)

go func() {

t := time.NewTimer(0)

defer t.Stop()

startTime := time.Now()

baseInterval := opts.FetchInterval

for {

select {

case <-ctx.Done():

return

case <-t.C:

agent, err := opts.Fetch(ctx, agentID)

select {

case <-fetchedAgent:

default:

}

if err != nil {

fetchedAgent <- fetchAgentResult{err: xerrors.Errorf("fetch workspace agent: %w", err)}

return

}

fetchedAgent <- fetchAgentResult{agent: agent}

// Adjust the interval based on how long we've been waiting.

elapsed := time.Since(startTime)

currentInterval := GetProgressiveInterval(baseInterval, elapsed)

t.Reset(currentInterval)

}

}

}()

fetch := func(ctx context.Context) (codersdk.WorkspaceAgent, error) {

select {

case <-ctx.Done():

return codersdk.WorkspaceAgent{}, ctx.Err()

case f := <-fetchedAgent:

if f.err != nil {

return codersdk.WorkspaceAgent{}, f.err

}

return f.agent, nil

}

}

agent, err := fetch(ctx)

if err != nil {

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

}

logSources := map[uuid.UUID]codersdk.WorkspaceAgentLogSource{}

for _, source := range agent.LogSources {

logSources[source.ID] = source

}

w := &agentWaiter{

opts: opts,

sw: &stageWriter{w: writer},

logSources: logSources,

fetchAgent: fetch,

}

return w.wait(ctx, agent, fetchedAgent)

}

// wait runs the main state machine loop.

func (aw *agentWaiter) wait(ctx context.Context, agent codersdk.WorkspaceAgent, fetchedAgent chan fetchAgentResult) error {

var err error

// Track whether we've gone through a wait state, which determines if we

// should show startup logs when connected.

waitedForConnection := false

for {

// It doesn't matter if we're connected or not, if the agent is

// shutting down, we don't know if it's coming back.

if agent.LifecycleState.ShuttingDown() {

return errAgentShuttingDown

}

switch agent.Status {

case codersdk.WorkspaceAgentConnecting, codersdk.WorkspaceAgentTimeout:

agent, err = aw.waitForConnection(ctx, agent)

if err != nil {

return err

}

// Since we were waiting for the agent to connect, also show

// startup logs if applicable.

waitedForConnection = true

case codersdk.WorkspaceAgentConnected:

return aw.handleConnected(ctx, agent, waitedForConnection, fetchedAgent)

case codersdk.WorkspaceAgentDisconnected:

agent, waitedForConnection, err = aw.waitForReconnection(ctx, agent)

if err != nil {

return err

}

}

}

}

// waitForConnection handles the Connecting/Timeout states.

// Returns when agent transitions to Connected or Disconnected.

func (aw *agentWaiter) waitForConnection(ctx context.Context, agent codersdk.WorkspaceAgent) (codersdk.WorkspaceAgent, error) {

stage := "Waiting for the workspace agent to connect"

aw.sw.Start(stage)

agent, err := aw.pollWhile(ctx, agent, func(agent codersdk.WorkspaceAgent) bool {

return agent.Status == codersdk.WorkspaceAgentConnecting

})

if err != nil {

return agent, err

}

if agent.Status == codersdk.WorkspaceAgentTimeout {

now := time.Now()

aw.sw.Log(now, codersdk.LogLevelInfo, "The workspace agent is having trouble connecting, wait for it to connect or restart your workspace.")

aw.sw.Log(now, codersdk.LogLevelInfo, troubleshootingMessage(agent, fmt.Sprintf("%s/admin/templates/troubleshooting#agent-connection-issues", aw.opts.DocsURL)))

agent, err = aw.pollWhile(ctx, agent, func(agent codersdk.WorkspaceAgent) bool {

return agent.Status == codersdk.WorkspaceAgentTimeout

})

if err != nil {

return agent, err

}

}

aw.sw.Complete(stage, agent.FirstConnectedAt.Sub(agent.CreatedAt))

return agent, nil

}

// handleConnected handles the Connected state and startup script logic.

// This is a terminal state, returns nil on success or error on failure.

//

//nolint:revive // Control flag is acceptable for internal method.

func (aw *agentWaiter) handleConnected(ctx context.Context, agent codersdk.WorkspaceAgent, showStartupLogs bool, fetchedAgent chan fetchAgentResult) error {

if !showStartupLogs && agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady {

// The workspace is ready, there's nothing to do but connect.

return nil

}

// Determine if we should follow/stream logs (blocking mode).

follow := aw.opts.Wait && agent.LifecycleState.Starting()

stage := "Running workspace agent startup scripts"

if !follow {

stage += " (non-blocking)"

}

aw.sw.Start(stage)

if follow {

aw.sw.Log(time.Time{}, codersdk.LogLevelInfo, "==> ℹ︎ To connect immediately, reconnect with --wait=no or CODER_SSH_WAIT=no, see --help for more information.")

}

// In non-blocking mode (Wait=false), we don't stream logs. This prevents

// dumping a wall of logs on users who explicitly pass --wait=no. The stage

// indicator is still shown, just not the log content. See issue #13580.

if aw.opts.Wait {

var err error

agent, err = aw.streamLogs(ctx, agent, follow, fetchedAgent)

if err != nil {

return err

}

// If we were following, wait until startup completes.

if follow {

agent, err = aw.pollWhile(ctx, agent, func(agent codersdk.WorkspaceAgent) bool {

return agent.LifecycleState.Starting()

})

if err != nil {

return err

}

}

}

// Handle final lifecycle state.

switch agent.LifecycleState {

case codersdk.WorkspaceAgentLifecycleReady:

aw.sw.Complete(stage, safeDuration(aw.sw, agent.ReadyAt, agent.StartedAt))

case codersdk.WorkspaceAgentLifecycleStartTimeout:

// Backwards compatibility: Avoid printing warning if

// coderd is old and doesn't set ReadyAt for timeouts.

if agent.ReadyAt == nil {

aw.sw.Fail(stage, 0)

} else {

aw.sw.Fail(stage, safeDuration(aw.sw, agent.ReadyAt, agent.StartedAt))

}

aw.sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: A startup script timed out and your workspace may be incomplete.")

case codersdk.WorkspaceAgentLifecycleStartError:

aw.sw.Fail(stage, safeDuration(aw.sw, agent.ReadyAt, agent.StartedAt))

aw.sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: A startup script exited with an error and your workspace may be incomplete.")

aw.sw.Log(time.Time{}, codersdk.LogLevelWarn, troubleshootingMessage(agent, fmt.Sprintf("%s/admin/templates/troubleshooting#startup-script-exited-with-an-error", aw.opts.DocsURL)))

default:

switch {

case agent.LifecycleState.Starting():

aw.sw.Log(time.Time{}, codersdk.LogLevelWarn, "Notice: The startup scripts are still running and your workspace may be incomplete.")

aw.sw.Log(time.Time{}, codersdk.LogLevelWarn, troubleshootingMessage(agent, fmt.Sprintf("%s/admin/templates/troubleshooting#your-workspace-may-be-incomplete", aw.opts.DocsURL)))

// Note: We don't complete or fail the stage here, it's

// intentionally left open to indicate this stage didn't

// complete.

case agent.LifecycleState.ShuttingDown():

// We no longer know if the startup script failed or not,

// but we need to tell the user something.

aw.sw.Complete(stage, safeDuration(aw.sw, agent.ReadyAt, agent.StartedAt))

return errAgentShuttingDown

}

}

return nil

}

// streamLogs handles streaming or fetching startup logs.

//

//nolint:revive // Control flag is acceptable for internal method.

func (aw *agentWaiter) streamLogs(ctx context.Context, agent codersdk.WorkspaceAgent, follow bool, fetchedAgent chan fetchAgentResult) (codersdk.WorkspaceAgent, error) {

logStream, logsCloser, err := aw.opts.FetchLogs(ctx, agent.ID, 0, follow)

if err != nil {

return agent, xerrors.Errorf("fetch workspace agent startup logs: %w", err)

}

defer logsCloser.Close()

var lastLog codersdk.WorkspaceAgentLog

// If not following, we don't need to watch for agent state changes.

var fetchedAgentWhileFollowing chan fetchAgentResult

if follow {

fetchedAgentWhileFollowing = fetchedAgent

}

for {

select {

case <-ctx.Done():

return agent, ctx.Err()

case f := <-fetchedAgentWhileFollowing:

if f.err != nil {

return agent, xerrors.Errorf("fetch: %w", f.err)

}

agent = f.agent

// If the agent is no longer starting, stop following

// logs because FetchLogs will keep streaming forever.

// We do one last non-follow request to ensure we have

// fetched all logs.

if !agent.LifecycleState.Starting() {

_ = logsCloser.Close()

fetchedAgentWhileFollowing = nil

logStream, logsCloser, err = aw.opts.FetchLogs(ctx, agent.ID, lastLog.ID, false)

if err != nil {

return agent, xerrors.Errorf("fetch workspace agent startup logs: %w", err)

}

// Logs are already primed, so we can call close.

_ = logsCloser.Close()

}

case logs, ok := <-logStream:

if !ok {

return agent, nil

}

for _, log := range logs {

source, hasSource := aw.logSources[log.SourceID]

output := log.Output

if hasSource && source.DisplayName != "" {

output = source.DisplayName + ": " + output

}

aw.sw.Log(log.CreatedAt, log.Level, output)

lastLog = log

}

}

}

}

// waitForReconnection handles the Disconnected state.

// Returns when agent reconnects along with whether to show startup logs.

func (aw *agentWaiter) waitForReconnection(ctx context.Context, agent codersdk.WorkspaceAgent) (codersdk.WorkspaceAgent, bool, error) {

// If the agent was still starting during disconnect, we'll

// show startup logs.

showStartupLogs := agent.LifecycleState.Starting()

stage := "The workspace agent lost connection"

aw.sw.Start(stage)

aw.sw.Log(time.Now(), codersdk.LogLevelWarn, "Wait for it to reconnect or restart your workspace.")

aw.sw.Log(time.Now(), codersdk.LogLevelWarn, troubleshootingMessage(agent, fmt.Sprintf("%s/admin/templates/troubleshooting#agent-connection-issues", aw.opts.DocsURL)))

disconnectedAt := agent.DisconnectedAt

agent, err := aw.pollWhile(ctx, agent, func(agent codersdk.WorkspaceAgent) bool {

return agent.Status == codersdk.WorkspaceAgentDisconnected

})

if err != nil {

return agent, showStartupLogs, err

}

aw.sw.Complete(stage, safeDuration(aw.sw, agent.LastConnectedAt, disconnectedAt))

return agent, showStartupLogs, nil

}

// pollWhile polls the agent while the condition is true. It fetches the agent

// on each iteration and returns the updated agent when the condition is false,

// the context is canceled, or an error occurs.

func (aw *agentWaiter) pollWhile(ctx context.Context, agent codersdk.WorkspaceAgent, cond func(agent codersdk.WorkspaceAgent) bool) (codersdk.WorkspaceAgent, error) {

var err error

for cond(agent) {

agent, err = aw.fetchAgent(ctx)

if err != nil {

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

}

}

if err = ctx.Err(); err != nil {

return agent, err

}

return agent, nil

}

func troubleshootingMessage(agent codersdk.WorkspaceAgent, url string) string {

m := "For more information and troubleshooting, see " + url

if agent.TroubleshootingURL != "" {

m += " and " + agent.TroubleshootingURL

}

return m

}

// safeDuration returns a-b. If a or b is nil, it returns 0.

// This is because we often dereference a time pointer, which can

// cause a panic. These dereferences are used to calculate durations,

// which are not critical, and therefor should not break things

// when it fails.

// A panic has been observed in a test.

func safeDuration(sw *stageWriter, a, b *time.Time) time.Duration {

if a == nil || b == nil {

if sw != nil {

// Ideally the message includes which fields are <nil>, but you can

// use the surrounding log lines to figure that out. And passing more

// params makes this unwieldy.

sw.Log(time.Now(), codersdk.LogLevelWarn, "Warning: Failed to calculate duration from a time being <nil>.")

}

return 0

}

return a.Sub(*b)

}

// GetProgressiveInterval returns an interval that increases over time.

// The interval starts at baseInterval and increases to

// a maximum of baseInterval * 16 over time.

func GetProgressiveInterval(baseInterval time.Duration, elapsed time.Duration) time.Duration {

switch {

case elapsed < 60*time.Second:

return baseInterval // 500ms for first 60 seconds

case elapsed < 2*time.Minute:

return baseInterval * 2 // 1s for next 1 minute

case elapsed < 5*time.Minute:

return baseInterval * 4 // 2s for next 3 minutes

case elapsed < 10*time.Minute:

return baseInterval * 8 // 4s for next 5 minutes

default:

return baseInterval * 16 // 8s after 10 minutes

}

}

type closeFunc func() error

func (c closeFunc) Close() error {

return c()

}

func PeerDiagnostics(w io.Writer, d tailnet.PeerDiagnostics) {

if d.PreferredDERP > 0 {

rn, ok := d.DERPRegionNames[d.PreferredDERP]

if !ok {

rn = "unknown"

}

_, _ = fmt.Fprintf(w, "✔ preferred DERP region: %d (%s)\n", d.PreferredDERP, rn)

} else {

_, _ = fmt.Fprint(w, "✘ not connected to DERP\n")

}

if d.SentNode {

_, _ = fmt.Fprint(w, "✔ sent local data to Coder networking coordinator\n")

} else {

_, _ = fmt.Fprint(w, "✘ have not sent local data to Coder networking coordinator\n")

}

if d.ReceivedNode != nil {

dp := d.ReceivedNode.DERP

dn := ""

// should be 127.3.3.40:N where N is the DERP region

ap := strings.Split(dp, ":")

if len(ap) == 2 {

dp = ap[1]

di, err := strconv.Atoi(dp)

if err == nil {

var ok bool

dn, ok = d.DERPRegionNames[di]

if ok {

dn = fmt.Sprintf("(%s)", dn)

} else {

dn = "(unknown)"

}

}

}

_, _ = fmt.Fprintf(w,

"✔ received remote agent data from Coder networking coordinator\n preferred DERP region: %s %s\n endpoints: %s\n",

dp, dn, strings.Join(d.ReceivedNode.Endpoints, ", "))

} else {

_, _ = fmt.Fprint(w, "✘ have not received remote agent data from Coder networking coordinator\n")

}

if !d.LastWireguardHandshake.IsZero() {

ago := time.Since(d.LastWireguardHandshake)

symbol := "✔"

// wireguard is supposed to refresh handshake on 5 minute intervals

if ago > 5*time.Minute {

symbol = "⚠"

}

_, _ = fmt.Fprintf(w, "%s Wireguard handshake %s ago\n", symbol, ago.Round(time.Second))

} else {

_, _ = fmt.Fprint(w, "✘ Wireguard is not connected\n")

}

}

type ConnDiags struct {

ConnInfo workspacesdk.AgentConnectionInfo

PingP2P bool

DisableDirect bool

LocalNetInfo *tailcfg.NetInfo

LocalInterfaces *healthsdk.InterfacesReport

AgentNetcheck *healthsdk.AgentNetcheckReport

ClientIPIsAWS bool

AgentIPIsAWS bool

Verbose bool

TroubleshootingURL string

}

func (d ConnDiags) Write(w io.Writer) {

_, _ = fmt.Fprintln(w, "")

general, client, agent := d.splitDiagnostics()

for _, msg := range general {

_, _ = fmt.Fprintln(w, msg)

}

if len(general) > 0 {

_, _ = fmt.Fprintln(w, "")

}

if len(client) > 0 {

_, _ = fmt.Fprint(w, "Possible client-side issues with direct connection:\n\n")

for _, msg := range client {

_, _ = fmt.Fprintf(w, " - %s\n\n", msg)

}

}

if len(agent) > 0 {

_, _ = fmt.Fprint(w, "Possible agent-side issues with direct connections:\n\n")

for _, msg := range agent {

_, _ = fmt.Fprintf(w, " - %s\n\n", msg)

}

}

}

func (d ConnDiags) splitDiagnostics() (general, client, agent []string) {

if d.AgentNetcheck != nil {

for _, msg := range d.AgentNetcheck.Interfaces.Warnings {

agent = append(agent, msg.Message)

}

if len(d.AgentNetcheck.Interfaces.Warnings) > 0 {

agent[len(agent)-1] += fmt.Sprintf("\n%s#low-mtu", d.TroubleshootingURL)

}

}

if d.LocalInterfaces != nil {

for _, msg := range d.LocalInterfaces.Warnings {

client = append(client, msg.Message)

}

if len(d.LocalInterfaces.Warnings) > 0 {

client[len(client)-1] += fmt.Sprintf("\n%s#low-mtu", d.TroubleshootingURL)

}

}

if d.PingP2P && !d.Verbose {

return general, client, agent

}

if d.DisableDirect {

general = append(general, "❗ Direct connections are disabled locally, by `--disable-direct-connections` or `CODER_DISABLE_DIRECT_CONNECTIONS`.\n"+

" They may still be established over a private network.")

if !d.Verbose {

return general, client, agent

}

}

if d.ConnInfo.DisableDirectConnections {

general = append(general,

fmt.Sprintf("❗ Your Coder administrator has blocked direct connections\n %s#disabled-deployment-wide", d.TroubleshootingURL))

if !d.Verbose {

return general, client, agent

}

}

if !d.ConnInfo.DERPMap.HasSTUN() {

general = append(general,

fmt.Sprintf("❗ The DERP map is not configured to use STUN\n %s#no-stun-servers", d.TroubleshootingURL))

} else if d.LocalNetInfo != nil && !d.LocalNetInfo.UDP {

client = append(client,

fmt.Sprintf("Client could not connect to STUN over UDP\n %s#udp-blocked", d.TroubleshootingURL))

}

if d.LocalNetInfo != nil && d.LocalNetInfo.MappingVariesByDestIP.EqualBool(true) {

client = append(client,

fmt.Sprintf("Client is potentially behind a hard NAT, as multiple endpoints were retrieved from different STUN servers\n %s#endpoint-dependent-nat-hard-nat", d.TroubleshootingURL))

}

if d.AgentNetcheck != nil && d.AgentNetcheck.NetInfo != nil {

if d.AgentNetcheck.NetInfo.MappingVariesByDestIP.EqualBool(true) {

agent = append(agent,

fmt.Sprintf("Agent is potentially behind a hard NAT, as multiple endpoints were retrieved from different STUN servers\n %s#endpoint-dependent-nat-hard-nat", d.TroubleshootingURL))

}

if !d.AgentNetcheck.NetInfo.UDP {

agent = append(agent,

fmt.Sprintf("Agent could not connect to STUN over UDP\n %s#udp-blocked", d.TroubleshootingURL))

}

}

if d.ClientIPIsAWS {

client = append(client,

fmt.Sprintf("Client IP address is within an AWS range (AWS uses hard NAT)\n %s#endpoint-dependent-nat-hard-nat", d.TroubleshootingURL))

}

if d.AgentIPIsAWS {

agent = append(agent,

fmt.Sprintf("Agent IP address is within an AWS range (AWS uses hard NAT)\n %s#endpoint-dependent-nat-hard-nat", d.TroubleshootingURL))

}

return general, client, agent

}