package tailnet

import (

"context"

"errors"

"fmt"

"io"

"maps"

"math"

"net"

"net/netip"

"slices"

"strings"

"sync"

"time"

"github.com/google/uuid"

"golang.org/x/xerrors"

"storj.io/drpc"

"storj.io/drpc/drpcerr"

"tailscale.com/net/tsaddr"

"tailscale.com/tailcfg"

"tailscale.com/util/dnsname"

"cdr.dev/slog/v3"

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

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

"github.com/coder/quartz"

"github.com/coder/retry"

)

// A Controller connects to the tailnet control plane, and then uses the control protocols to

// program a tailnet.Conn in production (in test it could be an interface simulating the Conn). It

// delegates this task to sub-controllers responsible for the main areas of the tailnet control

// protocol: coordination, DERP map updates, resume tokens, telemetry, and workspace updates.

type Controller struct {

Dialer ControlProtocolDialer

CoordCtrl CoordinationController

DERPCtrl DERPController

ResumeTokenCtrl ResumeTokenController

TelemetryCtrl TelemetryController

WorkspaceUpdatesCtrl WorkspaceUpdatesController

ctx context.Context

gracefulCtx context.Context

cancelGracefulCtx context.CancelFunc

logger slog.Logger

closedCh chan struct{}

// Testing only

clock quartz.Clock

gracefulTimeout time.Duration

}

type CloserWaiter interface {

Close(context.Context) error

Wait() <-chan error

}

// CoordinatorClient is an abstraction of the Coordinator's control protocol interface from the

// perspective of a protocol client (i.e. the Coder Agent is also a client of this interface).

type CoordinatorClient interface {

Close() error

Send(*proto.CoordinateRequest) error

Recv() (*proto.CoordinateResponse, error)

}

// A CoordinationController accepts connections to the control plane, and handles the Coordination

// protocol on behalf of some Coordinatee (tailnet.Conn in production). This is the "glue" code

// between them.

type CoordinationController interface {

New(CoordinatorClient) CloserWaiter

}

// DERPClient is an abstraction of the stream of DERPMap updates from the control plane.

type DERPClient interface {

Close() error

Recv() (*tailcfg.DERPMap, error)

}

// A DERPController accepts connections to the control plane, and handles the DERPMap updates

// delivered over them by programming the data plane (tailnet.Conn or some test interface).

type DERPController interface {

New(DERPClient) CloserWaiter

}

type ResumeTokenClient interface {

RefreshResumeToken(ctx context.Context, in *proto.RefreshResumeTokenRequest) (*proto.RefreshResumeTokenResponse, error)

}

type ResumeTokenController interface {

New(ResumeTokenClient) CloserWaiter

Token() (string, bool)

}

type TelemetryClient interface {

PostTelemetry(ctx context.Context, in *proto.TelemetryRequest) (*proto.TelemetryResponse, error)

}

type TelemetryController interface {

New(TelemetryClient)

}

type WorkspaceUpdatesClient interface {

Close() error

Recv() (*proto.WorkspaceUpdate, error)

}

type WorkspaceUpdatesController interface {

New(WorkspaceUpdatesClient) CloserWaiter

}

// DNSHostsSetter is something that you can set a mapping of DNS names to IPs on. It's the subset

// of the tailnet.Conn that we use to configure DNS records.

type DNSHostsSetter interface {

SetDNSHosts(hosts map[dnsname.FQDN][]netip.Addr) error

}

// UpdatesHandler is anything that expects a stream of workspace update diffs.

type UpdatesHandler interface {

Update(WorkspaceUpdate) error

}

// ControlProtocolClients represents an abstract interface to the tailnet control plane via a set

// of protocol clients. The Closer should close all the clients (e.g. by closing the underlying

// connection).

type ControlProtocolClients struct {

Closer io.Closer

Coordinator CoordinatorClient

DERP DERPClient

ResumeToken ResumeTokenClient

Telemetry TelemetryClient

WorkspaceUpdates WorkspaceUpdatesClient

}

type ControlProtocolDialer interface {

// Dial connects to the tailnet control plane and returns clients for the different control

// sub-protocols (coordination, DERP maps, resume tokens, and telemetry). If the

// ResumeTokenController is not nil, the dialer should query for a resume token and use it to

// dial, if available.

Dial(ctx context.Context, r ResumeTokenController) (ControlProtocolClients, error)

}

