[wip] capture output from run to use elsewhere

codegen/chain.go

@@ -15,6 +15,7 @@ import (
 	gateway "github.com/moby/buildkit/frontend/gateway/client"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/solver/pb"
+	digest "github.com/opencontainers/go-digest"
 	"github.com/openllb/hlb/local"
 	"github.com/openllb/hlb/parser"
 	"github.com/openllb/hlb/solver"
@@ -349,13 +350,22 @@ func (cg *CodeGen) EmitFilesystemBuiltinChainStmt(ctx context.Context, scope *pa
 				// to be in the context of a specific function run is in.
 			case with.Expr.FuncLit != nil:
 				for _, stmt := range with.Expr.FuncLit.Body.NonEmptyStmts() {
-					if stmt.Call.Func.Name() != "mount" || stmt.Call.Alias == nil {
+					if stmt.Call.Alias == nil {
 						continue
 					}
 
-					target, err := cg.EmitStringExpr(ctx, scope, stmt.Call.Args[1])
-					if err != nil {
-						return fc, err
+					var target string
+					switch stmt.Call.Func.Name() {
+					case "mount":
+						target, err = cg.EmitStringExpr(ctx, scope, stmt.Call.Args[1])
+						if err != nil {
+							return fc, err
+						}
+					case "capture":
+						target = "capture"
+					}
+					if target == "" {
+						continue
 					}
 
 					calls[target] = stmt.Call
@@ -381,7 +391,47 @@ func (cg *CodeGen) EmitFilesystemBuiltinChainStmt(ctx context.Context, scope *pa
 				for _, target := range targets {
 					// Mounts are unique by its mountpoint, and its vertex representing the
 					// mount after execing can be aliased.
-					cont := ac(calls[target], exec.GetMount(target))
+					var cont bool
+					switch calls[target].Func.Name() {
+					case "mount":
+						cont = ac(calls[target], exec.GetMount(target))
+					case "capture":
+						cont = ac(calls[target], func() (string, error) {
+							st := exec.Root()
+							pw := cg.mw.WithPrefix("", false)
+
+							s, err := cg.newSession(ctx)
+							if err != nil {
+								return "", err
+							}
+
+							g, ctx := errgroup.WithContext(ctx)
+
+							g.Go(func() error {
+								return s.Run(ctx, cg.cln.Dialer())
+							})
+
+							var captureBuf strings.Builder
+							g.Go(func() error {
+								opts, err := cg.SolveOptions(ctx, st)
+								if err != nil {
+									return err
+								}
+								def, err := st.Marshal(ctx, llb.LinuxAmd64)
+								if err != nil {
+									return err
+								}
+
+								// * last def is for constraint
+								// * second to last def is for the vertex we are capturing
+								dgst := digest.FromBytes(def.Def[len(def.Def)-2])
+								opts = append(opts, solver.WithOutputCapture(dgst, &captureBuf))
+								return solver.Solve(ctx, cg.cln, s, pw, def, opts...)
+							})
+							err = g.Wait()
+							return captureBuf.String(), err
+						})
+					}
 					if !cont {
 						return exec.Root(), ErrAliasReached
 					}
@@ -940,11 +990,13 @@ func (cg *CodeGen) EmitStringChainStmt(ctx context.Context, scope *parser.Scope,
 			return nil, err
 		}
 		return func(_ string) (string, error) {
-			str, ok := v.(string)
-			if !ok {
-				return str, errors.WithStack(ErrCodeGen{obj.Node, ErrBadCast})
+			switch s := v.(type) {
+			case string:
+				return s, nil
+			case func() (string, error):
+				return s()
 			}
-			return str, nil
+			return "", errors.WithStack(ErrCodeGen{obj.Node, ErrBadCast})
 		}, nil
 	}
 }

codegen/codegen.go

@@ -1075,6 +1075,8 @@ func (cg *CodeGen) EmitExecOptions(ctx context.Context, scope *parser.Scope, op
 				opts = append(opts, llb.Security(securityMode))
 			case "shlex":
 				opts = append(opts, &shlexOption{})
+			case "capture":
+				// no op, only relevant if aliased, handled in alias callback
 			case "host":
 				host, err := cg.EmitStringExpr(ctx, scope, args[0])
 				if err != nil {

docs/reference.md

@@ -480,6 +480,7 @@ If more than one arg is given, it will be executed directly, without a shell.
 	#!hlb
 	fs default() {
 		run "arg" with option {
+			capture
 			dir "path"
 			env "key" "value"
 			forward "src" "dest"
@@ -497,6 +498,11 @@ If more than one arg is given, it will be executed directly, without a shell.
 	}
 
 
+#### <span class='hlb-type'>option::run</span> <span class='hlb-name'>capture</span>()
+
+
+
+
 #### <span class='hlb-type'>option::run</span> <span class='hlb-name'>dir</span>(<span class='hlb-type'>string</span> <span class='hlb-variable'>path</span>)
 
 !!! info "<span class='hlb-type'>string</span> <span class='hlb-variable'>path</span>"

language/builtin.hlb

@@ -120,6 +120,8 @@ fs shell(variadic string arg)
 # @return the filesystem after the command has executed.
 fs run(variadic string arg)
 
+option::run capture()
+
 # Sets the rootfs as read-only for the duration of the run command.
 #
 # @return an option to set the rootfs as read-only.

solver/solve.go

@@ -3,6 +3,7 @@ package solver
 import (
 	"context"
 	"encoding/json"
+	"io"
 
 	"github.com/docker/buildx/util/progress"
 	"github.com/moby/buildkit/client"
@@ -11,6 +12,7 @@ import (
 	gateway "github.com/moby/buildkit/frontend/gateway/client"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/util/entitlements"
+	digest "github.com/opencontainers/go-digest"
 	specs "github.com/opencontainers/image-spec/specs-go/v1"
 	"golang.org/x/sync/errgroup"
 )
@@ -26,6 +28,8 @@ type SolveInfo struct {
 	Callbacks             []func() error `json:"-"`
 	ImageSpec             *specs.Image
 	Entitlements          []entitlements.Entitlement
+	OutputCapture         io.Writer
+	OutputCaptureDigest   digest.Digest
 }
 
 func WithDownloadDockerTarball(ref string) SolveOption {
@@ -63,6 +67,14 @@ func WithDownloadOCITarball() SolveOption {
 	}
 }
 
+func WithOutputCapture(dgst digest.Digest, w io.Writer) SolveOption {
+	return func(info *SolveInfo) error {
+		info.OutputCaptureDigest = dgst
+		info.OutputCapture = w
+		return nil
+	}
+}
+
 func WithCallback(fn func() error) SolveOption {
 	return func(info *SolveInfo) error {
 		info.Callbacks = append(info.Callbacks, fn)
@@ -174,6 +186,38 @@ func Build(ctx context.Context, c *client.Client, s *session.Session, pw progres
 		statusCh = pw.Status()
 	}
 
+	if info.OutputCapture != nil {
+		captureStatusCh := make(chan *client.SolveStatus)
+		go func(origStatusCh chan *client.SolveStatus) {
+			defer func() {
+				if origStatusCh != nil {
+					close(origStatusCh)
+				}
+			}()
+			for {
+				select {
+				case <-pw.Done():
+					return
+				case <-ctx.Done():
+					return
+				case status, ok := <-captureStatusCh:
+					if !ok {
+						return
+					}
+					for _, log := range status.Logs {
+						if log.Vertex.String() == info.OutputCaptureDigest.String() {
+							info.OutputCapture.Write(log.Data)
+						}
+					}
+					if origStatusCh != nil {
+						origStatusCh <- status
+					}
+				}
+			}
+		}(statusCh)
+		statusCh = captureStatusCh
+	}
+
 	g.Go(func() error {
 		_, err := c.Build(ctx, solveOpt, "", f, statusCh)
 		return err