Enable to run build and invoke in background

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>

commands/controller/client.go

@@ -0,0 +1,262 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/containerd/console"
+	"github.com/containerd/containerd/defaults"
+	"github.com/containerd/containerd/pkg/dialer"
+	"github.com/docker/buildx/commands/controller/pb"
+	"github.com/docker/buildx/util/progress"
+	"github.com/moby/buildkit/client"
+	"github.com/moby/buildkit/identity"
+	"github.com/pkg/errors"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+func NewClient(addr string) (*Client, error) {
+	backoffConfig := backoff.DefaultConfig
+	backoffConfig.MaxDelay = 3 * time.Second
+	connParams := grpc.ConnectParams{
+		Backoff: backoffConfig,
+	}
+	gopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithConnectParams(connParams),
+		grpc.WithContextDialer(dialer.ContextDialer),
+		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(defaults.DefaultMaxRecvMsgSize)),
+		grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(defaults.DefaultMaxSendMsgSize)),
+	}
+	conn, err := grpc.Dial(dialer.DialAddress(addr), gopts...)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{conn: conn}, nil
+}
+
+type Client struct {
+	conn      *grpc.ClientConn
+	closeOnce sync.Once
+}
+
+func (c *Client) Close() (err error) {
+	c.closeOnce.Do(func() {
+		err = c.conn.Close()
+	})
+	return
+}
+
+func (c *Client) Version(ctx context.Context) (string, string, string, error) {
+	res, err := c.client().Info(ctx, &pb.InfoRequest{})
+	if err != nil {
+		return "", "", "", err
+	}
+	v := res.BuildxVersion
+	return v.Package, v.Version, v.Revision, nil
+}
+
+func (c *Client) List(ctx context.Context) (keys []string, retErr error) {
+	res, err := c.client().List(ctx, &pb.ListRequest{})
+	if err != nil {
+		return nil, err
+	}
+	return res.Keys, nil
+}
+
+func (c *Client) Disconnect(ctx context.Context, key string) error {
+	_, err := c.client().Disconnect(ctx, &pb.DisconnectRequest{Ref: key})
+	return err
+}
+
+func (c *Client) Invoke(ctx context.Context, ref string, containerConfig pb.ContainerConfig, in io.ReadCloser, stdout io.WriteCloser, stderr io.WriteCloser) error {
+	if ref == "" {
+		return fmt.Errorf("build reference must be specified")
+	}
+	stream, err := c.client().Invoke(ctx)
+	if err != nil {
+		return err
+	}
+	return attachIO(ctx, stream, &pb.InitMessage{Ref: ref, ContainerConfig: &containerConfig}, ioAttachConfig{
+		stdin:  in,
+		stdout: stdout,
+		stderr: stderr,
+		// TODO: Signal, Resize
+	})
+}
+
+func (c *Client) Build(ctx context.Context, options pb.BuildOptions, in io.ReadCloser, w io.Writer, out console.File, progressMode string) (string, error) {
+	ref := identity.NewID()
+	pw, err := progress.NewPrinter(context.TODO(), w, out, progressMode)
+	if err != nil {
+		return "", err
+	}
+	statusChan := make(chan *client.SolveStatus)
+	statusDone := make(chan struct{})
+	eg, egCtx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		defer close(statusChan)
+		return c.build(egCtx, ref, options, in, statusChan)
+	})
+	eg.Go(func() error {
+		defer close(statusDone)
+		for s := range statusChan {
+			st := s
+			pw.Write(st)
+		}
+		return nil
+	})
+	eg.Go(func() error {
+		<-statusDone
+		return pw.Wait()
+	})
+	return ref, eg.Wait()
+}
+
+func (c *Client) build(ctx context.Context, ref string, options pb.BuildOptions, in io.ReadCloser, statusChan chan *client.SolveStatus) error {
+	eg, egCtx := errgroup.WithContext(ctx)
+	done := make(chan struct{})
+	eg.Go(func() error {
+		defer close(done)
+		if _, err := c.client().Build(egCtx, &pb.BuildRequest{
+			Ref:     ref,
+			Options: &options,
+		}); err != nil {
+			return err
+		}
+		return nil
+	})
+	eg.Go(func() error {
+		stream, err := c.client().Status(egCtx, &pb.StatusRequest{
+			Ref: ref,
+		})
+		if err != nil {
+			return err
+		}
+		for {
+			resp, err := stream.Recv()
+			if err != nil {
+				if err == io.EOF {
+					return nil
+				}
+				return errors.Wrap(err, "failed to receive status")
+			}
+			s := client.SolveStatus{}
+			for _, v := range resp.Vertexes {
+				s.Vertexes = append(s.Vertexes, &client.Vertex{
+					Digest:        v.Digest,
+					Inputs:        v.Inputs,
+					Name:          v.Name,
+					Started:       v.Started,
+					Completed:     v.Completed,
+					Error:         v.Error,
+					Cached:        v.Cached,
+					ProgressGroup: v.ProgressGroup,
+				})
+			}
+			for _, v := range resp.Statuses {
+				s.Statuses = append(s.Statuses, &client.VertexStatus{
+					ID:        v.ID,
+					Vertex:    v.Vertex,
+					Name:      v.Name,
+					Total:     v.Total,
+					Current:   v.Current,
+					Timestamp: v.Timestamp,
+					Started:   v.Started,
+					Completed: v.Completed,
+				})
+			}
+			for _, v := range resp.Logs {
+				s.Logs = append(s.Logs, &client.VertexLog{
+					Vertex:    v.Vertex,
+					Stream:    int(v.Stream),
+					Data:      v.Msg,
+					Timestamp: v.Timestamp,
+				})
+			}
+			for _, v := range resp.Warnings {
+				s.Warnings = append(s.Warnings, &client.VertexWarning{
+					Vertex:     v.Vertex,
+					Level:      int(v.Level),
+					Short:      v.Short,
+					Detail:     v.Detail,
+					URL:        v.Url,
+					SourceInfo: v.Info,
+					Range:      v.Ranges,
+				})
+			}
+			statusChan <- &s
+		}
+	})
+	if in != nil {
+		eg.Go(func() error {
+			stream, err := c.client().Input(egCtx)
+			if err != nil {
+				return err
+			}
+			if err := stream.Send(&pb.InputMessage{
+				Input: &pb.InputMessage_Init{
+					Init: &pb.InputInitMessage{
+						Ref: ref,
+					},
+				},
+			}); err != nil {
+				return fmt.Errorf("failed to init input: %w", err)
+			}
+
+			inReader, inWriter := io.Pipe()
+			eg2, _ := errgroup.WithContext(ctx)
+			eg2.Go(func() error {
+				<-done
+				return inWriter.Close()
+			})
+			go func() {
+				// do not wait for read completion but return here and let the caller send EOF
+				// this allows us to return on ctx.Done() without being blocked by this reader.
+				io.Copy(inWriter, in)
+				inWriter.Close()
+			}()
+			eg2.Go(func() error {
+				for {
+					buf := make([]byte, 32*1024)
+					n, err := inReader.Read(buf)
+					if err != nil {
+						if err == io.EOF {
+							break // break loop and send EOF
+						}
+						return err
+					} else if n > 0 {
+						if stream.Send(&pb.InputMessage{
+							Input: &pb.InputMessage_Data{
+								Data: &pb.DataMessage{
+									Data: buf[:n],
+								},
+							},
+						}); err != nil {
+							return err
+						}
+					}
+				}
+				return stream.Send(&pb.InputMessage{
+					Input: &pb.InputMessage_Data{
+						Data: &pb.DataMessage{
+							EOF: true,
+						},
+					},
+				})
+			})
+			return eg2.Wait()
+		})
+	}
+	return eg.Wait()
+}
+
+func (c *Client) client() pb.ControllerClient {
+	return pb.NewControllerClient(c.conn)
+}