// BasicCoordinationController handles the basic coordination operations common to all types of

// tailnet consumers:

//

// 1. sending local node updates to the Coordinator

// 2. receiving peer node updates and programming them into the Coordinatee (e.g. tailnet.Conn)

// 3. (optionally) sending ReadyToHandshake acknowledgements for peer updates.

//

// It is designed to be used on its own, or composed into more advanced CoordinationControllers.

type BasicCoordinationController struct {

Logger slog.Logger

Coordinatee Coordinatee

SendAcks bool

// Initiator labels which side of the coordination this controller

// represents (e.g. "agent", "client", "server"). It is attached as

// a slog field to disconnect-related logs emitted by this

// coordination so operators can attribute connection-close events

// without inspecting call sites. Leave unset on synthetic

// controllers (tests, in-memory fakes) and the field will be

// omitted from log output.

Initiator codersdk.DisconnectInitiator

// Direction labels the connection layer (server_to_agent,

// agent_to_client, client_to_server) for disconnect logs.

Direction codersdk.ConnectionDirection

}

// New satisfies the method on the CoordinationController interface

func (c *BasicCoordinationController) New(client CoordinatorClient) CloserWaiter {

return c.NewCoordination(client)

}

// NewCoordination creates a BasicCoordination

func (c *BasicCoordinationController) NewCoordination(client CoordinatorClient) *BasicCoordination {

b := &BasicCoordination{

logger: c.Logger,

errChan: make(chan error, 1),

coordinatee: c.Coordinatee,

client: client,

respLoopDone: make(chan struct{}),

sendAcks: c.SendAcks,

initiator: c.Initiator,

direction: c.Direction,

}

c.Coordinatee.SetNodeCallback(func(node *Node) {

pn, err := NodeToProto(node)

if err != nil {

b.logger.Critical(context.Background(), "failed to convert node", slog.Error(err))

b.SendErr(err)

return

}

b.Lock()

defer b.Unlock()

if b.closed {

b.logger.Debug(context.Background(), "ignored node update because coordination is closed")

return

}

err = b.client.Send(&proto.CoordinateRequest{UpdateSelf: &proto.CoordinateRequest_UpdateSelf{Node: pn}})

if err != nil {

b.SendErr(xerrors.Errorf("write: %w", err))

}

})

go b.respLoop()

return b

}

// BasicCoordination handles:

//

// 1. Sending local node updates to the control plane

// 2. Reading remote updates from the control plane and programming them into the Coordinatee.

//

// It does *not* handle adding any Tunnels, but these can be handled by composing

// BasicCoordinationController with a more advanced controller.

type BasicCoordination struct {

sync.Mutex

closed bool

errChan chan error

coordinatee Coordinatee

logger slog.Logger

client CoordinatorClient

respLoopDone chan struct{}

sendAcks bool

initiator codersdk.DisconnectInitiator

direction codersdk.ConnectionDirection

}

// CloseClient forcibly closes the underlying coordinator client connection

// without sending a graceful Disconnect message. Use this when you need to

// tear down the connection immediately, for example after a send error.

func (c *BasicCoordination) CloseClient() error {

return c.client.Close()

}

// SendRequest sends a coordinate request on the client connection, holding

// the coordination lock to prevent concurrent writes on the dRPC stream.

func (c *BasicCoordination) SendRequest(req *proto.CoordinateRequest) error {

c.Lock()

defer c.Unlock()

if c.closed {

return xerrors.New("coordination is closed")

}

return c.client.Send(req)

}

// Close the coordination gracefully. If the context expires before the remote API server has hung

// up on us, we forcibly close the Client connection.

