uncloud/pkg/client/image.go

package client

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/charmbracelet/lipgloss"
	"github.com/docker/compose/v2/pkg/progress"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/image"
	dockerclient "github.com/docker/docker/client"
	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/docker/go-connections/nat"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/psviderski/uncloud/internal/docker"
	"github.com/psviderski/uncloud/internal/machine/api/pb"
	"github.com/psviderski/uncloud/internal/machine/constants"
	"github.com/psviderski/uncloud/internal/machine/network"
	"github.com/psviderski/uncloud/internal/proxy"
	"github.com/psviderski/uncloud/internal/secret"
	"github.com/psviderski/uncloud/pkg/api"
	netproxy "golang.org/x/net/proxy"
)

func (cli *Client) InspectImage(ctx context.Context, id string) ([]api.MachineImage, error) {
	images, err := cli.Docker.InspectImage(ctx, id)
	if dockerclient.IsErrNotFound(err) {
		err = api.ErrNotFound
	}

	return images, err
}

func (cli *Client) InspectRemoteImage(ctx context.Context, id string) ([]api.MachineRemoteImage, error) {
	return cli.Docker.InspectRemoteImage(ctx, id)
}

type PushImageOptions struct {
	// Machines is a list of machine names or IDs to push the image to. If empty, pushes to the machine
	// the client is connected to.
	Machines []string
	// Platform to push for a multi-platform image. Local Docker must use containerd image store
	// to support multi-platform images.
	Platform *ocispec.Platform
}

// PushImage pushes a local Docker image to the specified machines. If no machines are specified,
// it pushes to the machine the client is connected to.
func (cli *Client) PushImage(ctx context.Context, image string, opts PushImageOptions) error {
	dockerCliWrapped, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv,
		dockerclient.WithAPIVersionNegotiation())
	if err != nil {
		return fmt.Errorf("create Docker client: %w", err)
	}
	dockerCli := &docker.Client{Client: dockerCliWrapped}
	defer dockerCli.Close()

	// Check if Docker image exists locally.
	if _, _, err = dockerCli.ImageInspectWithRaw(ctx, image); err != nil {
		if dockerclient.IsErrNotFound(err) {
			return fmt.Errorf("image '%s' not found locally", image)
		}
		return fmt.Errorf("inspect image '%s' locally: %w", image, err)
	}

	// Get the machine info for the specified machines or the connected machine if none are specified.
	var machines []*pb.MachineInfo
	if len(opts.Machines) > 0 {
		machineMembers, err := cli.ListMachines(ctx, &api.MachineFilter{
			NamesOrIDs: opts.Machines,
		})
		if err != nil {
			return fmt.Errorf("list machines: %w", err)
		}

		// Check if all specified machines were found.
		if len(machineMembers) != len(opts.Machines) {
			var notFound []string
			for _, nameOrID := range opts.Machines {
				if machineMembers.FindByNameOrID(nameOrID) == nil {
					notFound = append(notFound, nameOrID)
				}
			}

			return fmt.Errorf("machines not found: %s", strings.Join(notFound, ", "))
		}

		for _, mm := range machineMembers {
			machines = append(machines, mm.Machine)
		}
	} else {
		// No machines specified, use the connected machine.
		m, err := cli.MachineClient.Inspect(ctx, nil)
		if err != nil {
			return fmt.Errorf("inspect connected machine: %w", err)
		}

		// If the machine has been renamed, the new name will only be stored in the cluster store. .Inspect will return
		// the old name from the machine config. So we need to fetch the machine info from the cluster.
		// TODO: make one source of truth for machine info.
		mm, err := cli.InspectMachine(ctx, m.Id)
		if err != nil {
			return fmt.Errorf("inspect machine: %w", err)
		}
		machines = append(machines, mm.Machine)
	}

	// Push image to all specified machines.
	var wg sync.WaitGroup
	errCh := make(chan error, len(machines))

	// TODO: detect the target machine platform and figure out how to handle scenarios when local and target
	//  platforms differ.
	for _, m := range machines {
		wg.Go(func() {
			if err := cli.pushImageToMachine(ctx, dockerCli, image, m, opts.Platform); err != nil {
				errCh <- fmt.Errorf("push image to machine '%s': %w", m.Name, err)
			}
		})
	}

	wg.Wait()
	close(errCh)

	var errs []error
	for err = range errCh {
		errs = append(errs, err)
	}

	return errors.Join(errs...)
}

