llb: automatically carry platform with state chain

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

client/llb/exec.go

@@ -28,10 +28,13 @@ func NewExecOp(root Output, meta Meta, readOnly bool, c Constraints) *ExecOp {
 	if readOnly {
 		e.root = root
 	} else {
-		e.root = &output{vertex: e, getIndex: e.getMountIndexFn(rootMount)}
+		o := &output{vertex: e, getIndex: e.getMountIndexFn(rootMount)}
+		if p := c.Platform; p != nil {
+			o.platform = p
+		}
+		e.root = o
 	}
 	rootMount.output = e.root
-
 	return e
 }
 
@@ -70,7 +73,11 @@ func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Outp
 	} else if m.tmpfs {
 		m.output = &output{vertex: e, err: errors.Errorf("tmpfs mount for %s can't be used as a parent", target)}
 	} else {
-		m.output = &output{vertex: e, getIndex: e.getMountIndexFn(m)}
+		o := &output{vertex: e, getIndex: e.getMountIndexFn(m)}
+		if p := e.constraints.Platform; p != nil {
+			o.platform = p
+		}
+		m.output = o
 	}
 	e.Store(nil, nil, nil)
 	e.isValidated = false

client/llb/meta.go

@@ -4,16 +4,19 @@ import (
 	"fmt"
 	"path"
 
+	"github.com/containerd/containerd/platforms"
 	"github.com/google/shlex"
+	specs "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
 type contextKeyT string
 
 var (
-	keyArgs = contextKeyT("llb.exec.args")
-	keyDir  = contextKeyT("llb.exec.dir")
-	keyEnv  = contextKeyT("llb.exec.env")
-	keyUser = contextKeyT("llb.exec.user")
+	keyArgs     = contextKeyT("llb.exec.args")
+	keyDir      = contextKeyT("llb.exec.dir")
+	keyEnv      = contextKeyT("llb.exec.env")
+	keyUser     = contextKeyT("llb.exec.user")
+	keyPlatform = contextKeyT("llb.platform")
 )
 
 func addEnv(key, value string) StateOption {
@@ -106,6 +109,21 @@ func shlexf(str string, v ...interface{}) StateOption {
 	}
 }
 
+func platform(p specs.Platform) StateOption {
+	return func(s State) State {
+		return s.WithValue(keyPlatform, platforms.Normalize(p))
+	}
+}
+
+func getPlatform(s State) *specs.Platform {
+	v := s.Value(keyPlatform)
+	if v != nil {
+		p := v.(specs.Platform)
+		return &p
+	}
+	return nil
+}
+
 type EnvList []KeyValue
 
 type KeyValue struct {

client/llb/source.go

@@ -11,6 +11,7 @@ import (
 	"github.com/docker/distribution/reference"
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
+	specs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
 
@@ -33,6 +34,10 @@ func NewSource(id string, attrs map[string]string, c Constraints) *SourceOp {
 	return s
 }
 
+func (s *SourceOp) Platform() *specs.Platform {
+	return s.constraints.Platform
+}
+
 func (s *SourceOp) Validate() error {
 	if s.err != nil {
 		return s.err
@@ -82,7 +87,7 @@ func Image(ref string, opts ...ImageOption) State {
 	for _, opt := range opts {
 		opt.SetImageOption(&info)
 	}
-	src := NewSource("docker-image://"+ref, nil, info.Metadata()) // controversial
+	src := NewSource("docker-image://"+ref, nil, info.Constraints) // controversial
 	if err != nil {
 		src.err = err
 	}
@@ -164,7 +169,7 @@ func Git(remote, ref string, opts ...GitOption) State {
 	if url != "" {
 		attrs[pb.AttrFullRemoteURL] = url
 	}
-	source := NewSource("git://"+id, attrs, gi.Metadata())
+	source := NewSource("git://"+id, attrs, gi.Constraints)
 	return NewState(source.Output())
 }
 
@@ -215,7 +220,7 @@ func Local(name string, opts ...LocalOption) State {
 		attrs[pb.AttrSharedKeyHint] = gi.SharedKeyHint
 	}
 
-	source := NewSource("local://"+name, attrs, gi.Metadata())
+	source := NewSource("local://"+name, attrs, gi.Constraints)
 	return NewState(source.Output())
 }
 
@@ -305,7 +310,7 @@ func HTTP(url string, opts ...HTTPOption) State {
 		attrs[pb.AttrHTTPGID] = strconv.Itoa(hi.GID)
 	}
 
-	source := NewSource(url, attrs, hi.Metadata())
+	source := NewSource(url, attrs, hi.Constraints)
 	return NewState(source.Output())
 }
 

client/llb/state.go

@@ -31,6 +31,14 @@ func NewState(o Output) State {
 	}
 	s = dir("/")(s)
 	s = addEnv("PATH", system.DefaultPathEnv)(s)
+
+	if o, ok := o.(interface {
+		Platform() *specs.Platform
+	}); ok {
+		if p := o.Platform(); p != nil {
+			s = platform(*p)(s)
+		}
+	}
 	return s
 }
 
@@ -128,6 +136,9 @@ func (s State) WithOutput(o Output) State {
 
 func (s State) Run(ro ...RunOption) ExecState {
 	ei := &ExecInfo{State: s}
+	if p := s.GetPlatform(); p != nil {
+		ei.Constraints.Platform = p
+	}
 	for _, o := range ro {
 		o.SetRunOption(ei)
 	}
@@ -139,7 +150,7 @@ func (s State) Run(ro ...RunOption) ExecState {
 		ProxyEnv: ei.ProxyEnv,
 	}
 
-	exec := NewExecOp(s.Output(), meta, ei.ReadonlyRootFS, ei.Metadata())
+	exec := NewExecOp(s.Output(), meta, ei.ReadonlyRootFS, ei.Constraints)
 	for _, m := range ei.Mounts {
 		exec.AddMount(m.Target, m.Source, m.Opts...)
 	}
@@ -185,6 +196,14 @@ func (s State) User(v string) State {
 	return user(v)(s)
 }
 
+func (s State) Platform(p specs.Platform) State {
+	return platform(p)(s)
+}
+
+func (s State) GetPlatform() *specs.Platform {
+	return getPlatform(s)
+}
+
 func (s State) With(so ...StateOption) State {
 	for _, o := range so {
 		s = o(s)
@@ -196,6 +215,7 @@ type output struct {
 	vertex   Vertex
 	getIndex func() (pb.OutputIndex, error)
 	err      error
+	platform *specs.Platform
 }
 
 func (o *output) ToInput(c *Constraints) (*pb.Input, error) {
@@ -221,6 +241,10 @@ func (o *output) Vertex() Vertex {
 	return o.vertex
 }
 
+func (o *output) Platform() *specs.Platform {
+	return o.platform
+}
+
 type ConstraintsOpt interface {
 	SetConstraintsOption(*Constraints)
 	RunOption
@@ -318,10 +342,6 @@ func (cw *constraintsWrapper) applyConstraints(f func(c *Constraints)) {
 	f(&cw.Constraints)
 }
 
-func (cw *constraintsWrapper) Metadata() Constraints {
-	return cw.Constraints
-}
-
 type Constraints struct {
 	Platform          *specs.Platform
 	WorkerConstraints []string