package-updater/internal/agent/inventory/collector.go

// Package inventory turns what the local Docker daemon runs into the
// services an agent reports to the master.
package inventory

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/api/types/system"

	"github.com/shcizo/package-updater/internal/imageref"
	"github.com/shcizo/package-updater/internal/protocol"
)

const (
	labelProject    = "com.docker.compose.project"
	labelService    = "com.docker.compose.service"
	labelWorkingDir = "com.docker.compose.project.working_dir"
	labelSwarmTask  = "com.docker.swarm.service.id"
)

// DockerClient is the subset of the Docker SDK the collector needs.
// *client.Client satisfies it.
type DockerClient interface {
	ContainerList(ctx context.Context, opts container.ListOptions) ([]container.Summary, error)
	ContainerInspect(ctx context.Context, containerID string) (container.InspectResponse, error)
	ImageList(ctx context.Context, opts image.ListOptions) ([]image.Summary, error)
	ServiceList(ctx context.Context, opts swarm.ServiceListOptions) ([]swarm.Service, error)
	Info(ctx context.Context) (system.Info, error)
}

// Collector reads the local daemon's inventory.
type Collector struct {
	cli    DockerClient
	selfID string
}

// New returns a Collector. selfID is a prefix of the agent's own container
// ID (inside a container, the hostname); that container is marked Self.
func New(cli DockerClient, selfID string) *Collector {
	return &Collector{cli: cli, selfID: selfID}
}

// Hello returns the machine facts sent on connect. AgentVersion is left for
// the caller to fill in.
func (c *Collector) Hello(ctx context.Context) (protocol.Hello, error) {
	info, err := c.cli.Info(ctx)
	if err != nil {
		return protocol.Hello{}, fmt.Errorf("docker info: %w", err)
	}
	h := protocol.Hello{Hostname: info.Name, SwarmManager: info.Swarm.ControlAvailable}
	if info.Swarm.Cluster != nil {
		h.SwarmClusterID = info.Swarm.Cluster.ID
	}
	return h, nil
}

// Snapshot returns every service on this machine, sorted by Key.
func (c *Collector) Snapshot(ctx context.Context) ([]protocol.Service, error) {
	info, err := c.cli.Info(ctx)
	if err != nil {
		return nil, fmt.Errorf("docker info: %w", err)
	}
	containers, err := c.cli.ContainerList(ctx, container.ListOptions{All: true})
	if err != nil {
		return nil, fmt.Errorf("docker container list: %w", err)
	}
	images, err := c.cli.ImageList(ctx, image.ListOptions{})
	if err != nil {
		return nil, fmt.Errorf("docker image list: %w", err)
	}
	repoDigests := make(map[string][]string, len(images))
	for _, img := range images {
		repoDigests[img.ID] = img.RepoDigests
	}

	byKey := make(map[string]protocol.Service)
	for _, ct := range containers {
		if _, isTask := ct.Labels[labelSwarmTask]; isTask {
			continue
		}
		ref, err := c.imageRef(ctx, ct)
		if err != nil {
			return nil, err
		}
		svc := c.fromContainer(ct, ref, repoDigests[ct.ImageID])
		if prev, ok := byKey[svc.Key()]; ok {
			svc = mergeReplicas(prev, svc)
		}
		byKey[svc.Key()] = svc
	}

	if info.Swarm.ControlAvailable {
		services, err := c.cli.ServiceList(ctx, swarm.ServiceListOptions{Status: true})
		if err != nil {
			return nil, fmt.Errorf("docker service list: %w", err)
		}
		for _, s := range services {
			if s.Spec.TaskTemplate.ContainerSpec == nil {
				continue
			}
			svc := fromSwarmService(s)
			byKey[svc.Key()] = svc
		}
	}

	out := make([]protocol.Service, 0, len(byKey))
	for _, s := range byKey {
		out = append(out, s)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Key() < out[j].Key() })
	return out, nil
}

// imageRef returns the image reference the container was created from.
// When the tag has since moved to another image, ContainerList reports the
// image ID ("sha256:...") instead, so the reference is read from the
// container's config.
func (c *Collector) imageRef(ctx context.Context, ct container.Summary) (string, error) {
	if !strings.HasPrefix(ct.Image, "sha256:") {
		return ct.Image, nil
	}
	insp, err := c.cli.ContainerInspect(ctx, ct.ID)
	if err != nil {
		return "", fmt.Errorf("docker container inspect %s: %w", ct.ID, err)
	}
	if insp.Config == nil {
		return ct.Image, nil
	}
	return insp.Config.Image, nil
}

func (c *Collector) fromContainer(ct container.Summary, ref string, repoDigests []string) protocol.Service {
	svc := protocol.Service{
		Image:  ref,
		Digest: digestFor(ref, repoDigests),
		Status: containerStatus(ct.State),
		Self:   c.selfID != "" && strings.HasPrefix(ct.ID, c.selfID),
	}
	project, service := ct.Labels[labelProject], ct.Labels[labelService]
	if project != "" && service != "" {
		svc.Kind = protocol.KindCompose
		svc.ComposeProject = project
		svc.ComposeService = service
		svc.WorkingDir = ct.Labels[labelWorkingDir]
		return svc
	}
	svc.Kind = protocol.KindContainer
	svc.ContainerName = ct.ID
	if len(ct.Names) > 0 {
		svc.ContainerName = strings.TrimPrefix(ct.Names[0], "/")
	}
	return svc
}

func fromSwarmService(s swarm.Service) protocol.Service {
	ref := s.Spec.TaskTemplate.ContainerSpec.Image
	return protocol.Service{
		Kind:         protocol.KindSwarm,
		SwarmService: s.Spec.Name,
		Image:        ref,
		Digest:       imageref.Parse(ref).Digest,
		Status:       swarmStatus(s.ServiceStatus),
	}
}

// mergeReplicas folds a second container of the same Compose service into
// the first: the service counts as running if any replica runs.
func mergeReplicas(a, b protocol.Service) protocol.Service {
	if b.Status == "running" {
		a.Status = "running"
	}
	a.Self = a.Self || b.Self
	return a
}

// digestFor returns the registry digest the container's image was pulled
// as. A reference that already pins a digest wins; otherwise the image's
// RepoDigests entry for the same repository is used.
func digestFor(ref string, repoDigests []string) string {
	r := imageref.Parse(ref)
	if r.Digest != "" {
		return r.Digest
	}
	for _, rd := range repoDigests {
		if p := imageref.Parse(rd); p.Repo == r.Repo {
			return p.Digest
		}
	}
	return ""
}

func containerStatus(state string) string {
	if state == "running" {
		return "running"
	}
	return "stopped"
}

func swarmStatus(st *swarm.ServiceStatus) string {
	switch {
	case st == nil:
		return "unknown"
	case st.DesiredTasks == 0:
		return "stopped"
	case st.RunningTasks >= st.DesiredTasks:
		return "running"
	default:
		return "degraded"
	}
}