command.go

// Package command implements helpers useful for when building cobra commands.
package command

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	"github.com/blang/semver"
	"github.com/logrusorgru/aurora"
	"github.com/spf13/cobra"
	"github.com/superfly/flyctl/api"
	"github.com/superfly/flyctl/iostreams"

	"github.com/superfly/flyctl/client"
	"github.com/superfly/flyctl/internal/appconfig"
	"github.com/superfly/flyctl/internal/buildinfo"
	"github.com/superfly/flyctl/internal/cache"
	"github.com/superfly/flyctl/internal/cmdutil/preparers"
	"github.com/superfly/flyctl/internal/config"
	"github.com/superfly/flyctl/internal/env"
	"github.com/superfly/flyctl/internal/flag"
	"github.com/superfly/flyctl/internal/logger"
	"github.com/superfly/flyctl/internal/metrics"
	"github.com/superfly/flyctl/internal/state"
	"github.com/superfly/flyctl/internal/task"
	"github.com/superfly/flyctl/internal/update"
)

type Runner func(context.Context) error

func New(usage, short, long string, fn Runner, p ...preparers.Preparer) *cobra.Command {
	return &cobra.Command{
		Use:   usage,
		Short: short,
		Long:  long,
		RunE:  newRunE(fn, p...),
	}
}

// Preparers are split between here and the preparers package because
// tab-completion needs to run *some* of them, and importing this package from there
// would create a circular dependency. Likewise, if *all* the preparers were in the preparers module,
// that would also create a circular dependency.
// I don't like this, but it's shippable until someone else fixes it
var commonPreparers = []preparers.Preparer{
	preparers.ApplyAliases,
	determineHostname,
	determineWorkingDir,
	preparers.DetermineUserHomeDir,
	preparers.DetermineConfigDir,
	ensureConfigDirExists,
	ensureConfigDirPerms,
	loadCache,
	preparers.LoadConfig,
	initTaskManager,
	startQueryingForNewRelease,
	promptToUpdate,
	preparers.InitClient,
	killOldAgent,
	recordMetricsCommandContext,
}

func sendOsMetric(ctx context.Context, state string) {
	// Send /runs/[os_name]/[state]
	osName := ""
	switch runtime.GOOS {
	case "darwin":
		osName = "macos"
	case "linux":
		osName = "linux"
	case "windows":
		osName = "windows"
	default:
		osName = "other"
	}
	metrics.SendNoData(ctx, fmt.Sprintf("runs/%s/%s", osName, state))
}

func newRunE(fn Runner, preparers ...preparers.Preparer) func(*cobra.Command, []string) error {
	if fn == nil {
		return nil
	}

	return func(cmd *cobra.Command, _ []string) (err error) {
		ctx := cmd.Context()
		ctx = NewContext(ctx, cmd)
		ctx = flag.NewContext(ctx, cmd.Flags())

		// run the common preparers
		if ctx, err = prepare(ctx, commonPreparers...); err != nil {
			return
		}

		sendOsMetric(ctx, "started")
		defer func() {
			if err == nil {
				sendOsMetric(ctx, "successful")
			}
		}()

		// run the preparers specific to the command
		if ctx, err = prepare(ctx, preparers...); err != nil {
			return
		}

		// run the command
		if err = fn(ctx); err == nil {
			// and finally, run the finalizer
			finalize(ctx)
		}

		return
	}
}

func prepare(parent context.Context, preparers ...preparers.Preparer) (ctx context.Context, err error) {
	ctx = parent

	for _, p := range preparers {
		if ctx, err = p(ctx); err != nil {
			break
		}
	}

	return
}

func finalize(ctx context.Context) {
	// shutdown async tasks
	task.FromContext(ctx).Shutdown()

	// flush the cache to disk if required
	if c := cache.FromContext(ctx); c.Dirty() {
		path := filepath.Join(state.ConfigDirectory(ctx), cache.FileName)

		if err := c.Save(path); err != nil {
			logger.FromContext(ctx).
				Warnf("failed saving cache to %s: %v", path, err)
		}
	}
}

func determineHostname(ctx context.Context) (context.Context, error) {
	h, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("failed determining hostname: %w", err)
	}

	logger.FromContext(ctx).
		Debugf("determined hostname: %q", h)

	return state.WithHostname(ctx, h), nil
}

func determineWorkingDir(ctx context.Context) (context.Context, error) {
	wd, err := os.Getwd()
	if err != nil {
		return nil, fmt.Errorf("failed determining working directory: %w", err)
	}

	logger.FromContext(ctx).
		Debugf("determined working directory: %q", wd)

	return state.WithWorkingDirectory(ctx, wd), nil
}