commands/controller/controller.go

@@ -0,0 +1,440 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/docker/buildx/build"
+	"github.com/docker/buildx/commands/controller/pb"
+	"github.com/docker/buildx/util/ioset"
+	"github.com/docker/buildx/version"
+	controlapi "github.com/moby/buildkit/api/services/control"
+	"github.com/moby/buildkit/client"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
+)
+
+type BuildFunc func(ctx context.Context, options *pb.BuildOptions, stdin io.Reader, statusChan chan *client.SolveStatus) (res *build.ResultContext, err error)
+
+func New(buildFunc BuildFunc) *Controller {
+	return &Controller{
+		buildFunc: buildFunc,
+	}
+}
+
+type Controller struct {
+	buildFunc BuildFunc
+	session   map[string]session
+	sessionMu sync.Mutex
+}
+
+type session struct {
+	statusChan      chan *client.SolveStatus
+	result          *build.ResultContext
+	inputPipe       *io.PipeWriter
+	curInvokeCancel func()
+	curBuildCancel  func()
+}
+
+func (m *Controller) Info(ctx context.Context, req *pb.InfoRequest) (res *pb.InfoResponse, err error) {
+	return &pb.InfoResponse{
+		BuildxVersion: &pb.BuildxVersion{
+			Package:  version.Package,
+			Version:  version.Version,
+			Revision: version.Revision,
+		},
+	}, nil
+}
+
+func (m *Controller) List(ctx context.Context, req *pb.ListRequest) (res *pb.ListResponse, err error) {
+	keys := make(map[string]struct{})
+
+	m.sessionMu.Lock()
+	for k := range m.session {
+		keys[k] = struct{}{}
+	}
+	m.sessionMu.Unlock()
+
+	var keysL []string
+	for k := range keys {
+		keysL = append(keysL, k)
+	}
+	return &pb.ListResponse{
+		Keys: keysL,
+	}, nil
+}
+
+func (m *Controller) Disconnect(ctx context.Context, req *pb.DisconnectRequest) (res *pb.DisconnectResponse, err error) {
+	key := req.Ref
+	if key == "" {
+		return nil, fmt.Errorf("disconnect: empty key")
+	}
+
+	m.sessionMu.Lock()
+	if s, ok := m.session[key]; ok {
+		if s.curBuildCancel != nil {
+			s.curBuildCancel()
+		}
+		if s.curInvokeCancel != nil {
+			s.curInvokeCancel()
+		}
+	}
+	delete(m.session, key)
+	m.sessionMu.Unlock()
+
+	return &pb.DisconnectResponse{}, nil
+}
+
+func (m *Controller) Close() error {
+	m.sessionMu.Lock()
+	for k := range m.session {
+		if s, ok := m.session[k]; ok {
+			if s.curBuildCancel != nil {
+				s.curBuildCancel()
+			}
+			if s.curInvokeCancel != nil {
+				s.curInvokeCancel()
+			}
+		}
+	}
+	m.sessionMu.Unlock()
+	return nil
+}
+
+func (m *Controller) Build(ctx context.Context, req *pb.BuildRequest) (*pb.BuildResponse, error) {
+	ref := req.Ref
+	if ref == "" {
+		return nil, fmt.Errorf("build: empty key")
+	}
+
+	// Prepare status channel and session if not exists
+	m.sessionMu.Lock()
+	if m.session == nil {
+		m.session = make(map[string]session)
+	}
+	s, ok := m.session[ref]
+	if ok && m.session[ref].statusChan != nil {
+		m.sessionMu.Unlock()
+		return &pb.BuildResponse{}, fmt.Errorf("build or status ongoing or status didn't called")
+	}
+	statusChan := make(chan *client.SolveStatus)
+	s.statusChan = statusChan
+	m.session[ref] = session{statusChan: statusChan}
+	m.sessionMu.Unlock()
+	defer func() {
+		close(statusChan)
+		m.sessionMu.Lock()
+		s, ok := m.session[ref]
+		if ok {
+			s.statusChan = nil
+		}
+		m.sessionMu.Unlock()
+	}()
+
+	// Prepare input stream pipe
+	inR, inW := io.Pipe()
+	m.sessionMu.Lock()
+	if s, ok := m.session[ref]; ok {
+		s.inputPipe = inW
+		m.session[ref] = s
+	} else {
+		m.sessionMu.Unlock()
+		return nil, fmt.Errorf("build: unknown key %v", ref)
+	}
+	m.sessionMu.Unlock()
+	defer inR.Close()
+
+	// Build the specified request
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	res, err := m.buildFunc(ctx, req.Options, inR, statusChan)
+	m.sessionMu.Lock()
+	if s, ok := m.session[ref]; ok {
+		s.result = res
+		s.curBuildCancel = cancel
+		m.session[ref] = s
+	} else {
+		m.sessionMu.Unlock()
+		return nil, fmt.Errorf("build: unknown key %v", ref)
+	}
+	m.sessionMu.Unlock()
+
+	return &pb.BuildResponse{}, err
+}
+
+func (m *Controller) Status(req *pb.StatusRequest, stream pb.Controller_StatusServer) error {
+	ref := req.Ref
+	if ref == "" {
+		return fmt.Errorf("status: empty key")
+	}
+
+	// Wait and get status channel prepared by Build()
+	var statusChan <-chan *client.SolveStatus
+	for {
+		// TODO: timeout?
+		m.sessionMu.Lock()
+		if _, ok := m.session[ref]; !ok || m.session[ref].statusChan == nil {
+			m.sessionMu.Unlock()
+			time.Sleep(time.Millisecond) // TODO: wait Build without busy loop and make it cancellable
+			continue
+		}
+		statusChan = m.session[ref].statusChan
+		m.sessionMu.Unlock()
+		break
+	}
+
+	// forward status
+	for ss := range statusChan {
+		if ss == nil {
+			break
+		}
+		cs := toControlStatus(ss)
+		if err := stream.Send(cs); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (m *Controller) Input(stream pb.Controller_InputServer) (err error) {
+	// Get the target ref from init message
+	msg, err := stream.Recv()
+	if err != nil {
+		if !errors.Is(err, io.EOF) {
+			return err
+		}
+		return nil
+	}
+	init := msg.GetInit()
+	if init == nil {
+		return fmt.Errorf("unexpected message: %T; wanted init", msg.GetInit())
+	}
+	ref := init.Ref
+	if ref == "" {
+		return fmt.Errorf("input: no ref is provided")
+	}
+
+	// Wait and get input stream pipe prepared by Build()
+	var inputPipeW *io.PipeWriter
+	for {
+		// TODO: timeout?
+		m.sessionMu.Lock()
+		if _, ok := m.session[ref]; !ok || m.session[ref].inputPipe == nil {
+			m.sessionMu.Unlock()
+			time.Sleep(time.Millisecond) // TODO: wait Build without busy loop and make it cancellable
+			continue
+		}
+		inputPipeW = m.session[ref].inputPipe
+		m.sessionMu.Unlock()
+		break
+	}
+
+	// Forward input stream
+	eg, ctx := errgroup.WithContext(context.TODO())
+	done := make(chan struct{})
+	msgCh := make(chan *pb.InputMessage)
+	eg.Go(func() error {
+		defer close(msgCh)
+		for {
+			msg, err := stream.Recv()
+			if err != nil {
+				if !errors.Is(err, io.EOF) {
+					return err
+				}
+				return nil
+			}
+			select {
+			case msgCh <- msg:
+			case <-done:
+				return nil
+			case <-ctx.Done():
+				return nil
+			}
+		}
+	})
+	eg.Go(func() (retErr error) {
+		defer close(done)
+		defer func() {
+			if retErr != nil {
+				inputPipeW.CloseWithError(retErr)
+				return
+			}
+			inputPipeW.Close()
+		}()
+		for {
+			var msg *pb.InputMessage
+			select {
+			case msg = <-msgCh:
+			case <-ctx.Done():
+				return fmt.Errorf("canceled: %w", ctx.Err())
+			}
+			if msg == nil {
+				return nil
+			}
+			if data := msg.GetData(); data != nil {
+				if len(data.Data) > 0 {
+					_, err := inputPipeW.Write(data.Data)
+					if err != nil {
+						return err
+					}
+				}
+				if data.EOF {
+					return nil
+				}
+			}
+		}
+	})
+
+	return eg.Wait()
+}
+
+func (m *Controller) Invoke(srv pb.Controller_InvokeServer) error {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+	containerIn, containerOut := ioset.Pipe()
+	waitInvokeDoneCh := make(chan struct{})
+	var cancelOnce sync.Once
+	curInvokeCancel := func() {
+		cancelOnce.Do(func() { containerOut.Close(); containerIn.Close(); cancel() })
+		<-waitInvokeDoneCh
+	}
+	defer curInvokeCancel()
+
+	var cfg *pb.ContainerConfig
+	var resultCtx *build.ResultContext
+	initDoneCh := make(chan struct{})
+	initErrCh := make(chan error)
+	eg, egCtx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		return serveIO(egCtx, srv, func(initMessage *pb.InitMessage) (retErr error) {
+			defer func() {
+				if retErr != nil {
+					initErrCh <- retErr
+				}
+				close(initDoneCh)
+			}()
+			ref := initMessage.Ref
+			cfg = initMessage.ContainerConfig
+
+			// Register cancel callback
+			m.sessionMu.Lock()
+			if s, ok := m.session[ref]; ok {
+				if cancel := s.curInvokeCancel; cancel != nil {
+					logrus.Warnf("invoke: cancelling ongoing invoke of %q", ref)
+					cancel()
+				}
+				s.curInvokeCancel = curInvokeCancel
+				m.session[ref] = s
+			} else {
+				m.sessionMu.Unlock()
+				return fmt.Errorf("invoke: unknown key %v", ref)
+			}
+			m.sessionMu.Unlock()
+
+			// Get the target result to invoke a container from
+			m.sessionMu.Lock()
+			if _, ok := m.session[ref]; !ok || m.session[ref].result == nil {
+				m.sessionMu.Unlock()
+				return fmt.Errorf("unknown reference: %q", ref)
+			}
+			resultCtx = m.session[ref].result
+			m.sessionMu.Unlock()
+			return nil
+		}, &ioServerConfig{
+			stdin:  containerOut.Stdin,
+			stdout: containerOut.Stdout,
+			stderr: containerOut.Stderr,
+			// TODO: signal, resize
+		})
+	})
+	eg.Go(func() error {
+		defer containerIn.Close()
+		defer cancel()
+		select {
+		case <-initDoneCh:
+		case err := <-initErrCh:
+			return err
+		}
+		if cfg == nil {
+			return fmt.Errorf("no container config is provided")
+		}
+		if resultCtx == nil {
+			return fmt.Errorf("no result is provided")
+		}
+		ccfg := build.ContainerConfig{
+			ResultCtx:  resultCtx,
+			Entrypoint: cfg.Entrypoint,
+			Cmd:        cfg.Cmd,
+			Env:        cfg.Env,
+			Tty:        cfg.Tty,
+			Stdin:      containerIn.Stdin,
+			Stdout:     containerIn.Stdout,
+			Stderr:     containerIn.Stderr,
+		}
+		if !cfg.NoUser {
+			ccfg.User = &cfg.User
+		}
+		if !cfg.NoCwd {
+			ccfg.Cwd = &cfg.Cwd
+		}
+		return build.Invoke(egCtx, ccfg)
+	})
+	err := eg.Wait()
+	close(waitInvokeDoneCh)
+	curInvokeCancel()
+
+	return err
+}
+
+func toControlStatus(s *client.SolveStatus) *pb.StatusResponse {
+	resp := pb.StatusResponse{}
+	for _, v := range s.Vertexes {
+		resp.Vertexes = append(resp.Vertexes, &controlapi.Vertex{
+			Digest:        v.Digest,
+			Inputs:        v.Inputs,
+			Name:          v.Name,
+			Started:       v.Started,
+			Completed:     v.Completed,
+			Error:         v.Error,
+			Cached:        v.Cached,
+			ProgressGroup: v.ProgressGroup,
+		})
+	}
+	for _, v := range s.Statuses {
+		resp.Statuses = append(resp.Statuses, &controlapi.VertexStatus{
+			ID:        v.ID,
+			Vertex:    v.Vertex,
+			Name:      v.Name,
+			Total:     v.Total,
+			Current:   v.Current,
+			Timestamp: v.Timestamp,
+			Started:   v.Started,
+			Completed: v.Completed,
+		})
+	}
+	for _, v := range s.Logs {
+		resp.Logs = append(resp.Logs, &controlapi.VertexLog{
+			Vertex:    v.Vertex,
+			Stream:    int64(v.Stream),
+			Msg:       v.Data,
+			Timestamp: v.Timestamp,
+		})
+	}
+	for _, v := range s.Warnings {
+		resp.Warnings = append(resp.Warnings, &controlapi.VertexWarning{
+			Vertex: v.Vertex,
+			Level:  int64(v.Level),
+			Short:  v.Short,
+			Detail: v.Detail,
+			Url:    v.URL,
+			Info:   v.SourceInfo,
+			Ranges: v.Range,
+		})
+	}
+	return &resp
+}