func (c *BasicCoordination) Close(ctx context.Context) (retErr error) {

c.Lock()

if c.closed {

c.Unlock()

return nil

}

c.closed = true

err := c.client.Send(&proto.CoordinateRequest{Disconnect: &proto.CoordinateRequest_Disconnect{}})

c.Unlock()

if err != nil && !xerrors.Is(err, io.EOF) {

reason := codersdk.DisconnectReasonNetworkError

// Log but don't return early; we must still clean up below.

c.logger.Warn(context.Background(), "failed to send disconnect",

c.direction.SlogField(),

reason.SlogExpectedField(),

reason.SlogField(),

c.initiator.SlogField(),

slog.Error(err),

)

retErr = xerrors.Errorf("send disconnect: %w", err)

} else {

reason := codersdk.DisconnectReasonGraceful

c.logger.Debug(context.Background(), "sent disconnect",

c.direction.SlogField(),

reason.SlogExpectedField(),

reason.SlogField(),

c.initiator.SlogField(),

)

}

// We shouldn't just close the protocol right away, because the way dRPC streams work is

// that if you close them, that could take effect immediately, even before the Disconnect

// message is processed. Coordinators are supposed to hang up on us once they get a

// Disconnect message, so we should wait around for that until the context expires.

select {

case <-c.respLoopDone:

reason := codersdk.DisconnectReasonGraceful

c.logger.Debug(ctx, "responses closed after disconnect",

c.direction.SlogField(),

reason.SlogExpectedField(),

reason.SlogField(),

c.initiator.SlogField(),

)

return retErr

case <-ctx.Done():

reason := codersdk.DisconnectReasonNetworkError

c.logger.Warn(ctx, "context expired while waiting for coordinate responses to close",

c.direction.SlogField(),

reason.SlogExpectedField(),

reason.SlogField(),

c.initiator.SlogField(),

codersdk.SlogDisconnectDetail("context expired before coordinator hung up"),

)

}

// forcefully close the stream

protoErr := c.client.Close()

<-c.respLoopDone

if retErr == nil {

retErr = protoErr

}

return retErr

}

// Wait for the Coordination to complete

func (c *BasicCoordination) Wait() <-chan error {

return c.errChan

}

// SendErr is not part of the CloserWaiter interface, and is intended to be called internally, or

// by Controllers that use BasicCoordinationController in composition. It triggers Wait() to

// report the error if an error has not already been reported.

func (c *BasicCoordination) SendErr(err error) {

select {

case c.errChan <- err:

default:

}

}

func (c *BasicCoordination) respLoop() {

defer func() {

cErr := c.client.Close()

if cErr != nil {

c.logger.Debug(context.Background(),

"failed to close coordinate client after respLoop exit", slog.Error(cErr))

}

c.coordinatee.SetAllPeersLost()

close(c.respLoopDone)

}()

for {

resp, err := c.client.Recv()

if err != nil {

c.logger.Debug(context.Background(),

"failed to read from protocol", slog.Error(err))

c.SendErr(xerrors.Errorf("read: %w", err))

return

}

if resp.Error != "" {

// ReadyForHandshake error can occur during race conditions, where we send a ReadyForHandshake message,

// but the source has already disconnected from the tunnel by the time we do. So, just log at warning.

if strings.HasPrefix(resp.Error, ReadyForHandshakeError) {

c.logger.Warn(context.Background(), "coordination warning", slog.F("msg", resp.Error))

} else {

c.logger.Error(context.Background(),

"coordination protocol error", slog.F("error", resp.Error))

}

}

err = c.coordinatee.UpdatePeers(resp.GetPeerUpdates())

if err != nil {

c.logger.Debug(context.Background(), "failed to update peers", slog.Error(err))

c.SendErr(xerrors.Errorf("update peers: %w", err))

return

}

// Only send ReadyForHandshake acks from peers without a target.

if c.sendAcks {

// Send an ack back for all received peers. This could

// potentially be smarter to only send an ACK once per client,

// but there's nothing currently stopping clients from reusing

// IDs.

rfh := []*proto.CoordinateRequest_ReadyForHandshake{}

for _, peer := range resp.GetPeerUpdates() {

if peer.Kind != proto.CoordinateResponse_PeerUpdate_NODE {

continue

}

rfh = append(rfh, &proto.CoordinateRequest_ReadyForHandshake{Id: peer.Id})

}

if len(rfh) > 0 {

err := c.SendRequest(&proto.CoordinateRequest{

ReadyForHandshake: rfh,

})

if err != nil {

c.logger.Debug(context.Background(),

"failed to send ready for handshake", slog.Error(err))

c.SendErr(xerrors.Errorf("send: %w", err))

return

}

}

}

}

}

