coder/codersdk/provisionerdaemons.go

package codersdk

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/hashicorp/yamux"
	"golang.org/x/exp/maps"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/buildinfo"
	"github.com/coder/coder/v2/codersdk/drpcsdk"
	"github.com/coder/coder/v2/codersdk/wsjson"
	"github.com/coder/coder/v2/provisionerd/proto"
	"github.com/coder/coder/v2/provisionerd/runner"
	"github.com/coder/websocket"
)

type LogSource string

type LogLevel string

const (
	LogSourceProvisionerDaemon LogSource = "provisioner_daemon"
	LogSourceProvisioner       LogSource = "provisioner"

	LogLevelTrace LogLevel = "trace"
	LogLevelDebug LogLevel = "debug"
	LogLevelInfo  LogLevel = "info"
	LogLevelWarn  LogLevel = "warn"
	LogLevelError LogLevel = "error"
)

// ProvisionerDaemonStatus represents the status of a provisioner daemon.
type ProvisionerDaemonStatus string

// ProvisionerDaemonStatus enums.
const (
	ProvisionerDaemonOffline ProvisionerDaemonStatus = "offline"
	ProvisionerDaemonIdle    ProvisionerDaemonStatus = "idle"
	ProvisionerDaemonBusy    ProvisionerDaemonStatus = "busy"
)

func ProvisionerDaemonStatusEnums() []ProvisionerDaemonStatus {
	return []ProvisionerDaemonStatus{
		ProvisionerDaemonOffline,
		ProvisionerDaemonIdle,
		ProvisionerDaemonBusy,
	}
}

type ProvisionerDaemon struct {
	ID             uuid.UUID         `json:"id" format:"uuid" table:"id"`
	OrganizationID uuid.UUID         `json:"organization_id" format:"uuid" table:"organization id"`
	KeyID          uuid.UUID         `json:"key_id" format:"uuid" table:"-"`
	CreatedAt      time.Time         `json:"created_at" format:"date-time" table:"created at"`
	LastSeenAt     NullTime          `json:"last_seen_at,omitempty" format:"date-time" table:"last seen at"`
	Name           string            `json:"name" table:"name,default_sort"`
	Version        string            `json:"version" table:"version"`
	APIVersion     string            `json:"api_version" table:"api version"`
	Provisioners   []ProvisionerType `json:"provisioners" table:"-"`
	Tags           map[string]string `json:"tags" table:"tags"`

	// Optional fields.
	KeyName     *string                  `json:"key_name" table:"key name"`
	Status      *ProvisionerDaemonStatus `json:"status" enums:"offline,idle,busy" table:"status"`
	CurrentJob  *ProvisionerDaemonJob    `json:"current_job" table:"current job,recursive"`
	PreviousJob *ProvisionerDaemonJob    `json:"previous_job" table:"previous job,recursive"`
}

type ProvisionerDaemonJob struct {
	ID                  uuid.UUID            `json:"id" format:"uuid" table:"id"`
	Status              ProvisionerJobStatus `json:"status" enums:"pending,running,succeeded,canceling,canceled,failed" table:"status"`
	TemplateName        string               `json:"template_name" table:"template name"`
	TemplateIcon        string               `json:"template_icon" table:"template icon"`
	TemplateDisplayName string               `json:"template_display_name" table:"template display name"`
}

// MatchedProvisioners represents the number of provisioner daemons
// available to take a job at a specific point in time.
// Introduced in Coder version 2.18.0.
type MatchedProvisioners struct {
	// Count is the number of provisioner daemons that matched the given
	// tags. If the count is 0, it means no provisioner daemons matched the
	// requested tags.
	Count int `json:"count"`
	// Available is the number of provisioner daemons that are available to
	// take jobs. This may be less than the count if some provisioners are
	// busy or have been stopped.
	Available int `json:"available"`
	// MostRecentlySeen is the most recently seen time of the set of matched
	// provisioners. If no provisioners matched, this field will be null.
	MostRecentlySeen NullTime `json:"most_recently_seen,omitempty" format:"date-time"`
}

// ProvisionerJobStatus represents the at-time state of a job.
type ProvisionerJobStatus string

// Active returns whether the job is still active or not.
// It returns true if canceling as well, since the job isn't
// in an entirely inactive state yet.
func (p ProvisionerJobStatus) Active() bool {
	return p == ProvisionerJobPending ||
		p == ProvisionerJobRunning ||
		p == ProvisionerJobCanceling
}

