Merge pull request #11498 from tianyuzhou95:albert/cpuid

PiperOrigin-RevId: 794363944

Author: gvisor-bot

pkg/cpuid/native_amd64.go

@@ -163,6 +163,23 @@ func (fs FeatureSet) query(fn cpuidFunction) (uint32, uint32, uint32, uint32) {
 	return out.Eax, out.Ebx, out.Ecx, out.Edx
 }
 
+// Intersect returns the intersection of features between self and allowedFeatures.
+func (fs FeatureSet) Intersect(allowedFeatures map[Feature]struct{}) (FeatureSet, error) {
+	hs := fs.ToStatic()
+
+	// only keep features inside allowedFeatures.
+	for f := range allFeatures {
+		if fs.HasFeature(f) {
+			if _, ok := allowedFeatures[f]; !ok {
+				log.Infof("Removing CPU feature %v as it is not allowed.", f)
+				hs.Remove(f)
+			}
+		}
+	}
+
+	return hs.ToFeatureSet(), nil
+}
+
 var hostFeatureSet FeatureSet
 
 // HostFeatureSet returns a host CPUID.

pkg/cpuid/native_arm64.go

@@ -18,6 +18,7 @@
 package cpuid
 
 import (
+	"fmt"
 	"os"
 	"runtime"
 	"strconv"
@@ -41,6 +42,13 @@ func (fs FeatureSet) Fixed() FeatureSet {
 	return fs
 }
 
+// Intersect returns the intersection of features between self and allowedFeatures.
+//
+// Just return error as there is no ARM64 equivalent to cpuid.Static.Remove().
+func (fs FeatureSet) Intersect(allowedFeatures map[Feature]struct{}) (FeatureSet, error) {
+	return FeatureSet{}, fmt.Errorf("FeatureSet intersection is not supported on ARM64")
+}
+
 // Reads CPU information from host /proc/cpuinfo.
 //
 // Must run before syscall filter installation. This value is used to create

runsc/boot/loader.go

@@ -21,6 +21,7 @@ import (
 	"os"
 	"runtime"
 	"strconv"
+	"strings"
 	gtime "time"
 
 	"github.com/moby/sys/capability"
@@ -597,6 +598,26 @@ func New(args Args) (*Loader, error) {
 		log.Infof("Setting total memory to %.2f GB", float64(args.TotalMem)/(1<<30))
 	}
 
+	cpufs := cpuid.HostFeatureSet()
+	if value, ok := args.Spec.Annotations[specutils.AnnotationCPUFeatures]; ok {
+		allowedFeatures := make(map[cpuid.Feature]struct{})
+		for str := range strings.SplitSeq(value, ",") {
+			if str == "" {
+				continue // ignore empty feature name rather than return error
+			}
+			f, ok := cpuid.FeatureFromString(str)
+			if !ok {
+				return nil, fmt.Errorf("annotation %s contains unknown feature %s", specutils.AnnotationCPUFeatures, str)
+			}
+			allowedFeatures[f] = struct{}{}
+		}
+		afs, err := cpufs.Intersect(allowedFeatures) // err only needed for now because this isn't implemented on ARM64 yet
+		if err != nil {
+			return nil, err
+		}
+		cpufs = afs
+	}
+
 	maxFDLimit := kernel.MaxFdLimit
 	if args.Spec.Linux != nil && args.Spec.Linux.Sysctl != nil {
 		if val, ok := args.Spec.Linux.Sysctl["fs.nr_open"]; ok {
@@ -616,7 +637,7 @@ func New(args Args) (*Loader, error) {
 		DisconnectOnSave: args.Conf.NetDisconnectOk,
 	}
 	if err = l.k.Init(kernel.InitKernelArgs{
-		FeatureSet:           cpuid.HostFeatureSet().Fixed(),
+		FeatureSet:           cpufs.Fixed(),
 		Timekeeper:           tk,
 		RootUserNamespace:    creds.UserNamespace,
 		RootNetworkNamespace: netns,

runsc/cli/main.go

@@ -281,6 +281,7 @@ func forEachCmd(cb func(cmd subcommands.Command, group string)) {
 	cb(new(cmd.Uninstall), helperGroup)
 	cb(new(nvproxy.Nvproxy), helperGroup)
 	cb(new(trace.Trace), helperGroup)
+	cb(new(cmd.CPUFeatures), helperGroup)
 
 	const debugGroup = "debug"
 	cb(new(cmd.Debug), debugGroup)

runsc/cmd/BUILD

@@ -38,6 +38,7 @@ go_library(
         "checkpoint.go",
         "chroot.go",
         "cmd.go",
+        "cpu_features.go",
         "create.go",
         "debug.go",
         "delete.go",

runsc/cmd/cpu_features.go

@@ -0,0 +1,63 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/google/subcommands"
+	"gvisor.dev/gvisor/pkg/cpuid"
+	"gvisor.dev/gvisor/runsc/flag"
+)
+
+// CPUFeatures implements subcommands.Command for the "cpu-features" command.
+type CPUFeatures struct{}
+
+// Name implements subcommands.command.name.
+func (*CPUFeatures) Name() string {
+	return "cpu-features"
+}
+
+// Synopsis implements subcommands.Command.Synopsis.
+func (*CPUFeatures) Synopsis() string {
+	return "list CPU features supported on current machine"
+}
+
+// Usage implements subcommands.Command.Usage.
+func (*CPUFeatures) Usage() string {
+	return "cpu-features\n"
+}
+
+// SetFlags implements subcommands.Command.SetFlags.
+func (*CPUFeatures) SetFlags(*flag.FlagSet) {}
+
+// Execute implements subcommands.Command.Execute.
+func (*CPUFeatures) Execute(_ context.Context, _ *flag.FlagSet, args ...any) subcommands.ExitStatus {
+	cpuid.Initialize()
+	hfs := cpuid.HostFeatureSet().Fixed()
+	allFeatures := cpuid.AllFeatures()
+
+	features := []string{}
+	for _, v := range allFeatures {
+		if hfs.HasFeature(v) {
+			features = append(features, v.String())
+		}
+	}
+	fmt.Println(strings.Join(features, ","))
+
+	return subcommands.ExitSuccess
+}

runsc/specutils/specutils.go

@@ -70,6 +70,10 @@ const (
 
 	// AnnotationTPU is the annotation used to enable TPU proxy on a pod.
 	AnnotationTPU = "dev.gvisor.internal.tpuproxy"
+
+	// AnnotationCPUFeatures is the annotation used to control cpu features
+	// that exposed to user apps.
+	AnnotationCPUFeatures = "dev.gvisor.internal.cpufeatures"
 )
 
 // ExePath must point to runsc binary, which is normally the same binary. It's