// pushImageToMachine pushes a local Docker image to a specific machine using local port forwarding to its unregistry.
func (cli *Client) pushImageToMachine(
	ctx context.Context,
	dockerCli *docker.Client,
	imageName string,
	machine *pb.MachineInfo,
	platform *ocispec.Platform,
) error {
	pw := progress.ContextWriter(ctx)

	machineSubnet, _ := machine.Network.Subnet.ToPrefix()
	machineIP := network.MachineIP(machineSubnet)
	unregistryAddr := net.JoinHostPort(machineIP.String(), strconv.Itoa(constants.UnregistryPort))

	dialer, err := cli.connector.Dialer()
	if err != nil {
		return fmt.Errorf("get proxy dialer: %w", err)
	}

	boldStyle := lipgloss.NewStyle().Bold(true)
	proxyEventID := fmt.Sprintf("Proxy to unregistry on %s", boldStyle.Render(machine.Name))
	pw.Event(progress.StartingEvent(proxyEventID))

	// Forward local port 127.0.0.1:PORT to the machine's unregistry over the established client connection.
	unregProxy, err := newUnregistryProxy(ctx, unregistryAddr, dialer, func(err error) {
		pw.Event(progress.NewEvent(proxyEventID, progress.Error, err.Error()))
	})
	if err != nil {
		pw.Event(progress.NewEvent(proxyEventID, progress.Error, err.Error()))
		return fmt.Errorf("create local proxy to unregistry on machine '%s': %w", machine.Name, err)
	}
	// Get the local port the unregistry proxy is listening on.
	proxyPort := unregProxy.Listener.Addr().(*net.TCPAddr).Port

	proxyCtx, cancelProxy := context.WithCancel(ctx)
	proxyCtrID := ""
	pushImageTag := ""

	// Cleanup function to remove temporary resources and stop proxies.
	cleanup := func() {
		// Remove temporary image tag.
		if pushImageTag != "" {
			dockerCli.ImageRemove(ctx, pushImageTag, image.RemoveOptions{})
		}

		// Remove socat proxy container.
		if proxyCtrID != "" {
			dockerCli.ContainerRemove(ctx, proxyCtrID, container.RemoveOptions{Force: true})
		}

		cancelProxy()
	}
	defer cleanup()

	go unregProxy.Run(proxyCtx)

	dockerVirtualised, err := isDockerVirtualised(ctx, dockerCli)
	if err != nil {
		return err
	}

	if dockerVirtualised {
		// Run socat proxy container to forward a localhost port from within the Docker VM to the host machine.
		pw.Event(progress.Event{
			ID:         proxyEventID,
			Status:     progress.Working,
			StatusText: "Starting",
			Text:       "(detected virtualised Docker locally, starting socat container to proxy to Docker VM)",
		})

		proxyCtrID, proxyPort, err = runDockerVMProxyContainer(ctx, dockerCli, proxyPort)
		if err != nil {
			pw.Event(progress.NewEvent(proxyEventID, progress.Error, err.Error()))
			return fmt.Errorf("run socat container to proxy unregistry to Docker VM: %w", err)
		}
	}

	pw.Event(progress.Event{
		ID:         proxyEventID,
		Status:     progress.Done,
		StatusText: "Started",
		Text:       fmt.Sprintf("(localhost:%d → %s)", proxyPort, unregistryAddr),
	})

	// Tag the image for pushing through the proxy.
	pushImageTag = fmt.Sprintf("127.0.0.1:%d/%s", proxyPort, imageName)
	if err = dockerCli.ImageTag(ctx, imageName, pushImageTag); err != nil {
		return fmt.Errorf("tag image for push: %w", err)
	}

	// Push the image through the proxy.
	pushEventID := fmt.Sprintf("Pushing %s to %s", boldStyle.Render(imageName), boldStyle.Render(machine.Name))
	pw.Event(progress.NewEvent(pushEventID, progress.Working, "Pushing"))

	pushCh, err := dockerCli.PushImage(ctx, pushImageTag, image.PushOptions{
		Platform: platform,
	})
	if err != nil {
		pw.Event(progress.NewEvent(pushEventID, progress.Error, err.Error()))
		return fmt.Errorf("push image: %w", err)
	}

	// Wait for push to complete by reading all progress messages and converting them to events.
	// If the context is cancelled, the pushCh will receive a context cancellation error.
	for msg := range pushCh {
		if msg.Err != nil {
			pw.Event(progress.NewEvent(pushEventID, progress.Error, msg.Err.Error()))
			return fmt.Errorf("push image: %w", msg.Err)
		}

		// TODO: support quite mode like in compose: --quiet Push without printing progress information
		if e := toPushProgressEvent(msg.Message); e != nil {
			e.ID = fmt.Sprintf("Layer %s on %s:", e.ID, boldStyle.Render(machine.Name))
			e.ParentID = pushEventID
			pw.Event(*e)
		}
	}
	pw.Event(progress.NewEvent(pushEventID, progress.Done, "Pushed"))

	return nil
}

