src/pkg/kernel/deploy.go

package kernel

import (
	"context"
	_ "embed"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

//go:embed scripts/cos-kernel-deploy
var kernelDeployScript string

// CopyFileToGcpInstance copies a file from a local path to a remote path on a GCP instance.
//
// The copy is done via `gcloud compute ssh`.
//
// The gcloudArgs args are passed directly to `gcloud compute ssh` in order to specify
// the instance the file should be copied to.
func CopyFileToGcpInstance(ctx context.Context, localPath string, remotePath string, gcloudArgs []string) error {
	fp, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("could not read local file %s: %w", localPath, err)
	}

	// This is a gross hack. If we used `compute scp` instead of `compute ssh`, we'd have to inspect
	// the user-provided gcloud arguments in order to concatenate the remote path to the instance
	// name. By using tee and consuming the input file through stdin, we just copy the file directly
	// to the output via ssh, sidestepping the issue.
	gcloudCmdArgs := append([]string{"compute", "ssh", fmt.Sprintf("--command=tee '%s'", remotePath)}, gcloudArgs...)

	cmd := exec.CommandContext(ctx, "gcloud", gcloudCmdArgs...)
	cmd.Stdin = fp
	// We don't map stdout since it'll just be the file's contents.
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to copy file to GCP instance: %w", err)
	}

	return nil
}

// DeployKernelToGcpInstance replaces the kernel on a GCP instance with a kernel from a remote path on that instance.
//
// The gcloudArgs args are passed directly to `gcloud compute ssh` in order to specify
// the instance where the kernel will be deployed.
func DeployKernelToGcpInstance(ctx context.Context, path string, gcloudArgs []string) error {
	// This allows us to read the script from stdin while also passing positional args.
	gcloudCmdArgs := append([]string{"compute", "ssh", fmt.Sprintf("--command=sudo bash -s -- '%s'", path)}, gcloudArgs...)

	cmd := exec.CommandContext(ctx, "gcloud", gcloudCmdArgs...)
	cmd.Stdin = strings.NewReader(kernelDeployScript)
	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to deploy COS kernel: %w", err)
	}

	return nil
}

// DeployKernelToLocalInstance replaces the kernel on a local instance with the kernel at the provided path.
func DeployKernelToLocalInstance(path string) error {
	cmd := exec.Command("bash", "-s", "--", path)
	cmd.Stdin = strings.NewReader(kernelDeployScript)

	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to set dev mode: %w", err)
	}

	return nil
}