src/cmd/cos_kernel/actions/utils.go

package actions

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/golang/glog"
)

// SSHOptions configures retry and timeout behavior for remotely executed SSH commands.
type SSHOptions struct {
	// Number of times to retry the SSH command if the initial attempt fails. A value of 0 means the
	// command will only be attmpted once with no retries.
	MaxRetries int
	// Specifies the delay to occur before attempting to execute the SSH command again.
	RetryDelay time.Duration
	// Specifies the maximum duration allowed for an SSH command. If the command's runtime exceeds
	// this duration, the command will be terminated.
	Timeout time.Duration
}

// runGcloudComputeSSH wraps the `gcloud compute ssh` call to manage timeouts and reconnections
// associated with the command. Defaults to 5 retries at a 10 second delay with a timeout of 2
// minutes. It accepts SSHOptions to override these defaults.
func runGcloudComputeSSH(ctx context.Context, args []string, opts *SSHOptions) (string, error) {
	if opts == nil {
		opts = &SSHOptions{
			MaxRetries: 5,
			RetryDelay: 10 * time.Second,
			Timeout:    2 * time.Minute,
		}
	}

	totalAttempts := opts.MaxRetries + 1
	var lastErr error

	for attempt := 1; attempt <= totalAttempts; attempt++ {
		out, err := executeGcloudAttempt(ctx, args, opts.Timeout)
		if err == nil {
			return out, nil
		}

		lastErr = err
		isRetryable := false

		if strings.Contains(err.Error(), "timed out") {
			isRetryable = true
		} else if exitErr, ok := errors.AsType[*exec.ExitError](err); ok && exitErr.ExitCode() == 255 {
			// Exit code 255 is standard for SSH failures. Other codes indicate gcloud CLI errors.
			isRetryable = true
		}

		if !isRetryable {
			return "", fmt.Errorf("fatal execution error (aborting retries): %w", err)
		}

		if opts.MaxRetries > 0 {
			glog.Warningf("attempt %d/%d failed: %v", attempt, totalAttempts, err)
		}

		if attempt < totalAttempts {
			glog.Warningf("retrying in %v", opts.RetryDelay)
			time.Sleep(opts.RetryDelay)
		}
	}

	return "", fmt.Errorf("failed to connect via gcloud compute ssh after %d attempts: %w", totalAttempts, lastErr)
}

// executeGcloudAttempt manages an exec call with context timeout support.
func executeGcloudAttempt(ctx context.Context, args []string, timeout time.Duration) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	cmdArgs := append([]string{"compute", "ssh"}, args...)
	cmd := exec.CommandContext(ctx, "gcloud", cmdArgs...)

	var outBuf bytes.Buffer
	cmd.Stdout = &outBuf
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	if ctx.Err() == context.DeadlineExceeded {
		return "", fmt.Errorf("ssh command timed out after %v", timeout)
	}

	return outBuf.String(), err
}

// executeScript executes the given scriptContent with `bash -s` and redirects the output to bash
// via stdin locally. It returns the captured stdout as a string.
func executeScript(scriptContent string) (string, error) {
	cmd := exec.Command("bash", "-s")
	cmd.Stdin = strings.NewReader(scriptContent)

	var outBuf bytes.Buffer

	cmd.Stdout = &outBuf
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	return outBuf.String(), err
}