commands/controller/io.go

@@ -0,0 +1,431 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"syscall"
+	"time"
+
+	"github.com/docker/buildx/commands/controller/pb"
+	"github.com/moby/sys/signal"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
+)
+
+type msgStream interface {
+	Send(*pb.Message) error
+	Recv() (*pb.Message, error)
+}
+
+type ioServerConfig struct {
+	stdin          io.WriteCloser
+	stdout, stderr io.ReadCloser
+
+	// signalFn is a callback function called when a signal is reached to the client.
+	signalFn func(context.Context, syscall.Signal) error
+
+	// resizeFn is a callback function called when a resize event is reached to the client.
+	resizeFn func(context.Context, winSize) error
+}
+
+func serveIO(attachCtx context.Context, srv msgStream, initFn func(*pb.InitMessage) error, ioConfig *ioServerConfig) (err error) {
+	stdin, stdout, stderr := ioConfig.stdin, ioConfig.stdout, ioConfig.stderr
+	stream := &debugStream{srv, "server=" + time.Now().String()}
+	eg, ctx := errgroup.WithContext(attachCtx)
+	done := make(chan struct{})
+
+	msg, err := receive(ctx, stream)
+	if err != nil {
+		return err
+	}
+	init := msg.GetInit()
+	if init == nil {
+		return fmt.Errorf("unexpected message: %T; wanted init", msg.GetInput())
+	}
+	ref := init.Ref
+	if ref == "" {
+		return fmt.Errorf("no ref is provided")
+	}
+	if err := initFn(init); err != nil {
+		return fmt.Errorf("failed to initialize IO server: %w", err)
+	}
+
+	if stdout != nil {
+		stdoutReader, stdoutWriter := io.Pipe()
+		eg.Go(func() error {
+			<-done
+			return stdoutWriter.Close()
+		})
+
+		go func() {
+			// do not wait for read completion but return here and let the caller send EOF
+			// this allows us to return on ctx.Done() without being blocked by this reader.
+			io.Copy(stdoutWriter, stdout)
+			stdoutWriter.Close()
+		}()
+
+		eg.Go(func() error {
+			defer stdoutReader.Close()
+			return copyToStream(1, stream, stdoutReader)
+		})
+	}
+
+	if stderr != nil {
+		stderrReader, stderrWriter := io.Pipe()
+		eg.Go(func() error {
+			<-done
+			return stderrWriter.Close()
+		})
+
+		go func() {
+			// do not wait for read completion but return here and let the caller send EOF
+			// this allows us to return on ctx.Done() without being blocked by this reader.
+			io.Copy(stderrWriter, stderr)
+			stderrWriter.Close()
+		}()
+
+		eg.Go(func() error {
+			defer stderrReader.Close()
+			return copyToStream(2, stream, stderrReader)
+		})
+	}
+
+	msgCh := make(chan *pb.Message)
+	eg.Go(func() error {
+		defer close(msgCh)
+		for {
+			msg, err := receive(ctx, stream)
+			if err != nil {
+				return err
+			}
+			select {
+			case msgCh <- msg:
+			case <-done:
+				return nil
+			case <-ctx.Done():
+				return nil
+			}
+		}
+	})
+
+	eg.Go(func() error {
+		defer close(done)
+		for {
+			var msg *pb.Message
+			select {
+			case msg = <-msgCh:
+			case <-ctx.Done():
+				return nil
+			}
+			if msg == nil {
+				return nil
+			}
+			if file := msg.GetFile(); file != nil {
+				if file.Fd != 0 {
+					return fmt.Errorf("unexpected fd: %v", file.Fd)
+				}
+				if stdin == nil {
+					continue // no stdin destination is specified so ignore the data
+				}
+				if len(file.Data) > 0 {
+					_, err := stdin.Write(file.Data)
+					if err != nil {
+						return err
+					}
+				}
+				if file.EOF {
+					stdin.Close()
+				}
+			} else if resize := msg.GetResize(); resize != nil {
+				if ioConfig.resizeFn != nil {
+					ioConfig.resizeFn(ctx, winSize{
+						cols: resize.Cols,
+						rows: resize.Rows,
+					})
+				}
+			} else if sig := msg.GetSignal(); sig != nil {
+				if ioConfig.signalFn != nil {
+					syscallSignal, ok := signal.SignalMap[sig.Name]
+					if !ok {
+						continue
+					}
+					ioConfig.signalFn(ctx, syscallSignal)
+				}
+			} else {
+				return fmt.Errorf("unexpected message: %T", msg.GetInput())
+			}
+		}
+	})
+
+	return eg.Wait()
+}
+
+type ioAttachConfig struct {
+	stdin          io.ReadCloser
+	stdout, stderr io.WriteCloser
+	signal         <-chan syscall.Signal
+	resize         <-chan winSize
+}
+
+type winSize struct {
+	rows uint32
+	cols uint32
+}
+
+func attachIO(ctx context.Context, stream msgStream, initMessage *pb.InitMessage, cfg ioAttachConfig) (retErr error) {
+	eg, ctx := errgroup.WithContext(ctx)
+	done := make(chan struct{})
+
+	if err := stream.Send(&pb.Message{
+		Input: &pb.Message_Init{
+			Init: initMessage,
+		},
+	}); err != nil {
+		return fmt.Errorf("failed to init: %w", err)
+	}
+
+	if cfg.stdin != nil {
+		stdinReader, stdinWriter := io.Pipe()
+		eg.Go(func() error {
+			<-done
+			return stdinWriter.Close()
+		})
+
+		go func() {
+			// do not wait for read completion but return here and let the caller send EOF
+			// this allows us to return on ctx.Done() without being blocked by this reader.
+			io.Copy(stdinWriter, cfg.stdin)
+			stdinWriter.Close()
+		}()
+
+		eg.Go(func() error {
+			defer stdinReader.Close()
+			return copyToStream(0, stream, stdinReader)
+		})
+	}
+
+	if cfg.signal != nil {
+		eg.Go(func() error {
+			for {
+				var sig syscall.Signal
+				select {
+				case sig = <-cfg.signal:
+				case <-done:
+					return nil
+				case <-ctx.Done():
+					return nil
+				}
+				name := sigToName[sig]
+				if name == "" {
+					continue
+				}
+				if err := stream.Send(&pb.Message{
+					Input: &pb.Message_Signal{
+						Signal: &pb.SignalMessage{
+							Name: name,
+						},
+					},
+				}); err != nil {
+					return fmt.Errorf("failed to send signal: %w", err)
+				}
+			}
+		})
+	}
+
+	if cfg.resize != nil {
+		eg.Go(func() error {
+			for {
+				var win winSize
+				select {
+				case win = <-cfg.resize:
+				case <-done:
+					return nil
+				case <-ctx.Done():
+					return nil
+				}
+				if err := stream.Send(&pb.Message{
+					Input: &pb.Message_Resize{
+						Resize: &pb.ResizeMessage{
+							Rows: win.rows,
+							Cols: win.cols,
+						},
+					},
+				}); err != nil {
+					return fmt.Errorf("failed to send resize: %w", err)
+				}
+			}
+		})
+	}
+
+	msgCh := make(chan *pb.Message)
+	eg.Go(func() error {
+		defer close(msgCh)
+		for {
+			msg, err := receive(ctx, stream)
+			if err != nil {
+				return err
+			}
+			select {
+			case msgCh <- msg:
+			case <-done:
+				return nil
+			case <-ctx.Done():
+				return nil
+			}
+		}
+	})
+
+	eg.Go(func() error {
+		eofs := make(map[uint32]struct{})
+		defer close(done)
+		for {
+			var msg *pb.Message
+			select {
+			case msg = <-msgCh:
+			case <-ctx.Done():
+				return nil
+			}
+			if msg == nil {
+				return nil
+			}
+			if file := msg.GetFile(); file != nil {
+				if _, ok := eofs[file.Fd]; ok {
+					continue
+				}
+				var out io.WriteCloser
+				switch file.Fd {
+				case 1:
+					out = cfg.stdout
+				case 2:
+					out = cfg.stderr
+				default:
+					return fmt.Errorf("unsupported fd %d", file.Fd)
+
+				}
+				if out == nil {
+					logrus.Warnf("attachIO: no writer for fd %d", file.Fd)
+					continue
+				}
+				if len(file.Data) > 0 {
+					if _, err := out.Write(file.Data); err != nil {
+						return err
+					}
+				}
+				if file.EOF {
+					eofs[file.Fd] = struct{}{}
+				}
+			} else {
+				return fmt.Errorf("unexpected message: %T", msg.GetInput())
+			}
+		}
+	})
+
+	return eg.Wait()
+}
+
+func receive(ctx context.Context, stream msgStream) (*pb.Message, error) {
+	msgCh := make(chan *pb.Message)
+	errCh := make(chan error)
+	go func() {
+		msg, err := stream.Recv()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return
+			}
+			errCh <- err
+			return
+		}
+		msgCh <- msg
+	}()
+	select {
+	case msg := <-msgCh:
+		return msg, nil
+	case err := <-errCh:
+		return nil, err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+func copyToStream(fd uint32, snd msgStream, r io.Reader) error {
+	for {
+		buf := make([]byte, 32*1024)
+		n, err := r.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break // break loop and send EOF
+			}
+			return err
+		} else if n > 0 {
+			if snd.Send(&pb.Message{
+				Input: &pb.Message_File{
+					File: &pb.FdMessage{
+						Fd:   fd,
+						Data: buf[:n],
+					},
+				},
+			}); err != nil {
+				return err
+			}
+		}
+	}
+	return snd.Send(&pb.Message{
+		Input: &pb.Message_File{
+			File: &pb.FdMessage{
+				Fd:  fd,
+				EOF: true,
+			},
+		},
+	})
+}
+
+var sigToName = map[syscall.Signal]string{}
+
+func init() {
+	for name, value := range signal.SignalMap {
+		sigToName[value] = name
+	}
+}
+
+type debugStream struct {
+	msgStream
+	prefix string
+}
+
+func (s *debugStream) Send(msg *pb.Message) error {
+	switch m := msg.GetInput().(type) {
+	case *pb.Message_File:
+		if m.File.EOF {
+			logrus.Debugf("|---> File Message (sender:%v) fd=%d, EOF", s.prefix, m.File.Fd)
+		} else {
+			logrus.Debugf("|---> File Message (sender:%v) fd=%d, %d bytes", s.prefix, m.File.Fd, len(m.File.Data))
+		}
+	case *pb.Message_Resize:
+		logrus.Debugf("|---> Resize Message (sender:%v): %+v", s.prefix, m.Resize)
+	case *pb.Message_Signal:
+		logrus.Debugf("|---> Signal Message (sender:%v): %s", s.prefix, m.Signal.Name)
+	}
+	return s.msgStream.Send(msg)
+}
+
+func (s *debugStream) Recv() (*pb.Message, error) {
+	msg, err := s.msgStream.Recv()
+	if err != nil {
+		return nil, err
+	}
+	switch m := msg.GetInput().(type) {
+	case *pb.Message_File:
+		if m.File.EOF {
+			logrus.Debugf("|<--- File Message (receiver:%v) fd=%d, EOF", s.prefix, m.File.Fd)
+		} else {
+			logrus.Debugf("|<--- File Message (receiver:%v) fd=%d, %d bytes", s.prefix, m.File.Fd, len(m.File.Data))
+		}
+	case *pb.Message_Resize:
+		logrus.Debugf("|<--- Resize Message (receiver:%v): %+v", s.prefix, m.Resize)
+	case *pb.Message_Signal:
+		logrus.Debugf("|<--- Signal Message (receiver:%v): %s", s.prefix, m.Signal.Name)
+	}
+	return msg, nil
+}