func ensureConfigDirExists(ctx context.Context) (context.Context, error) {
	dir := state.ConfigDirectory(ctx)

	switch fi, err := os.Stat(dir); {
	case errors.Is(err, fs.ErrNotExist):
		if err := os.MkdirAll(dir, 0o700); err != nil {
			return nil, fmt.Errorf("failed creating config directory: %w", err)
		}
	case err != nil:
		return nil, fmt.Errorf("failed stat-ing config directory: %w", err)
	case !fi.IsDir():
		return nil, fmt.Errorf("the path to the config directory (%s) is occupied by not a directory", dir)
	}

	logger.FromContext(ctx).
		Debug("ensured config directory exists.")

	return ctx, nil
}

func ensureConfigDirPerms(parent context.Context) (ctx context.Context, err error) {
	defer func() {
		if err != nil {
			ctx = nil
			err = fmt.Errorf("failed ensuring config directory perms: %w", err)

			return
		}

		logger.FromContext(ctx).
			Debug("ensured config directory perms.")
	}()

	ctx = parent
	dir := state.ConfigDirectory(parent)

	var f *os.File
	if f, err = os.CreateTemp(dir, "perms.*"); err != nil {
		return
	}
	defer func() {
		if e := os.Remove(f.Name()); err == nil {
			err = e
		}
	}()

	err = f.Close()

	return
}

func loadCache(ctx context.Context) (context.Context, error) {
	logger := logger.FromContext(ctx)

	path := filepath.Join(state.ConfigDirectory(ctx), cache.FileName)

	c, err := cache.Load(path)
	if err != nil {
		c = cache.New()

		if !errors.Is(err, fs.ErrNotExist) {
			logger.Warnf("failed loading cache file from %s: %v", path, err)
		}
	}

	logger.Debug("cache loaded.")

	return cache.NewContext(ctx, c), nil
}

func initTaskManager(ctx context.Context) (context.Context, error) {
	tm := task.New(ctx)

	logger.FromContext(ctx).Debug("initialized task manager.")

	return task.NewContext(ctx, tm), nil
}

func IsMachinesPlatform(ctx context.Context, appName string) (bool, error) {
	apiClient := client.FromContext(ctx).API()
	app, err := apiClient.GetAppBasic(ctx, appName)
	if err != nil {
		return false, fmt.Errorf("failed to retrieve app: %w", err)
	}

	return app.PlatformVersion == appconfig.MachinesPlatform, nil
}

func startQueryingForNewRelease(ctx context.Context) (context.Context, error) {
	logger := logger.FromContext(ctx)

	cache := cache.FromContext(ctx)
	if !update.Check() || time.Since(cache.LastCheckedAt()) < time.Hour {
		logger.Debug("skipped querying for new release")

		return ctx, nil
	}

	channel := cache.Channel()
	tm := task.FromContext(ctx)

	tm.Run(func(parent context.Context) {
		ctx, cancel := context.WithTimeout(parent, time.Second)
		defer cancel()

		switch r, err := update.LatestRelease(ctx, channel); {
		case err == nil:
			if r == nil {
				break
			}

			cache.SetLatestRelease(channel, r)

			logger.Debugf("querying for release resulted to %v", r.Version)
		case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
			break
		default:
			logger.Warnf("failed querying for new release: %v", err)
		}
	})

	logger.Debug("started querying for new release")

	return ctx, nil
}

// shouldIgnore allows a preparer to disable itself for specific commands
// E.g. `shouldIgnore([][]string{{"version", "upgrade"}, {"machine", "status"}})`
// would return true for "fly version upgrade" and "fly machine status"
func shouldIgnore(ctx context.Context, cmds [][]string) bool {
	cmd := FromContext(ctx)
	for _, ignoredCmd := range cmds {
		match := true
		currentCmd := cmd
		// The shape of the ignoredCmd slice is something like ["version", "upgrade"],
		// but we're walking up the tree from the end, so we have to iterate that in reverse
		for i := len(ignoredCmd) - 1; i >= 0; i-- {
			if !currentCmd.HasParent() || currentCmd.Use != ignoredCmd[i] {
				match = false
				break
			}
			currentCmd = currentCmd.Parent()
		}
		// Ensure that we have the root node, so that a filter on ["y"] wouldn't be tripped by "fly x y"
		if match {
			if !currentCmd.HasParent() {
				return true
			}
		}
	}
	return false
}