const (
	ProvisionerJobPending   ProvisionerJobStatus = "pending"
	ProvisionerJobRunning   ProvisionerJobStatus = "running"
	ProvisionerJobSucceeded ProvisionerJobStatus = "succeeded"
	ProvisionerJobCanceling ProvisionerJobStatus = "canceling"
	ProvisionerJobCanceled  ProvisionerJobStatus = "canceled"
	ProvisionerJobFailed    ProvisionerJobStatus = "failed"
	ProvisionerJobUnknown   ProvisionerJobStatus = "unknown"
)

func ProvisionerJobStatusEnums() []ProvisionerJobStatus {
	return []ProvisionerJobStatus{
		ProvisionerJobPending,
		ProvisionerJobRunning,
		ProvisionerJobSucceeded,
		ProvisionerJobCanceling,
		ProvisionerJobCanceled,
		ProvisionerJobFailed,
		ProvisionerJobUnknown,
	}
}

// ProvisionerJobInput represents the input for the job.
type ProvisionerJobInput struct {
	TemplateVersionID *uuid.UUID `json:"template_version_id,omitempty" format:"uuid" table:"template version id"`
	WorkspaceBuildID  *uuid.UUID `json:"workspace_build_id,omitempty" format:"uuid" table:"workspace build id"`
	Error             string     `json:"error,omitempty" table:"-"`
}

// ProvisionerJobMetadata contains metadata for the job.
type ProvisionerJobMetadata struct {
	TemplateVersionName      string              `json:"template_version_name" table:"template version name"`
	TemplateID               uuid.UUID           `json:"template_id" format:"uuid" table:"template id"`
	TemplateName             string              `json:"template_name" table:"template name"`
	TemplateDisplayName      string              `json:"template_display_name" table:"template display name"`
	TemplateIcon             string              `json:"template_icon" table:"template icon"`
	WorkspaceID              *uuid.UUID          `json:"workspace_id,omitempty" format:"uuid" table:"workspace id"`
	WorkspaceName            string              `json:"workspace_name,omitempty" table:"workspace name"`
	WorkspaceBuildTransition WorkspaceTransition `json:"workspace_build_transition,omitempty" table:"workspace build transition"`
}

// ProvisionerJobType represents the type of job.
type ProvisionerJobType string

const (
	ProvisionerJobTypeTemplateVersionImport ProvisionerJobType = "template_version_import"
	ProvisionerJobTypeWorkspaceBuild        ProvisionerJobType = "workspace_build"
	ProvisionerJobTypeTemplateVersionDryRun ProvisionerJobType = "template_version_dry_run"
)

// JobErrorCode defines the error code returned by job runner.
type JobErrorCode string

const (
	RequiredTemplateVariables JobErrorCode = "REQUIRED_TEMPLATE_VARIABLES"
	InsufficientQuota         JobErrorCode = "INSUFFICIENT_QUOTA"
)

// JobIsMissingParameterErrorCode returns whether the error is a missing parameter error.
// This can indicate to consumers that they should check parameters.
func JobIsMissingParameterErrorCode(code JobErrorCode) bool {
	return string(code) == runner.MissingParameterErrorCode
}

// JobIsMissingRequiredTemplateVariableErrorCode returns whether the error is a missing a required template
// variable error. This can indicate to consumers that they need to provide required template variables.
func JobIsMissingRequiredTemplateVariableErrorCode(code JobErrorCode) bool {
	return string(code) == runner.RequiredTemplateVariablesErrorCode
}

// JobIsInsufficientQuotaErrorCode returns whether the error is an insufficient
// quota error. This can indicate to consumers that they should explain quota
// recovery options instead of treating the failure as a generic build error.
func JobIsInsufficientQuotaErrorCode(code JobErrorCode) bool {
	return string(code) == runner.InsufficientQuotaErrorCode
}