type TunnelSrcCoordController struct {

*BasicCoordinationController

mu sync.Mutex

dests map[uuid.UUID]struct{}

coordination *BasicCoordination

}

// NewTunnelSrcCoordController creates a CoordinationController for peers that are exclusively

// tunnel sources (that is, they create tunnel --- Coder clients not workspaces).

func NewTunnelSrcCoordController(

logger slog.Logger, coordinatee Coordinatee,

) *TunnelSrcCoordController {

return &TunnelSrcCoordController{

BasicCoordinationController: &BasicCoordinationController{

Logger: logger,

Coordinatee: coordinatee,

SendAcks: false,

Initiator: codersdk.DisconnectInitiatorClient,

Direction: codersdk.ConnectionDirectionClientToServer,

},

dests: make(map[uuid.UUID]struct{}),

}

}

func (c *TunnelSrcCoordController) New(client CoordinatorClient) CloserWaiter {

c.mu.Lock()

defer c.mu.Unlock()

b := c.NewCoordination(client)

c.coordination = b

// resync destinations on reconnect

for dest := range c.dests {

err := b.SendRequest(&proto.CoordinateRequest{

AddTunnel: &proto.CoordinateRequest_Tunnel{Id: UUIDToByteSlice(dest)},

})

if err != nil {

b.SendErr(err)

c.coordination = nil

cErr := client.Close()

if cErr != nil {

c.Logger.Debug(

context.Background(),

"failed to close coordinator client after add tunnel failure",

slog.Error(cErr),

)

}

break

}

}

return b

}

func (c *TunnelSrcCoordController) AddDestination(dest uuid.UUID) {

c.mu.Lock()

defer c.mu.Unlock()

c.Coordinatee.SetTunnelDestination(dest) // this prepares us for an ack

c.dests[dest] = struct{}{}

if c.coordination == nil {

return

}

err := c.coordination.SendRequest(

&proto.CoordinateRequest{

AddTunnel: &proto.CoordinateRequest_Tunnel{Id: UUIDToByteSlice(dest)},

})

if err != nil {

c.coordination.SendErr(err)

cErr := c.coordination.client.Close() // close the client so we don't gracefully disconnect

if cErr != nil {

c.Logger.Debug(context.Background(),

"failed to close coordinator client after add tunnel failure",

slog.Error(cErr))

}

c.coordination = nil

}

}

func (c *TunnelSrcCoordController) RemoveDestination(dest uuid.UUID) {

c.mu.Lock()

defer c.mu.Unlock()

delete(c.dests, dest)

if c.coordination == nil {

return

}

err := c.coordination.SendRequest(

&proto.CoordinateRequest{

RemoveTunnel: &proto.CoordinateRequest_Tunnel{Id: UUIDToByteSlice(dest)},

})

if err != nil {

c.coordination.SendErr(err)

cErr := c.coordination.client.Close() // close the client so we don't gracefully disconnect

if cErr != nil {

c.Logger.Debug(context.Background(),

"failed to close coordinator client after remove tunnel failure",

slog.Error(cErr))

}

c.coordination = nil

}

}

func (c *TunnelSrcCoordController) SyncDestinations(destinations []uuid.UUID) {

c.mu.Lock()

defer c.mu.Unlock()

toAdd := make(map[uuid.UUID]struct{})

toRemove := maps.Clone(c.dests)

all := make(map[uuid.UUID]struct{})

for _, dest := range destinations {

all[dest] = struct{}{}

delete(toRemove, dest)

if _, ok := c.dests[dest]; !ok {

toAdd[dest] = struct{}{}

}

}

c.dests = all

if c.coordination == nil {

return

}

var err error

defer func() {

if err != nil {

c.coordination.SendErr(err)

cErr := c.coordination.client.Close() // don't gracefully disconnect

if cErr != nil {

c.Logger.Debug(context.Background(),

"failed to close coordinator client during sync destinations",

slog.Error(cErr))

}

c.coordination = nil

}

}()

for dest := range toAdd {

c.Coordinatee.SetTunnelDestination(dest)

err = c.coordination.SendRequest(

&proto.CoordinateRequest{

AddTunnel: &proto.CoordinateRequest_Tunnel{Id: UUIDToByteSlice(dest)},

})

if err != nil {

return

}

}

for dest := range toRemove {

err = c.coordination.SendRequest(

&proto.CoordinateRequest{

RemoveTunnel: &proto.CoordinateRequest_Tunnel{Id: UUIDToByteSlice(dest)},

})

if err != nil {

return

}

}

}