func promptToUpdate(ctx context.Context) (context.Context, error) {
	cfg := config.FromContext(ctx)
	if cfg.JSONOutput || shouldIgnore(ctx, [][]string{
		{"version", "upgrade"},
	}) {
		return ctx, nil
	}

	if !update.Check() {
		return ctx, nil
	}

	c := cache.FromContext(ctx)

	r := c.LatestRelease()
	if r == nil {
		return ctx, nil
	}

	logger := logger.FromContext(ctx)

	current := buildinfo.Info().Version

	switch latest, err := semver.ParseTolerant(r.Version); {
	case err != nil:
		logger.Warnf("error parsing version number '%s': %s", r.Version, err)

		return ctx, nil
	case latest.LTE(current):
		return ctx, nil
	}

	io := iostreams.FromContext(ctx)
	colorize := io.ColorScheme()

	fmt.Fprintln(io.ErrOut, colorize.Yellow(fmt.Sprintf("Update available %s -> %s.", current, r.Version)))

	if cfg.AutoUpdate && !env.IsCI() {
		fmt.Fprintln(io.ErrOut, colorize.Green("Automatically updating..."))

		if err := update.UpgradeAndRelaunch(ctx, io, r.Prerelease); err != nil {
			fmt.Fprintf(io.ErrOut, fmt.Sprintf("failed to update: %s\n", err))

			return ctx, nil
		}
	} else {
		fmt.Fprintln(io.ErrOut, colorize.Yellow(fmt.Sprintf("Run \"%s\" to upgrade.", colorize.Bold(buildinfo.Name()+" version upgrade"))))
	}

	return ctx, nil
}

func PromptToMigrate(ctx context.Context, app *api.AppCompact) {
	if app.PlatformVersion == "nomad" {
		config := appconfig.ConfigFromContext(ctx)
		if config != nil {
			io := iostreams.FromContext(ctx)
			fmt.Fprintf(io.ErrOut, "%s Apps v1 Platform is deprecated. We recommend migrating your app with:\nfly migrate-to-v2 -c %s\n", aurora.Yellow("WARN"), config.ConfigFilePath())
		}
	}
}

func killOldAgent(ctx context.Context) (context.Context, error) {
	path := filepath.Join(state.ConfigDirectory(ctx), "agent.pid")

	data, err := os.ReadFile(path)
	if errors.Is(err, fs.ErrNotExist) {
		return ctx, nil // no old agent running or can't access that file
	} else if err != nil {
		return nil, fmt.Errorf("failed reading old agent's PID file: %w", err)
	}

	pid, err := strconv.Atoi(string(data))
	if err != nil {
		return nil, fmt.Errorf("failed determining old agent's PID: %w", err)
	}

	logger := logger.FromContext(ctx)
	unlink := func() (err error) {
		if err = os.Remove(path); err != nil {
			err = fmt.Errorf("failed removing old agent's PID file: %w", err)

			return
		}

		logger.Debug("removed old agent's PID file.")

		return
	}

	p, err := os.FindProcess(pid)
	if err != nil {
		return nil, fmt.Errorf("failed retrieving old agent's process: %w", err)
	} else if p == nil {
		return ctx, unlink()
	}

	if err := p.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
		return nil, fmt.Errorf("failed killing old agent process: %w", err)
	}

	logger.Debugf("killed old agent (PID: %d)", pid)

	if err := unlink(); err != nil {
		return nil, err
	}

	time.Sleep(time.Second) // we've killed and removed the pid file

	return ctx, nil
}

func recordMetricsCommandContext(ctx context.Context) (context.Context, error) {
	metrics.RecordCommandContext(ctx)
	return ctx, nil
}

func ExcludeFromMetrics(ctx context.Context) (context.Context, error) {
	metrics.Enabled = false
	return ctx, nil
}

// RequireSession is a Preparer which makes sure a session exists.
func RequireSession(ctx context.Context) (context.Context, error) {
	if !client.FromContext(ctx).Authenticated() {
		return nil, client.ErrNoAuthToken
	}

	return ctx, nil
}

// LoadAppConfigIfPresent is a Preparer which loads the application's
// configuration file from the path the user has selected via command line args
// or the current working directory.
func LoadAppConfigIfPresent(ctx context.Context) (context.Context, error) {
	// Shortcut to avoid unmarshaling and querying Web when
	// LoadAppConfigIfPresent is chained with RequireAppName
	if cfg := appconfig.ConfigFromContext(ctx); cfg != nil {
		return ctx, nil
	}

	logger := logger.FromContext(ctx)
	for _, path := range appConfigFilePaths(ctx) {
		switch cfg, err := appconfig.LoadConfig(path); {
		case err == nil:
			logger.Debugf("app config loaded from %s", path)

			// Query Web API for platform version
			platformVersion, _ := determinePlatform(ctx, cfg.AppName)
			if platformVersion != "" {
				err := cfg.SetPlatformVersion(platformVersion)
				if err != nil {
					logger.Warnf("WARNING the config file at '%s' is not valid: %s", path, err)
				}
			}

			return appconfig.WithConfig(ctx, cfg), nil // we loaded a configuration file
		case errors.Is(err, fs.ErrNotExist):
			logger.Debugf("no app config found at %s; skipped.", path)
			continue
		default:
			return nil, fmt.Errorf("failed loading app config from %s: %w", path, err)
		}
	}

	return ctx, nil
}