// ProvisionerJob describes the job executed by the provisioning daemon.
type ProvisionerJob struct {
	ID               uuid.UUID              `json:"id" format:"uuid" table:"id"`
	CreatedAt        time.Time              `json:"created_at" format:"date-time" table:"created at"`
	StartedAt        *time.Time             `json:"started_at,omitempty" format:"date-time" table:"started at"`
	CompletedAt      *time.Time             `json:"completed_at,omitempty" format:"date-time" table:"completed at"`
	CanceledAt       *time.Time             `json:"canceled_at,omitempty" format:"date-time" table:"canceled at"`
	Error            string                 `json:"error,omitempty" table:"error"`
	ErrorCode        JobErrorCode           `json:"error_code,omitempty" enums:"REQUIRED_TEMPLATE_VARIABLES,INSUFFICIENT_QUOTA" table:"error code"`
	Status           ProvisionerJobStatus   `json:"status" enums:"pending,running,succeeded,canceling,canceled,failed" table:"status"`
	WorkerID         *uuid.UUID             `json:"worker_id,omitempty" format:"uuid" table:"worker id"`
	WorkerName       string                 `json:"worker_name,omitempty" table:"worker name"`
	FileID           uuid.UUID              `json:"file_id" format:"uuid" table:"file id"`
	Tags             map[string]string      `json:"tags" table:"tags"`
	QueuePosition    int                    `json:"queue_position" table:"queue position"`
	QueueSize        int                    `json:"queue_size" table:"queue size"`
	OrganizationID   uuid.UUID              `json:"organization_id" format:"uuid" table:"organization id"`
	InitiatorID      uuid.UUID              `json:"initiator_id" format:"uuid" table:"initiator id"`
	Input            ProvisionerJobInput    `json:"input" table:"input,recursive_inline"`
	Type             ProvisionerJobType     `json:"type" table:"type"`
	AvailableWorkers []uuid.UUID            `json:"available_workers,omitempty" format:"uuid" table:"available workers"`
	Metadata         ProvisionerJobMetadata `json:"metadata" table:"metadata,recursive_inline"`
	LogsOverflowed   bool                   `json:"logs_overflowed" table:"logs overflowed"`
}

// ProvisionerJobLog represents the provisioner log entry annotated with source and level.
type ProvisionerJobLog struct {
	ID        int64     `json:"id"`
	CreatedAt time.Time `json:"created_at" format:"date-time"`
	Source    LogSource `json:"log_source"`
	Level     LogLevel  `json:"log_level" enums:"trace,debug,info,warn,error"`
	Stage     string    `json:"stage"`
	Output    string    `json:"output"`
}

// Text formats the log entry as human-readable text.
func (l ProvisionerJobLog) Text() string {
	var sb strings.Builder
	_, _ = sb.WriteString(l.CreatedAt.Format(time.RFC3339))
	_, _ = sb.WriteString(" [")
	_, _ = sb.WriteString(string(l.Level))
	_, _ = sb.WriteString("] [provisioner|")
	_, _ = sb.WriteString(l.Stage)
	_, _ = sb.WriteString("] ")
	_, _ = sb.WriteString(l.Output)
	return sb.String()
}

// provisionerJobLogsAfter streams logs that occurred after a specific time.
func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after int64) (<-chan ProvisionerJobLog, io.Closer, error) {
	afterQuery := ""
	if after != 0 {
		afterQuery = fmt.Sprintf("&after=%d", after)
	}
	followURL, err := c.URL.Parse(fmt.Sprintf("%s?follow%s", path, afterQuery))
	if err != nil {
		return nil, nil, err
	}
	httpClient := &http.Client{
		Transport: c.HTTPClient.Transport,
	}
	conn, res, err := websocket.Dial(ctx, followURL.String(), &websocket.DialOptions{
		HTTPClient: httpClient,
		HTTPHeader: http.Header{
			SessionTokenHeader: []string{c.SessionToken()},
		},
		CompressionMode: websocket.CompressionDisabled,
	})
	if err != nil {
		if res == nil {
			return nil, nil, err
		}
		return nil, nil, ReadBodyAsError(res)
	}
	d := wsjson.NewDecoder[ProvisionerJobLog](conn, websocket.MessageText, c.logger)
	return d.Chan(), d, nil
}

// ServeProvisionerDaemonRequest are the parameters to call ServeProvisionerDaemon with
// @typescript-ignore ServeProvisionerDaemonRequest
type ServeProvisionerDaemonRequest struct {
	// ID is a unique ID for a provisioner daemon.
	// Deprecated: this field has always been ignored.
	ID uuid.UUID `json:"id" format:"uuid"`
	// Name is the human-readable unique identifier for the daemon.
	Name string `json:"name" example:"my-cool-provisioner-daemon"`
	// Organization is the organization for the URL. If no orgID is provided,
	// then it is assumed to use the default organization.
	Organization uuid.UUID `json:"organization" format:"uuid"`
	// Provisioners is a list of provisioner types hosted by the provisioner daemon
	Provisioners []ProvisionerType `json:"provisioners"`
	// Tags is a map of key-value pairs that tag the jobs this provisioner daemon can handle
	Tags map[string]string `json:"tags"`
	// PreSharedKey is an authentication key to use on the API instead of the normal session token from the client.
	PreSharedKey string `json:"pre_shared_key"`
	// ProvisionerKey is an authentication key to use on the API instead of the normal session token from the client.
	ProvisionerKey string `json:"provisioner_key"`
}