// NewAgentCoordinationController creates a CoordinationController for Coder Agents, which never

// create tunnels and always send ReadyToHandshake acknowledgements.

func NewAgentCoordinationController(

logger slog.Logger, coordinatee Coordinatee,

) CoordinationController {

return &BasicCoordinationController{

Logger: logger,

Coordinatee: coordinatee,

SendAcks: true,

Initiator: codersdk.DisconnectInitiatorAgent,

Direction: codersdk.ConnectionDirectionServerToAgent,

}

}

type inMemoryCoordClient struct {

sync.Mutex

ctx context.Context

cancel context.CancelFunc

closed bool

logger slog.Logger

resps <-chan *proto.CoordinateResponse

reqs chan<- *proto.CoordinateRequest

}

func (c *inMemoryCoordClient) Close() error {

c.cancel()

c.Lock()

defer c.Unlock()

if c.closed {

return nil

}

c.closed = true

close(c.reqs)

return nil

}

func (c *inMemoryCoordClient) Send(request *proto.CoordinateRequest) error {

c.Lock()

defer c.Unlock()

if c.closed {

return drpc.ClosedError.New("in-memory coordinator client closed")

}

select {

case c.reqs <- request:

return nil

case <-c.ctx.Done():

return drpc.ClosedError.New("in-memory coordinator client closed")

}

}

func (c *inMemoryCoordClient) Recv() (*proto.CoordinateResponse, error) {

select {

case resp, ok := <-c.resps:

if ok {

return resp, nil

}

// response from Coordinator was closed, so close the send direction as well, so that the

// Coordinator won't be waiting for us while shutting down.

_ = c.Close()

return nil, io.EOF

case <-c.ctx.Done():

return nil, drpc.ClosedError.New("in-memory coord client closed")

}

}

// NewInMemoryCoordinatorClient creates a coordination client that uses channels to connect to a

// local Coordinator. (The typical alternative is a DRPC-based client.)

func NewInMemoryCoordinatorClient(

logger slog.Logger,

clientID uuid.UUID,

auth CoordinateeAuth,

coordinator Coordinator,

) CoordinatorClient {

logger = logger.With(slog.F("client_id", clientID))

c := &inMemoryCoordClient{logger: logger}

c.ctx, c.cancel = context.WithCancel(context.Background())

// use the background context since we will depend exclusively on closing the req channel to

// tell the coordinator we are done.

c.reqs, c.resps = coordinator.Coordinate(context.Background(),

clientID, fmt.Sprintf("inmemory%s", clientID),

auth,

)

return c

}

type DERPMapSetter interface {

SetDERPMap(derpMap *tailcfg.DERPMap)

}

type basicDERPController struct {

logger slog.Logger

rewriter DERPMapRewriter // optional

setter DERPMapSetter

}

func (b *basicDERPController) New(client DERPClient) CloserWaiter {

l := &derpSetLoop{

logger: b.logger,

rewriter: b.rewriter,

setter: b.setter,

client: client,

errChan: make(chan error, 1),

recvLoopDone: make(chan struct{}),

}

go l.recvLoop()

return l

}

// NewBasicDERPController creates a DERP controller that rewrites the DERP map

// with the provided rewriter before setting it on the provided setter.

//

// The rewriter is optional and can be nil.

func NewBasicDERPController(logger slog.Logger, rewriter DERPMapRewriter, setter DERPMapSetter) DERPController {

return &basicDERPController{

logger: logger,

rewriter: rewriter,

setter: setter,

}

}

type derpSetLoop struct {

logger slog.Logger

rewriter DERPMapRewriter // optional

setter DERPMapSetter

client DERPClient

sync.Mutex

closed bool

errChan chan error

recvLoopDone chan struct{}

}