func determinePlatform(ctx context.Context, appName string) (string, error) {
	client := client.FromContext(ctx)
	if appName == "" {
		return "", fmt.Errorf("Can't determine platform without an application name")
	}

	basicApp, err := client.API().GetAppBasic(ctx, appName)
	if err != nil {
		return "", err
	}
	return basicApp.PlatformVersion, nil
}

// appConfigFilePaths returns the possible paths at which we may find a fly.toml
// in order of preference. it takes into consideration whether the user has
// specified a command-line path to a config file.
func appConfigFilePaths(ctx context.Context) (paths []string) {
	if p := flag.GetAppConfigFilePath(ctx); p != "" {
		paths = append(paths, p, filepath.Join(p, appconfig.DefaultConfigFileName))

		return
	}

	wd := state.WorkingDirectory(ctx)
	paths = append(paths, filepath.Join(wd, appconfig.DefaultConfigFileName))

	return
}

var ErrRequireAppName = fmt.Errorf("the config for your app is missing an app name, add an app field to the fly.toml file or specify with the -a flag`")

// RequireAppName is a Preparer which makes sure the user has selected an
// application name via command line arguments, the environment or an application
// config file (fly.toml). It embeds LoadAppConfigIfPresent.
func RequireAppName(ctx context.Context) (context.Context, error) {
	ctx, err := LoadAppConfigIfPresent(ctx)
	if err != nil {
		return nil, err
	}

	name := flag.GetApp(ctx)
	if name == "" {
		// if there's no flag present, first consult with the environment
		if name = env.First("FLY_APP"); name == "" {
			// and then with the config file (if any)
			if cfg := appconfig.ConfigFromContext(ctx); cfg != nil {
				name = cfg.AppName
			}
		}
	}

	if name == "" {
		return nil, ErrRequireAppName
	}

	return appconfig.WithName(ctx, name), nil
}

// RequireAppNameNoFlag is a Preparer which makes sure the user has selected an
// application name via the environment or an application
// config file (fly.toml). It embeds LoadAppConfigIfPresent.
//
// Identical to RequireAppName but does not check for the --app flag.
func RequireAppNameNoFlag(ctx context.Context) (context.Context, error) {
	ctx, err := LoadAppConfigIfPresent(ctx)
	if err != nil {
		return nil, err
	}

	// First consult with the environment
	name := env.First("FLY_APP")
	if name == "" {
		// and then with the config file (if any)
		if cfg := appconfig.ConfigFromContext(ctx); cfg != nil {
			name = cfg.AppName
		}
	}

	if name == "" {
		return nil, ErrRequireAppName
	}

	return appconfig.WithName(ctx, name), nil
}

// LoadAppNameIfPresent is a Preparer which adds app name if the user has used --app or there appConfig
// but unlike RequireAppName it does not error if the user has not specified an app name.
func LoadAppNameIfPresent(ctx context.Context) (context.Context, error) {
	localCtx, err := RequireAppName(ctx)

	if errors.Is(err, ErrRequireAppName) {
		return appconfig.WithName(ctx, ""), nil
	}

	return localCtx, err
}

// LoadAppNameIfPresentNoFlag is like LoadAppNameIfPresent, but it does not check for the --app flag.
func LoadAppNameIfPresentNoFlag(ctx context.Context) (context.Context, error) {
	localCtx, err := RequireAppNameNoFlag(ctx)

	if errors.Is(err, ErrRequireAppName) {
		return appconfig.WithName(ctx, ""), nil
	}

	return localCtx, err
}

func ChangeWorkingDirectoryToFirstArgIfPresent(ctx context.Context) (context.Context, error) {
	wd := flag.FirstArg(ctx)
	if wd == "" {
		return ctx, nil
	}
	return ChangeWorkingDirectory(ctx, wd)
}

func ChangeWorkingDirectory(ctx context.Context, wd string) (context.Context, error) {
	if !filepath.IsAbs(wd) {
		p, err := filepath.Abs(wd)
		if err != nil {
			return nil, fmt.Errorf("failed converting %s to an absolute path: %w", wd, err)
		}
		wd = p
	}

	if err := os.Chdir(wd); err != nil {
		return nil, fmt.Errorf("failed changing working directory: %w", err)
	}

	return state.WithWorkingDirectory(ctx, wd), nil
}