feat(provider): pull command and get-image protocol for image distribution

Providers opt into image distribution by declaring a pull block in their metadata — the presence of the command block is the declaration of support, generalizing the stop precedent. During its pull command a provider can request the service image from the local daemon with a get-image message (optionally platform-narrowed); compose answers on stdin with one image-stream JSON line then the tar as HTTP/1.1 chunked data (RFC 9112): length-prefixed blocks need no in-band delimiter in binary data, the zero-length chunk marks a COMPLETE transfer, and any stock chunked reader consumes it. There is deliberately no trailer nor final CRLF after the zero chunk — a stock reader stops there without consuming further bytes, so the next stdin answer starts clean for a provider requesting several images. An export failure is announced in-band through the error field so the provider is never left waiting; a failure after the announce closes the answer channel so the truncation is observable as EOF mid-chunk instead of a stream nobody will finish. The stream is exclusive on stdin for its whole duration — the framing invariant behind the lock. get-image is scoped to the pull command: emitted during up, down or stop it is a protocol error rather than served, so an image export can never stall another lifecycle command. Message dispatch moves out of executePlugin into handlePluginMessage, which also keeps the function under the complexity threshold. Signed-off-by: Nicolas De Loof <[email protected]>
2026-10-12 02:13:08 +00:00 · 2026-10-02 17:35:20 +02:00
parent 06f0a0c913
commit 5bc5b864fc
2 changed files with 158 additions and 7 deletions
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/http/httputil"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -33,8 +34,10 @@ import (

 	"github.com/compose-spec/compose-go/v2/types"
 	"github.com/containerd/errdefs"
+	"github.com/containerd/platforms"
 	"github.com/docker/cli/cli-plugins/manager"
 	"github.com/docker/cli/cli/config"
+	"github.com/moby/moby/client"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"