func (l *derpSetLoop) Close(ctx context.Context) error {

l.Lock()

defer l.Unlock()

if l.closed {

select {

case <-ctx.Done():

return ctx.Err()

case <-l.recvLoopDone:

return nil

}

}

l.closed = true

cErr := l.client.Close()

select {

case <-ctx.Done():

return ctx.Err()

case <-l.recvLoopDone:

return cErr

}

}

func (l *derpSetLoop) Wait() <-chan error {

return l.errChan

}

func (l *derpSetLoop) recvLoop() {

defer close(l.recvLoopDone)

for {

dm, err := l.client.Recv()

if err != nil {

l.logger.Debug(context.Background(), "failed to receive DERP message", slog.Error(err))

select {

case l.errChan <- err:

default:

}

return

}

l.logger.Debug(context.Background(), "got new DERP Map", slog.F("derp_map", dm))

if l.rewriter != nil {

l.rewriter.RewriteDERPMap(dm)

}

l.setter.SetDERPMap(dm)

}

}

type BasicTelemetryController struct {

logger slog.Logger

sync.Mutex

client TelemetryClient

unavailable bool

}

func (b *BasicTelemetryController) New(client TelemetryClient) {

b.Lock()

defer b.Unlock()

b.client = client

b.unavailable = false

b.logger.Debug(context.Background(), "new telemetry client connected to controller")

}

func (b *BasicTelemetryController) SendTelemetryEvent(event *proto.TelemetryEvent) {

b.Lock()

if b.client == nil {

b.Unlock()

b.logger.Debug(context.Background(),

"telemetry event dropped; no client", slog.F("event", event))

return

}

if b.unavailable {

b.Unlock()

b.logger.Debug(context.Background(),

"telemetry event dropped; unavailable", slog.F("event", event))

return

}

client := b.client

b.Unlock()

unavailable := sendTelemetry(b.logger, client, event)

if unavailable {

b.Lock()

defer b.Unlock()

if b.client == client {

b.unavailable = true

}

}

}

func NewBasicTelemetryController(logger slog.Logger) *BasicTelemetryController {

return &BasicTelemetryController{logger: logger}

}

var (

_ TelemetrySink = &BasicTelemetryController{}

_ TelemetryController = &BasicTelemetryController{}

)

func sendTelemetry(

logger slog.Logger, client TelemetryClient, event *proto.TelemetryEvent,

) (

unavailable bool,

) {

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

defer cancel()

_, err := client.PostTelemetry(ctx, &proto.TelemetryRequest{

Events: []*proto.TelemetryEvent{event},

})

if IsDRPCUnimplementedError(err) {

logger.Debug(

context.Background(),

"attempted to send telemetry to a server that doesn't support it",

slog.Error(err),

)

return true

} else if err != nil {

logger.Warn(

context.Background(),

"failed to post telemetry event",

slog.F("event", event), slog.Error(err),

)

}

return false

}

// IsDRPCUnimplementedError returns true if the error indicates the RPC called is not implemented

// by the server.

func IsDRPCUnimplementedError(err error) bool {

return drpcerr.Code(err) == drpcerr.Unimplemented ||

drpc.ProtocolError.Has(err) &&

strings.Contains(err.Error(), "unknown rpc: ")

}

type basicResumeTokenController struct {

logger slog.Logger

sync.Mutex

token *proto.RefreshResumeTokenResponse

refresher *basicResumeTokenRefresher

// for testing

clock quartz.Clock

}

func (b *basicResumeTokenController) New(client ResumeTokenClient) CloserWaiter {

b.Lock()

defer b.Unlock()

if b.refresher != nil {

cErr := b.refresher.Close(context.Background())

if cErr != nil {

b.logger.Debug(context.Background(), "closed previous refresher", slog.Error(cErr))

}

}

b.refresher = newBasicResumeTokenRefresher(b.logger, b.clock, b, client)

return b.refresher

}

func (b *basicResumeTokenController) Token() (string, bool) {

b.Lock()

defer b.Unlock()

if b.token == nil {

return "", false

}

if b.token.ExpiresAt.AsTime().Before(b.clock.Now()) {

return "", false

}

return b.token.Token, true

}