// ServeProvisionerDaemon returns the gRPC service for a provisioner daemon
// implementation. The context is during dial, not during the lifetime of the
// client. Client should be closed after use.
func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisionerDaemonRequest) (proto.DRPCProvisionerDaemonClient, error) {
	orgParam := req.Organization.String()
	if req.Organization == uuid.Nil {
		orgParam = DefaultOrganization
	}

	serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/organizations/%s/provisionerdaemons/serve", orgParam))
	if err != nil {
		return nil, xerrors.Errorf("parse url: %w", err)
	}
	query := serverURL.Query()
	query.Add("version", proto.CurrentVersion.String())
	query.Add("name", req.Name)
	query.Add("version", proto.CurrentVersion.String())

	for _, provisioner := range req.Provisioners {
		query.Add("provisioner", string(provisioner))
	}
	for key, value := range req.Tags {
		query.Add("tag", fmt.Sprintf("%s=%s", key, value))
	}
	serverURL.RawQuery = query.Encode()
	httpClient := &http.Client{
		Transport: c.HTTPClient.Transport,
	}
	headers := http.Header{}

	headers.Set(BuildVersionHeader, buildinfo.Version())

	if req.ProvisionerKey != "" {
		headers.Set(ProvisionerDaemonKey, req.ProvisionerKey)
	}
	if req.PreSharedKey != "" {
		headers.Set(ProvisionerDaemonPSK, req.PreSharedKey)
	}
	if req.ProvisionerKey == "" && req.PreSharedKey == "" {
		// Use session token if we don't have a PSK or provisioner key.
		headers.Set(SessionTokenHeader, c.SessionToken())
	}

	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
		HTTPClient: httpClient,
		// Need to disable compression to avoid a data-race.
		CompressionMode: websocket.CompressionDisabled,
		HTTPHeader:      headers,
	})
	if err != nil {
		if res == nil {
			return nil, err
		}
		return nil, ReadBodyAsError(res)
	}
	// Align with the frame size of yamux.
	conn.SetReadLimit(256 * 1024)

	config := yamux.DefaultConfig()
	config.LogOutput = io.Discard
	// Use background context because caller should close the client.
	_, wsNetConn := WebsocketNetConn(context.Background(), conn, websocket.MessageBinary)
	session, err := yamux.Client(wsNetConn, config)
	if err != nil {
		_ = conn.Close(websocket.StatusGoingAway, "")
		_ = wsNetConn.Close()
		return nil, xerrors.Errorf("multiplex client: %w", err)
	}
	return proto.NewDRPCProvisionerDaemonClient(drpcsdk.MultiplexedConn(session)), nil
}

type ProvisionerKeyTags map[string]string

func (p ProvisionerKeyTags) String() string {
	keys := maps.Keys(p)
	slices.Sort(keys)
	tags := []string{}
	for _, key := range keys {
		tags = append(tags, fmt.Sprintf("%s=%s", key, p[key]))
	}
	return strings.Join(tags, " ")
}

type ProvisionerKey struct {
	ID             uuid.UUID          `json:"id" table:"-" format:"uuid"`
	CreatedAt      time.Time          `json:"created_at" table:"created at" format:"date-time"`
	OrganizationID uuid.UUID          `json:"organization" table:"-" format:"uuid"`
	Name           string             `json:"name" table:"name,default_sort"`
	Tags           ProvisionerKeyTags `json:"tags" table:"tags"`
	// HashedSecret - never include the access token in the API response
}

type ProvisionerKeyDaemons struct {
	Key     ProvisionerKey      `json:"key"`
	Daemons []ProvisionerDaemon `json:"daemons"`
}

const (
	ProvisionerKeyIDBuiltIn  = "00000000-0000-0000-0000-000000000001"
	ProvisionerKeyIDUserAuth = "00000000-0000-0000-0000-000000000002"
	ProvisionerKeyIDPSK      = "00000000-0000-0000-0000-000000000003"
)