@@ -44,6 +47,9 @@ import (
 type JsonMessage struct {
 	Type    string `json:"type"`
 	Message string `json:"message,omitempty"`
+	// Platform optionally narrows a get-image request to one platform of a
+	// multi-platform image (e.g. "linux/arm64").
+	Platform string `json:"platform,omitempty"`
 }

 const (
@@ -70,6 +76,17 @@ const (
 	// resource has no use for it.
 	GetRelayInfoType = "get-relay-info"

+	// GetImageType is a message the provider sends during its pull command
+	// to receive the service image from the local daemon: compose answers on
+	// the provider's stdin with one ImageStreamType JSON line, then — unless
+	// that line carries an error — the image tar encoded as an HTTP/1.1
+	// chunked body (see streamImageTo).
+	GetImageType = "get-image"
+
+	// ImageStreamType is the type of the JSON line answering a get-image
+	// request.
+	ImageStreamType = "image-stream"
+
 	// ComposeProviderMessagesEnv announces to the provider process, as a
 	// comma-separated list, every message type this compose accepts on the
 	// provider's stdout — so a provider can adapt to the compose it runs
@@ -91,8 +108,19 @@ var providerMessageTypes = strings.Join([]string{
 	PublishEndpointType,
 	GetServiceConfigType,
 	GetRelayInfoType,
+	GetImageType,
 }, ",")

+// imageStreamAnswer is the stdin answer to a get-image request. On success
+// Encoding and MediaType describe the byte stream that follows the JSON line;
+// on failure Error carries the reason and no stream follows.
+type imageStreamAnswer struct {
+	Type      string `json:"type"`
+	Encoding  string `json:"encoding,omitempty"`
+	MediaType string `json:"media-type,omitempty"`
+	Error     string `json:"error,omitempty"`
+}
+
 type pluginVariables struct {
 	prefixed types.Mapping
 	raw      types.Mapping
@@ -105,7 +133,7 @@ type pluginVariables struct {

 var mux sync.Mutex

-func (s *composeService) runPlugin(ctx context.Context, project *types.Project, service types.ServiceConfig, command string) error {
+func (s *composeService) runPlugin(ctx context.Context, project *types.Project, service types.ServiceConfig, command string, extraArgs ...string) error {
 	provider := *service.Provider

 	plugin, err := s.getPluginBinaryPath(provider.Type)
@@ -113,7 +141,7 @@ func (s *composeService) runPlugin(ctx context.Context, project *types.Project,
 		return err
 	}

-	cmd, err := s.setupPluginCommand(ctx, project, service, plugin, command)
+	cmd, err := s.setupPluginCommand(ctx, project, service, plugin, command, extraArgs...)
 	if err != nil {
 		return err
 	}
@@ -126,7 +154,7 @@ func (s *composeService) runPlugin(ctx context.Context, project *types.Project,
 		return err
 	}

-	if command == "stop" {
+	if command == "stop" || command == "pull" {
 		return nil
 	}

@@ -203,6 +231,9 @@ func (s *composeService) executePlugin(ctx context.Context, project *types.Proje
 	case "stop":
 		s.events.On(stoppingEvent(service.Name))
 		action = "stop"
+	case "pull":
+		s.events.On(newEvent(service.Name, api.Working, "Pulling"))
+		action = "pull"
 	default:
 		return pluginVariables{}, fmt.Errorf("unsupported plugin command: %s", command)
 	}
@@ -275,7 +306,7 @@ func (s *composeService) executePlugin(ctx context.Context, project *types.Proje
 		if err != nil {
 			return pluginVariables{}, err
 		}
-		if err := s.handlePluginMessage(ctx, project, service, msg, &variables, stdin, &stdinMu, &answers); err != nil {
+		if err := s.handlePluginMessage(ctx, project, service, msg, command, &variables, stdin, &stdinMu, &answers); err != nil {
 			return pluginVariables{}, err
 		}
 	}
@@ -293,17 +324,77 @@ func (s *composeService) executePlugin(ctx context.Context, project *types.Proje
 		s.events.On(removedEvent(service.Name))
 	case "stop":
 		s.events.On(stoppedEvent(service.Name))
+	case "pull":
+		s.events.On(newEvent(service.Name, api.Done, "Pulled"))
 	}
 	return variables, nil
 }

+// imageStreamChunkSize is the fixed buffer compose fills from the image
+// export before flushing it as one chunk of the stream.
+const imageStreamChunkSize = 1024 * 1024
+
+// streamImageTo answers a get-image request on the provider's stdin: one
+// ImageStreamType JSON line announcing the transfer, then the image tar
+// encoded as HTTP/1.1 chunked data (RFC 9112 §7.1) — every block of data
+// prefixed by its length, the zero-length chunk marking a COMPLETE
+// transfer. Length-prefixed framing needs no in-band delimiter (any
+// byte value can appear inside a tar) and is consumable with any language's
+// chunked-body reader. A stream that ends without the terminating zero chunk
+// was aborted: the provider must discard it.
+// When the image cannot be exported, the announce line carries an error
+// instead and no stream follows.
+func (s *composeService) streamImageTo(ctx context.Context, w io.Writer, ref, platform string) error {
+	announce := func(a imageStreamAnswer) error {
+		payload, err := json.Marshal(a)
+		if err != nil {
+			return err
+		}
+		_, err = w.Write(append(payload, '\n'))
+		return err
+	}
+
+	var opts []client.ImageSaveOption
+	if platform != "" {
+		p, err := platforms.Parse(platform)
+		if err != nil {
+			_ = announce(imageStreamAnswer{Type: ImageStreamType, Error: fmt.Sprintf("invalid platform %q: %s", platform, err)})
+			return err
+		}
+		opts = append(opts, client.ImageSaveWithPlatforms(p))
+	}
+	tar, err := s.apiClient().ImageSave(ctx, []string{ref}, opts...)
+	if err != nil {
+		_ = announce(imageStreamAnswer{Type: ImageStreamType, Error: err.Error()})
+		return err
+	}
+	defer func() { _ = tar.Close() }()
+
+	if err := announce(imageStreamAnswer{Type: ImageStreamType, Encoding: "chunked", MediaType: "application/x-tar"}); err != nil {
+		return err
+	}
+	cw := httputil.NewChunkedWriter(w)
+	if _, err := io.CopyBuffer(cw, struct{ io.Reader }{tar}, make([]byte, imageStreamChunkSize)); err != nil {
+		return err
+	}
+	// ChunkedWriter.Close writes the zero-length chunk, which ends the
+	// stream. Unlike an HTTP message there is deliberately no trailer
+	// section nor final CRLF: a stock chunked reader stops at the zero
+	// chunk without consuming either, and leftover bytes would corrupt the
+	// next JSON answer a provider requesting several images reads.
+	return cw.Close()
+}
+
 // handlePluginMessage processes one provider message, mutating variables in
 // place; a returned error is terminal for the run (the provider is killed
 // and the command fails). An unknown message type IS such an error: a
 // provider requiring a message this compose does not support must fail
 // loudly, not degrade silently — providers adapt through the
-// COMPOSE_PROVIDER_MESSAGES announcement instead.
-func (s *composeService) handlePluginMessage(ctx context.Context, project *types.Project, service types.ServiceConfig, msg JsonMessage,
+// COMPOSE_PROVIDER_MESSAGES announcement instead. Answer-bearing requests
+// (get-service-config, get-relay-info, get-image) are served from their own
+// goroutine — see the answers/stdinMu contract in executePlugin.
+func (s *composeService) handlePluginMessage(
+	ctx context.Context, project *types.Project, service types.ServiceConfig, msg JsonMessage, command string,
 	variables *pluginVariables, stdin io.WriteCloser, stdinMu *sync.Mutex, answers *sync.WaitGroup,
 ) error {
 	switch msg.Type {
@@ -339,6 +430,36 @@ func (s *composeService) handlePluginMessage(ctx context.Context, project *types
 			return fmt.Errorf("failed to answer get-relay-info: %w", err)
 		}
 		answerProvider(stdin, stdinMu, answers, payload)
+	case GetImageType:
+		// image distribution belongs to the pull command: answering it
+		// elsewhere would let an image export stall a down or a stop
+		if command != "pull" {
+			return fmt.Errorf("invalid message from plugin: %s is only supported during the pull command", GetImageType)
+		}
+		ref, platform := msg.Message, msg.Platform
+		answers.Add(1)
+		go func() {
+			defer answers.Done()
+			// stdinMu is deliberately held for the whole transfer: the
+			// chunked body must be contiguous on stdin, any concurrent
+			// answer interleaved into it would corrupt the framing. A
+			// provider must therefore drain the announced stream before
+			// expecting any other answer (documented contract); compose
+			// cannot hang forever on a provider that stops reading — the
+			// error path kills the process, which EPIPEs the write.
+			stdinMu.Lock()
+			defer stdinMu.Unlock()
+			if err := s.streamImageTo(ctx, stdin, ref, platform); err != nil {
+				logrus.Warnf("provider %q: get-image %q: %v", service.Name, ref, err)
+				// A failure after the success announce leaves the stream
+				// without its terminating chunk, and the channel cannot be
+				// resynchronized (any byte would read as chunk data).
+				// Closing stdin makes the truncation observable — the
+				// provider gets EOF mid-chunk and discards, instead of
+				// blocking forever on a stream nobody will finish.
+				_ = stdin.Close()
+			}
+		}()
 	case PublishEndpointType:
 		port, upstream, err := parseEndpointMessage(msg.Message)
 		if err != nil {
@@ -388,7 +509,7 @@ func (s *composeService) getPluginBinaryPath(provider string) (path string, err
 	return path, err
 }

-func (s *composeService) setupPluginCommand(ctx context.Context, project *types.Project, service types.ServiceConfig, path, command string) (*exec.Cmd, error) {
+func (s *composeService) setupPluginCommand(ctx context.Context, project *types.Project, service types.ServiceConfig, path, command string, extraArgs ...string) (*exec.Cmd, error) {
 	cmdOptionsMetadata := s.getPluginMetadata(path, service.Provider.Type, project)
 	var currentCommandMetadata CommandMetadata
 	switch command {
@@ -401,6 +522,13 @@ func (s *composeService) setupPluginCommand(ctx context.Context, project *types.
 			return nil, nil
 		}
 		currentCommandMetadata = *cmdOptionsMetadata.Stop
+	case "pull":
+		// image distribution is opt-in, declared like stop by the presence
+		// of the command block in the provider metadata
+		if cmdOptionsMetadata.Pull == nil {
+			return nil, nil
+		}
+		currentCommandMetadata = *cmdOptionsMetadata.Pull
 	}

 	provider := *service.Provider
@@ -417,6 +545,7 @@ func (s *composeService) setupPluginCommand(ctx context.Context, project *types.
 			}
 		}
 	}
+	args = append(args, extraArgs...)
 	args = append(args, service.Name)

 	cmd := exec.CommandContext(ctx, path, args...)
@@ -469,6 +598,9 @@ type ProviderMetadata struct {
 	Up          CommandMetadata  `json:"up"`
 	Down        CommandMetadata  `json:"down"`
 	Stop        *CommandMetadata `json:"stop,omitempty"`
+	// Pull declares support for the image-distribution command; like Stop,
+	// the presence of the block is what opts the provider in.
+	Pull *CommandMetadata `json:"pull,omitempty"`
 }

 func (p ProviderMetadata) IsEmpty() bool {
@@ -97,6 +97,25 @@ func TestProviderMetadata_StopAbsent(t *testing.T) {
 	assert.Assert(t, metadata.Stop == nil, "Stop should be nil when absent from JSON")
 }

+func TestProviderMetadata_PullAbsent(t *testing.T) {
+	raw := `{"description":"x","up":{"parameters":[]},"down":{"parameters":[]}}`
+
+	var metadata ProviderMetadata
+	err := json.Unmarshal([]byte(raw), &metadata)
+	assert.NilError(t, err)
+	assert.Assert(t, metadata.Pull == nil, "Pull should be nil when absent from JSON")
+}
+
+func TestProviderMetadata_PullPresent(t *testing.T) {
+	raw := `{"pull":{"parameters":[{"name":"registry"}]}}`
+
+	var metadata ProviderMetadata
+	err := json.Unmarshal([]byte(raw), &metadata)
+	assert.NilError(t, err)
+	assert.Assert(t, metadata.Pull != nil, "Pull should be non-nil when key present")
+	assert.Equal(t, metadata.Pull.Parameters[0].Name, "registry")
+}
+
 func TestProviderMetadata_StopAdvertisedWithoutParameters(t *testing.T) {
 	raw := `{"stop":{"parameters":null}}`