func NewBasicResumeTokenController(logger slog.Logger, clock quartz.Clock) ResumeTokenController {

return &basicResumeTokenController{

logger: logger,

clock: clock,

}

}

type basicResumeTokenRefresher struct {

logger slog.Logger

ctx context.Context

cancel context.CancelFunc

ctrl *basicResumeTokenController

client ResumeTokenClient

errCh chan error

sync.Mutex

closed bool

timer *quartz.Timer

}

func (r *basicResumeTokenRefresher) Close(_ context.Context) error {

r.cancel()

r.Lock()

defer r.Unlock()

if r.closed {

return nil

}

r.closed = true

r.timer.Stop()

select {

case r.errCh <- nil:

default: // already have an error

}

return nil

}

func (r *basicResumeTokenRefresher) Wait() <-chan error {

return r.errCh

}

const never time.Duration = math.MaxInt64

func newBasicResumeTokenRefresher(

logger slog.Logger, clock quartz.Clock,

ctrl *basicResumeTokenController, client ResumeTokenClient,

) *basicResumeTokenRefresher {

r := &basicResumeTokenRefresher{

logger: logger,

ctrl: ctrl,

client: client,

errCh: make(chan error, 1),

}

r.ctx, r.cancel = context.WithCancel(context.Background())

r.timer = clock.AfterFunc(never, r.refresh, "basicResumeTokenRefresher")

go r.refresh()

return r

}

func (r *basicResumeTokenRefresher) refresh() {

if r.ctx.Err() != nil {

return // context done, no need to refresh

}

res, err := r.client.RefreshResumeToken(r.ctx, &proto.RefreshResumeTokenRequest{})

if xerrors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded) {

// these can only come from being closed, no need to log

select {

case r.errCh <- nil:

default: // already have an error

}

return

}

if IsDRPCUnimplementedError(err) {

r.logger.Info(r.ctx, "resume token is not supported by the server")

select {

case r.errCh <- nil:

default: // already have an error

}

return

} else if err != nil {

r.logger.Error(r.ctx, "error refreshing coordinator resume token", slog.Error(err))

select {

case r.errCh <- err:

default: // already have an error

}

return

}

r.logger.Debug(r.ctx, "refreshed coordinator resume token",

slog.F("expires_at", res.GetExpiresAt()),

slog.F("refresh_in", res.GetRefreshIn()),

)

r.ctrl.Lock()

if r.ctrl.refresher == r { // don't overwrite if we're not the current refresher

r.ctrl.token = res

} else {

r.logger.Debug(context.Background(), "not writing token because we have a new client")

}

r.ctrl.Unlock()

dur := res.RefreshIn.AsDuration()

if dur <= 0 {

// A sensible delay to refresh again.

dur = 30 * time.Minute

}

r.Lock()

defer r.Unlock()

if r.closed {

return

}

r.timer.Reset(dur, "basicResumeTokenRefresher", "refresh")

}

type TunnelAllWorkspaceUpdatesController struct {

coordCtrl *TunnelSrcCoordController

dnsHostSetter DNSHostsSetter

dnsNameOptions DNSNameOptions

updateHandler UpdatesHandler

ownerUsername string

logger slog.Logger

mu sync.Mutex

updater *tunnelUpdater

}

type Workspace struct {

ID uuid.UUID

Name string

Status proto.Workspace_Status

ownerUsername string

agents map[uuid.UUID]*Agent

}

func (w *Workspace) Clone() Workspace {

agents := make(map[uuid.UUID]*Agent, len(w.agents))

for k, v := range w.agents {

clone := v.Clone()

agents[k] = &clone

}

return Workspace{

ID: w.ID,

Name: w.Name,

Status: w.Status,

ownerUsername: w.ownerUsername,

agents: agents,

}

}

type DNSNameOptions struct {

Suffix string

}

// updateDNSNames updates the DNS names for all agents in the workspace.

// DNS hosts must be all lowercase, or the resolver won't be able to find them.

// Usernames are globally unique & case-insensitive.

// Workspace names are unique per-user & case-insensitive.

// Agent names are unique per-workspace & case-insensitive.

func (w *Workspace) updateDNSNames(options DNSNameOptions) error {

wsName := strings.ToLower(w.Name)