func newUnregistryProxy(
	ctx context.Context, remoteAddr string, dialer netproxy.ContextDialer, onError func(error),
) (*proxy.Proxy, error) {
	// Test remote connectivity before creating a proxy.
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	testConn, err := dialer.DialContext(ctx, "tcp", remoteAddr)
	if err != nil {
		return nil, fmt.Errorf("connect to remote address '%s': %w", remoteAddr, err)
	}
	testConn.Close()

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, fmt.Errorf("listen on an available port on 127.0.0.1: %w", err)
	}

	p := &proxy.Proxy{
		Listener:    listener,
		RemoteAddr:  remoteAddr,
		DialContext: dialer.DialContext,
		OnError:     onError,
	}

	return p, nil
}

// isDockerVirtualised checks if Docker is running in a virtualised environment like Docker Desktop on macOS.
func isDockerVirtualised(ctx context.Context, dockerCli *docker.Client) (bool, error) {
	info, err := dockerCli.Info(ctx)
	if err != nil {
		return false, fmt.Errorf("get Docker info: %w", err)
	}

	virtualisedHostnames := []string{"docker-desktop", "colima"}
	for _, name := range virtualisedHostnames {
		if strings.Contains(strings.ToLower(info.Name), name) {
			return true, nil
		}
	}

	return false, nil
}

// runDockerVMProxyContainer creates a socat container to proxy an available localhost port within the Docker VM
// (e.g. Docker Desktop on macOS) to the specified target port on the host machine.
// Returns the container ID and the localhost port the container port is bound to.
// TODO: accept custom image name.
func runDockerVMProxyContainer(ctx context.Context, dockerCli *docker.Client, targetPort int) (string, int, error) {
	suffix, err := secret.RandomAlphaNumeric(4)
	if err != nil {
		return "", 0, fmt.Errorf("generate random suffix: %w", err)
	}
	containerName := fmt.Sprintf("uncloud-push-proxy-%s", suffix)

	containerPort := nat.Port("5000/tcp")
	config := &container.Config{
		// TODO: make image configurable.
		Image: "alpine/socat:latest",
		// Reset the default entrypoint "socat".
		Entrypoint: []string{},
		Cmd: []string{
			"timeout", "1800", // Auto-terminate socat after 30 minutes.
			"socat",
			"TCP-LISTEN:5000,fork,reuseaddr",
			fmt.Sprintf("TCP-CONNECT:host.docker.internal:%d", targetPort),
		},
		ExposedPorts: nat.PortSet{
			containerPort: {},
		},
	}

	// Get an available port on localhost to bind the container port to by creating a temporary listener and closing it.
	// We need to explicitly specify the host port and not rely on Docker mapping because if not specified,
	// 'docker push' from Docker Desktop is unable to reach the randomly mapped one for some reason.
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return "", 0, fmt.Errorf("reserve a local port: %w", err)
	}
	hostPort := listener.Addr().(*net.TCPAddr).Port
	listener.Close()

	hostConfig := &container.HostConfig{
		AutoRemove: true,
		PortBindings: nat.PortMap{
			containerPort: []nat.PortBinding{
				{
					HostIP:   "127.0.0.1",
					HostPort: strconv.Itoa(hostPort),
				},
			},
		},
	}

	resp, err := dockerCli.CreateContainerWithImagePull(ctx, containerName, config, hostConfig)
	if err != nil {
		return "", 0, fmt.Errorf("create socat proxy container: %w", err)
	}

	cleanup := func() {
		dockerCli.ContainerRemove(ctx, resp.ID, container.RemoveOptions{Force: true})
	}

	if err = dockerCli.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil {
		// Clean up if start fails.
		cleanup()
		return "", 0, fmt.Errorf("start socat proxy container %s: %w", resp.ID, err)
	}

	return resp.ID, hostPort, nil
}

// toPushProgressEvent converts a JSON progress message from the Docker API to a progress event.
// It's based on toPushProgressEvent from Docker Compose.
func toPushProgressEvent(jm jsonmessage.JSONMessage) *progress.Event {
	if jm.ID == "" || jm.Progress == nil {
		return nil
	}

	status := progress.Working
	percent := 0

	if jm.Progress.Total > 0 {
		percent = int(jm.Progress.Current * 100 / jm.Progress.Total)
	}

	switch jm.Status {
	case "Pushed", "Layer already exists":
		status = progress.Done
		percent = 100
	}

	return &progress.Event{
		ID:         jm.ID,
		Current:    jm.Progress.Current,
		Total:      jm.Progress.Total,
		Percent:    percent,
		Text:       jm.Status,
		Status:     status,
		StatusText: jm.Progress.String(),
	}
}