commands/controller/pb/controller.proto

@@ -0,0 +1,179 @@
+syntax = "proto3";
+
+package buildx.controller.v1;
+
+import "github.com/moby/buildkit/api/services/control/control.proto";
+
+option go_package = "pb";
+
+service Controller {
+  rpc Build(BuildRequest) returns (BuildResponse);
+  rpc Status(StatusRequest) returns (stream StatusResponse);
+  rpc Input(stream InputMessage) returns (InputResponse);
+  rpc Invoke(stream Message) returns (stream Message);
+  rpc List(ListRequest) returns (ListResponse);
+  rpc Disconnect(DisconnectRequest) returns (DisconnectResponse);
+  rpc Info(InfoRequest) returns (InfoResponse);
+}
+
+message BuildRequest {
+  string Ref = 1;
+  BuildOptions Options = 2;
+}
+
+message BuildOptions {
+  string ContextPath = 1;
+  string DockerfileName = 2;
+  string PrintFunc = 3;
+
+  repeated string Allow = 4;
+  repeated string Attests = 5; // TODO
+  repeated string BuildArgs = 6;
+  repeated string CacheFrom = 7;
+  repeated string CacheTo = 8;
+  string CgroupParent = 9;
+  repeated string Contexts = 10;
+  repeated string ExtraHosts = 11;
+  string ImageIDFile = 12;
+  repeated string Labels = 13;
+  string NetworkMode = 14;
+  repeated string NoCacheFilter = 15;
+  repeated string Outputs = 16;
+  repeated string Platforms = 17;
+  bool Quiet = 18;
+  repeated string Secrets = 19;
+  int64 ShmSize = 20;
+  repeated string SSH = 21;
+  repeated string Tags = 22;
+  string Target = 23;
+  UlimitOpt Ulimits = 24;
+  // string Invoke: provided via Invoke API
+  CommonOptions Opts = 25;
+}
+
+message UlimitOpt {
+  map<string, Ulimit> values = 1;
+}
+
+message Ulimit {
+  string Name = 1;
+  int64 Hard = 2;
+  int64 Soft = 3;
+}
+
+message CommonOptions {
+  string Builder = 1;
+  string MetadataFile = 2;
+  bool NoCache = 3;
+  // string Progress: no progress view on server side
+  bool Pull = 4;
+  bool ExportPush = 5;
+  bool ExportLoad = 6;
+  string SBOM       = 7; // TODO
+  string Provenance = 8; // TODO
+}
+
+message BuildResponse {}
+
+message DisconnectRequest {
+  string Ref = 1;
+}
+
+message DisconnectResponse {}
+
+message ListRequest {
+  string Ref = 1;
+}
+
+message ListResponse {
+  repeated string keys = 1;
+}
+
+message InputMessage {
+  oneof Input {
+    InputInitMessage Init = 1;
+    DataMessage Data = 2;
+  }
+}
+
+message InputInitMessage {
+  string Ref = 1;
+}
+
+message DataMessage {
+  bool EOF = 1;    // true if eof was reached
+  bytes Data = 2;  // should be chunked smaller than 4MB:
+                   // https://pkg.go.dev/google.golang.org/grpc#MaxRecvMsgSize
+}
+
+message InputResponse {}
+
+message Message {
+  oneof Input {
+    InitMessage Init = 1;
+    // FdMessage used from client to server for input (stdin) and
+    // from server to client for output (stdout, stderr)
+    FdMessage File = 2;
+    // ResizeMessage used from client to server for terminal resize events
+    ResizeMessage Resize = 3;
+    // SignalMessage is used from client to server to send signal events
+    SignalMessage Signal = 4;
+  }
+}
+
+message InitMessage {
+  string Ref = 1;
+  ContainerConfig ContainerConfig = 2;
+}
+
+message ContainerConfig {
+  repeated string Entrypoint = 1;
+  repeated string Cmd = 2;
+  repeated string Env = 3;
+  string User = 4;
+  bool NoUser = 5; // Do not set user but use the image's default
+  string Cwd = 6;
+  bool NoCwd = 7; // Do not set cwd but use the image's default
+  bool Tty = 8;
+}
+
+message FdMessage {
+  uint32 Fd = 1;   // what fd the data was from
+  bool EOF = 2;    // true if eof was reached
+  bytes Data = 3;  // should be chunked smaller than 4MB:
+                   // https://pkg.go.dev/google.golang.org/grpc#MaxRecvMsgSize
+}
+
+message ResizeMessage {
+  uint32 Rows = 1;
+  uint32 Cols = 2;
+}
+
+message SignalMessage {
+  // we only send name (ie HUP, INT) because the int values
+  // are platform dependent.
+  string Name = 1;
+}
+
+message StatusRequest {
+  string Ref = 1;
+}
+
+message StatusResponse {
+  repeated moby.buildkit.v1.Vertex vertexes = 1;
+  repeated moby.buildkit.v1.VertexStatus statuses = 2;
+  repeated moby.buildkit.v1.VertexLog logs = 3;
+  repeated moby.buildkit.v1.VertexWarning warnings = 4;
+}
+
+message InfoRequest {}
+
+message InfoResponse {
+  BuildxVersion buildxVersion = 1;
+}
+
+message BuildxVersion {
+  string package = 1;
+  string version = 2;
+  string revision = 3;
+}

commands/controller/pb/generate.go

@@ -0,0 +1,3 @@
+package pb
+
+//go:generate protoc -I=. -I=../../../vendor/ --gogo_out=plugins=grpc:. controller.proto