var (
	ProvisionerKeyUUIDBuiltIn  = uuid.MustParse(ProvisionerKeyIDBuiltIn)
	ProvisionerKeyUUIDUserAuth = uuid.MustParse(ProvisionerKeyIDUserAuth)
	ProvisionerKeyUUIDPSK      = uuid.MustParse(ProvisionerKeyIDPSK)
)

const (
	ProvisionerKeyNameBuiltIn  = "built-in"
	ProvisionerKeyNameUserAuth = "user-auth"
	ProvisionerKeyNamePSK      = "psk"
)

func ReservedProvisionerKeyNames() []string {
	return []string{
		ProvisionerKeyNameBuiltIn,
		ProvisionerKeyNameUserAuth,
		ProvisionerKeyNamePSK,
	}
}

type CreateProvisionerKeyRequest struct {
	Name string            `json:"name"`
	Tags map[string]string `json:"tags"`
}

type CreateProvisionerKeyResponse struct {
	Key string `json:"key"`
}

// CreateProvisionerKey creates a new provisioner key for an organization.
func (c *Client) CreateProvisionerKey(ctx context.Context, organizationID uuid.UUID, req CreateProvisionerKeyRequest) (CreateProvisionerKeyResponse, error) {
	res, err := c.Request(ctx, http.MethodPost,
		fmt.Sprintf("/api/v2/organizations/%s/provisionerkeys", organizationID.String()),
		req,
	)
	if err != nil {
		return CreateProvisionerKeyResponse{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusCreated {
		return CreateProvisionerKeyResponse{}, ReadBodyAsError(res)
	}
	var resp CreateProvisionerKeyResponse
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// ListProvisionerKeys lists all provisioner keys for an organization.
func (c *Client) ListProvisionerKeys(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/organizations/%s/provisionerkeys", organizationID.String()),
		nil,
	)
	if err != nil {
		return nil, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return nil, ReadBodyAsError(res)
	}
	var resp []ProvisionerKey
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// GetProvisionerKey returns the provisioner key.
func (c *Client) GetProvisionerKey(ctx context.Context, pk string) (ProvisionerKey, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/provisionerkeys/%s", pk), nil,
		func(req *http.Request) {
			req.Header.Add(ProvisionerDaemonKey, pk)
		},
	)
	if err != nil {
		return ProvisionerKey{}, xerrors.Errorf("request to fetch provisioner key failed: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return ProvisionerKey{}, ReadBodyAsError(res)
	}
	var resp ProvisionerKey
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// ListProvisionerKeyDaemons lists all provisioner keys with their associated daemons for an organization.
func (c *Client) ListProvisionerKeyDaemons(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKeyDaemons, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/organizations/%s/provisionerkeys/daemons", organizationID.String()),
		nil,
	)
	if err != nil {
		return nil, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return nil, ReadBodyAsError(res)
	}
	var resp []ProvisionerKeyDaemons
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// DeleteProvisionerKey deletes a provisioner key.
func (c *Client) DeleteProvisionerKey(ctx context.Context, organizationID uuid.UUID, name string) error {
	res, err := c.Request(ctx, http.MethodDelete,
		fmt.Sprintf("/api/v2/organizations/%s/provisionerkeys/%s", organizationID.String(), name),
		nil,
	)
	if err != nil {
		return xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusNoContent {
		return ReadBodyAsError(res)
	}
	return nil
}

func ConvertWorkspaceStatus(jobStatus ProvisionerJobStatus, transition WorkspaceTransition) WorkspaceStatus {
	switch jobStatus {
	case ProvisionerJobPending:
		return WorkspaceStatusPending
	case ProvisionerJobRunning:
		switch transition {
		case WorkspaceTransitionStart:
			return WorkspaceStatusStarting
		case WorkspaceTransitionStop:
			return WorkspaceStatusStopping
		case WorkspaceTransitionDelete:
			return WorkspaceStatusDeleting
		}
	case ProvisionerJobSucceeded:
		switch transition {
		case WorkspaceTransitionStart:
			return WorkspaceStatusRunning
		case WorkspaceTransitionStop:
			return WorkspaceStatusStopped
		case WorkspaceTransitionDelete:
			return WorkspaceStatusDeleted
		}
	case ProvisionerJobCanceling:
		return WorkspaceStatusCanceling
	case ProvisionerJobCanceled:
		return WorkspaceStatusCanceled
	case ProvisionerJobFailed:
		return WorkspaceStatusFailed
	}

	// return error status since we should never get here
	return WorkspaceStatusFailed
}