diff --git a/solver/llbsolver/cdidevices/manager.go b/solver/llbsolver/cdidevices/manager.go index 7df61f371a2a..472f5b246796 100644 --- a/solver/llbsolver/cdidevices/manager.go +++ b/solver/llbsolver/cdidevices/manager.go @@ -152,7 +152,25 @@ func (m *Manager) FindDevices(devs ...*pb.CDIDevice) ([]string, error) { return dedupSlice(out), nil } +// MatchDevices returns the registered devices that the request resolves to, +// including matches by class annotation. Unlike FindDevices it does not fail +// or warn when nothing matches. +func (m *Manager) MatchDevices(dev *pb.CDIDevice) []string { + return m.matchDevice(dev, m.cache.ListDevices()) +} + func (m *Manager) parseDevice(dev *pb.CDIDevice, all []string) ([]string, error) { + out := m.matchDevice(dev, all) + if len(out) == 0 { + if !dev.Optional { + return nil, errors.Errorf("required device %q is not registered", dev.Name) + } + bklog.G(context.TODO()).Warnf("Optional device %q is not registered", dev.Name) + } + return out, nil +} + +func (m *Manager) matchDevice(dev *pb.CDIDevice, all []string) []string { var out []string kind, name, _ := strings.Cut(dev.Name, "=") @@ -197,13 +215,7 @@ func (m *Manager) parseDevice(dev *pb.CDIDevice, all []string) ([]string, error) } } - if len(out) == 0 { - if !dev.Optional { - return nil, errors.Errorf("required device %q is not registered", dev.Name) - } - bklog.G(context.TODO()).Warnf("Optional device %q is not registered", dev.Name) - } - return out, nil + return out } func (m *Manager) hasDevice(name string) bool { diff --git a/solver/llbsolver/vertex.go b/solver/llbsolver/vertex.go index 3a661d94d3c6..99dd9f423072 100644 --- a/solver/llbsolver/vertex.go +++ b/solver/llbsolver/vertex.go @@ -167,26 +167,37 @@ func ValidateEntitlements(ent entitlements.Set, cdiManager *cdidevices.Manager) } var allowedDevices []*pb.CDIDevice - var nonAliasedDevices []*pb.CDIDevice + var devicesToResolve []*pb.CDIDevice + onDemand := make(map[string]struct{}) for _, d := range cdiManager.ListDevices() { if d.OnDemand && d.AutoAllow { - allowedDevices = append(allowedDevices, &pb.CDIDevice{Name: d.Name}) + onDemand[d.Name] = struct{}{} } } - if cfg != nil { - for _, d := range op.Exec.CdiDevices { + for _, d := range op.Exec.CdiDevices { + if d == nil { + continue + } + if cfg != nil { if newName, ok := cfg.Devices[d.Name]; ok && newName != "" { d.Name = newName allowedDevices = append(allowedDevices, d) - } else { - nonAliasedDevices = append(nonAliasedDevices, d) + continue } } - } else { - nonAliasedDevices = op.Exec.CdiDevices + kind, _, _ := strings.Cut(d.Name, "=") + // A kind with an on-demand installer has no registered devices + // yet, so a match here can only come from a class annotation + // on another device. That must go through the permission check. + if _, ok := onDemand[kind]; ok && len(cdiManager.MatchDevices(d)) == 0 { + // Resolve the requested name after the installer has registered its devices. + allowedDevices = append(allowedDevices, d) + } else { + devicesToResolve = append(devicesToResolve, d) + } } - mountedDevices, err := cdiManager.FindDevices(nonAliasedDevices...) + mountedDevices, err := cdiManager.FindDevices(devicesToResolve...) if err != nil { return err } @@ -198,6 +209,10 @@ func ValidateEntitlements(ent entitlements.Set, cdiManager *cdidevices.Manager) grantedDevices := make(map[string]struct{}) if cfg != nil { for d := range cfg.Devices { + kind, _, _ := strings.Cut(d, "=") + if _, ok := onDemand[kind]; ok && len(cdiManager.MatchDevices(&pb.CDIDevice{Name: d})) == 0 { + continue + } resolved, err := cdiManager.FindDevices(&pb.CDIDevice{Name: d}) if err != nil { return err diff --git a/solver/llbsolver/vertex_test.go b/solver/llbsolver/vertex_test.go index 49953a6e2e4f..5415505a3fd4 100644 --- a/solver/llbsolver/vertex_test.go +++ b/solver/llbsolver/vertex_test.go @@ -1,15 +1,22 @@ package llbsolver import ( + "context" _ "embed" "fmt" + "os" + "path/filepath" "testing" + "github.com/moby/buildkit/solver" + "github.com/moby/buildkit/solver/llbsolver/cdidevices" "github.com/moby/buildkit/solver/pb" "github.com/moby/buildkit/util/entitlements" digest "github.com/opencontainers/go-digest" + "github.com/pkg/errors" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "tags.cncf.io/container-device-interface/pkg/cdi" ) func TestRecomputeDigests(t *testing.T) { @@ -424,3 +431,197 @@ func requireVertexOp(t *testing.T, v interface{ Sys() any }) *pb.Op { require.True(t, ok) return op } + +func TestValidateEntitlementsOnDemandDevices(t *testing.T) { + const kind = "on-demand.test/gpu" + setup := &entitlementTestSetup{available: true} + cdidevices.Register(kind, setup) + t.Cleanup(func() { setup.available = false }) + + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "device.json"), []byte(`{ + "cdiVersion": "0.6.0", + "kind": "registered.test/device", + "devices": [{"name": "foo", "containerEdits": {"env": ["TEST_DEVICE=1"]}}] + }`), 0600)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "classed.json"), []byte(`{ + "cdiVersion": "0.6.0", + "kind": "classed.test/device", + "devices": [{"name": "bar", "annotations": {"org.mobyproject.buildkit.device.class": "on-demand.test/gpu=classed"}, "containerEdits": {"env": ["TEST_DEVICE=2"]}}] + }`), 0600)) + installedDir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(installedDir, "gpu.json"), []byte(`{ + "cdiVersion": "0.6.0", + "kind": "on-demand.test/gpu", + "devices": [{"name": "all", "containerEdits": {"env": ["TEST_GPU=1"]}}] + }`), 0600)) + + for _, tc := range []struct { + name string + devices []*pb.CDIDevice + cfg *entitlements.DevicesConfig + available bool + installed bool + want []*pb.CDIDevice + wantErr string + }{ + { + name: "required kind", + devices: []*pb.CDIDevice{{Name: kind}}, + available: true, + want: []*pb.CDIDevice{{Name: kind}}, + }, + { + name: "required named device", + devices: []*pb.CDIDevice{{Name: kind + "=all"}}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=all"}}, + }, + { + name: "optional named device", + devices: []*pb.CDIDevice{{Name: kind + "=1", Optional: true}}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=1", Optional: true}}, + }, + { + name: "wildcard", + devices: []*pb.CDIDevice{{Name: kind + "=*", Optional: true}}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=*", Optional: true}}, + }, + { + name: "unrelated registered device", + devices: []*pb.CDIDevice{{Name: "registered.test/device=foo"}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{"registered.test/device": ""}}, + available: true, + want: []*pb.CDIDevice{{Name: "registered.test/device=foo"}}, + }, + { + name: "unrelated missing optional device", + devices: []*pb.CDIDevice{{Name: "missing.test/device", Optional: true}}, + available: true, + }, + { + name: "unrelated missing required device", + devices: []*pb.CDIDevice{{Name: "missing.test/device"}}, + available: true, + wantErr: `required device "missing.test/device" is not registered`, + }, + { + name: "unavailable optional installer", + devices: []*pb.CDIDevice{{Name: kind, Optional: true}}, + }, + { + name: "unavailable required installer", + devices: []*pb.CDIDevice{{Name: kind}}, + wantErr: `required device "on-demand.test/gpu" is not registered`, + }, + { + name: "alias", + devices: []*pb.CDIDevice{{Name: "gpu", Optional: true}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{"gpu": kind + "=all"}}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=all", Optional: true}}, + }, + { + name: "allow all preserves selection", + devices: []*pb.CDIDevice{{Name: kind + "=all", Optional: true}}, + cfg: &entitlements.DevicesConfig{All: true}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=all", Optional: true}}, + }, + { + name: "mixed on-demand and registered grants", + devices: []*pb.CDIDevice{{Name: kind + "=all"}, {Name: "registered.test/device=foo"}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{kind: "", "registered.test/device": ""}}, + available: true, + want: []*pb.CDIDevice{{Name: kind + "=all"}, {Name: "registered.test/device=foo"}}, + }, + { + name: "class annotation on on-demand name still requires permission", + devices: []*pb.CDIDevice{{Name: kind + "=classed"}}, + available: true, + wantErr: `device classed.test/device=bar is requested by the build but not allowed`, + }, + { + name: "class annotation on on-demand name with explicit grant", + devices: []*pb.CDIDevice{{Name: kind + "=classed"}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{"classed.test/device=bar": ""}}, + available: true, + want: []*pb.CDIDevice{{Name: "classed.test/device=bar"}}, + }, + { + name: "class annotation on on-demand name with grant by class", + devices: []*pb.CDIDevice{{Name: kind + "=classed"}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{kind + "=classed": ""}}, + available: true, + want: []*pb.CDIDevice{{Name: "classed.test/device=bar"}}, + }, + { + name: "nil device is ignored", + devices: []*pb.CDIDevice{nil, {Name: kind}}, + available: true, + want: []*pb.CDIDevice{{Name: kind}}, + }, + { + name: "registered device still requires permission", + devices: []*pb.CDIDevice{{Name: kind}, {Name: "registered.test/device=foo"}}, + available: true, + wantErr: `device registered.test/device=foo is requested by the build but not allowed`, + }, + { + name: "installed kind is not implicitly allowed", + devices: []*pb.CDIDevice{{Name: kind + "=all"}}, + available: true, + installed: true, + wantErr: `device on-demand.test/gpu=all is requested by the build but not allowed`, + }, + { + name: "installed kind with explicit grant", + devices: []*pb.CDIDevice{{Name: kind + "=all"}}, + cfg: &entitlements.DevicesConfig{Devices: map[string]string{kind: ""}}, + available: true, + installed: true, + want: []*pb.CDIDevice{{Name: kind + "=all"}}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + setup.available = tc.available + dirs := []string{dir} + if tc.installed { + dirs = append(dirs, installedDir) + } + cache, err := cdi.NewCache(cdi.WithSpecDirs(dirs...), cdi.WithAutoRefresh(false)) + require.NoError(t, err) + manager := cdidevices.NewManager(cache, nil) + ent := entitlements.Set{} + if tc.cfg != nil { + ent[entitlements.EntitlementDevice] = tc.cfg + } + exec := &pb.ExecOp{CdiDevices: tc.devices} + op := &pb.Op{Op: &pb.Op_Exec{Exec: exec}} + err = ValidateEntitlements(ent, manager)(op, nil, &solver.VertexOptions{}) + if tc.wantErr != "" { + require.EqualError(t, err, tc.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, tc.want, exec.CdiDevices) + }) + } +} + +type entitlementTestSetup struct { + available bool +} + +func (s *entitlementTestSetup) Validate() error { + if !s.available { + return errors.New("device unavailable") + } + return nil +} + +func (*entitlementTestSetup) Run(context.Context) error { + panic("entitlement validation must not run the installer") +}