[SMT] Refactor runsc mititgate

Refactor mitigate to use /sys/devices/system/cpu/smt/control instead
of individual CPU control files.

PiperOrigin-RevId: 389215975

8 files changed, 306 insertions, 642 deletions

diff --git a/runsc/cmd/BUILD b/runsc/cmd/BUILD
index 39c8ff603..031ddd57e 100644
--- a/runsc/cmd/BUILD
+++ b/runsc/cmd/BUILD
@@ -95,10 +95,10 @@ go_test(
         "//runsc/config",
         "//runsc/container",
         "//runsc/mitigate",
-        "//runsc/mitigate/mock",
         "//runsc/specutils",
         "@com_github_google_go_cmp//cmp:go_default_library",
         "@com_github_google_go_cmp//cmp/cmpopts:go_default_library",
+        "@com_github_google_subcommands//:go_default_library",
         "@com_github_opencontainers_runtime_spec//specs-go:go_default_library",
         "@com_github_syndtr_gocapability//capability:go_default_library",
     ],
diff --git a/runsc/cmd/mitigate.go b/runsc/cmd/mitigate.go
index f4e65adb8..1aada5968 100644
--- a/runsc/cmd/mitigate.go
+++ b/runsc/cmd/mitigate.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
+	"os"
 	"runtime"
 
 	"github.com/google/subcommands"
@@ -29,8 +30,8 @@ import (
 const (
 	// cpuInfo is the path used to parse CPU info.
 	cpuInfo = "/proc/cpuinfo"
-	// allPossibleCPUs is the path used to enable CPUs.
-	allPossibleCPUs = "/sys/devices/system/cpu/possible"
+	// Path to enable/disable SMT.
+	smtPath = "/sys/devices/system/cpu/smt/control"
 )
 
 // Mitigate implements subcommands.Command for the "mitigate" command.
@@ -39,10 +40,10 @@ type Mitigate struct {
 	dryRun bool
 	// Reverse mitigate by turning on all CPU cores.
 	reverse bool
-	// Path to file to read to create CPUSet.
-	path string
 	// Extra data for post mitigate operations.
 	data string
+	// Control to mitigate/reverse smt.
+	control machineControl
 }
 
 // Name implements subcommands.command.name.
@@ -56,12 +57,12 @@ func (*Mitigate) Synopsis() string {
 }
 
 // Usage implements Usage for cmd.Mitigate.
-func (m Mitigate) Usage() string {
+func (m *Mitigate) Usage() string {
 	return fmt.Sprintf(`mitigate [flags]
 
-mitigate mitigates a system to the "MDS" vulnerability by implementing a manual shutdown of SMT. The command checks /proc/cpuinfo for cpus having the MDS vulnerability, and if found, shutdown all but one CPU per hyperthread pair via /sys/devices/system/cpu/cpu{N}/online. CPUs can be restored by writing "2" to each file in /sys/devices/system/cpu/cpu{N}/online or performing a system reboot.
+mitigate mitigates a system to the "MDS" vulnerability by writing "off" to %q. CPUs can be restored by writing "on" to the same file or rebooting your system.
 
-The command can be reversed with --reverse, which reads the total CPUs from /sys/devices/system/cpu/possible and enables all with /sys/devices/system/cpu/cpu{N}/online.%s`, m.usage())
+The command can be reversed with --reverse, which writes "on" to the file above.%s`, smtPath, m.usage())
 }
 
 // SetFlags sets flags for the command Mitigate.
@@ -74,104 +75,110 @@ func (m *Mitigate) SetFlags(f *flag.FlagSet) {
 // Execute implements subcommands.Command.Execute.
 func (m *Mitigate) Execute(_ context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
 	if runtime.GOARCH == "arm64" || runtime.GOARCH == "arm" {
-		log.Warningf("As ARM is not affected by MDS, mitigate does not support")
-		return subcommands.ExitFailure
+		log.Warningf("As ARM is not affected by MDS, mitigate does not support ARM machines.")
+		// Set reverse flag so that we still perform post mitigate operations. mitigate reverse is a noop in this case.
+		m.reverse = true
 	}
 
 	if f.NArg() != 0 {
 		f.Usage()
 		return subcommands.ExitUsageError
 	}
+	m.control = &machineControlImpl{}
+	return m.execute()
+}
 
-	m.path = cpuInfo
-	if m.reverse {
-		m.path = allPossibleCPUs
+// execute executes mitigate operations. Seperate from Execute method for
+// easier mocking.
+func (m *Mitigate) execute() subcommands.ExitStatus {
+	beforeSet, err := m.control.getCPUs()
+	if err != nil {
+		return Errorf("Get before CPUSet failed: %v", err)
 	}
+	log.Infof("CPUs before: %s", beforeSet.String())
 
-	set, err := m.doExecute()
-	if err != nil {
-		return Errorf("Execute failed: %v", err)
+	if err := m.doEnableDisable(beforeSet); err != nil {
+		return Errorf("Enabled/Disable action failed on %q: %v", smtPath, err)
 	}
 
-	if m.data == "" {
-		return subcommands.ExitSuccess
+	afterSet, err := m.control.getCPUs()
+	if err != nil {
+		return Errorf("Get after CPUSet failed: %v", err)
 	}
+	log.Infof("CPUs after: %s", afterSet.String())
 
-	if err = m.postMitigate(set); err != nil {
+	if err = m.postMitigate(afterSet); err != nil {
 		return Errorf("Post Mitigate failed: %v", err)
 	}
 
 	return subcommands.ExitSuccess
 }
 
-// Execute executes the Mitigate command.
-func (m *Mitigate) doExecute() (mitigate.CPUSet, error) {
-	if m.dryRun {
-		log.Infof("Running with DryRun. No cpu settings will be changed.")
-	}
-	data, err := ioutil.ReadFile(m.path)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read %s: %w", m.path, err)
-	}
+// doEnableDisable does either enable or disable operation based on flags.
+func (m *Mitigate) doEnableDisable(set mitigate.CPUSet) error {
 	if m.reverse {
-		set, err := m.doReverse(data)
-		if err != nil {
-			return nil, fmt.Errorf("reverse operation failed: %w", err)
+		if m.dryRun {
+			log.Infof("Skipping reverse action because dryrun is set.")
+			return nil
 		}
-		return set, nil
+		return m.control.enable()
 	}
-	set, err := m.doMitigate(data)
-	if err != nil {
-		return nil, fmt.Errorf("mitigate operation failed: %w", err)
+	if m.dryRun {
+		log.Infof("Skipping mitigate action because dryrun is set.")
+		return nil
 	}
-	return set, nil
+	if set.IsVulnerable() {
+		return m.control.disable()
+	}
+	log.Infof("CPUs not vulnerable. Skipping disable call.")
+	return nil
 }
 
-func (m *Mitigate) doMitigate(data []byte) (mitigate.CPUSet, error) {
-	set, err := mitigate.NewCPUSet(data)
-	if err != nil {
-		return nil, err
-	}
+// Interface to wrap interactions with underlying machine. Done
+// so testing with mocks can be done hermetically.
+type machineControl interface {
+	enable() error
+	disable() error
+	isEnabled() (bool, error)
+	getCPUs() (mitigate.CPUSet, error)
+}
 
-	log.Infof("Mitigate found the following CPUs...")
-	log.Infof("%s", set)
+// Implementation of SMT control interaction with the underlying machine.
+type machineControlImpl struct{}
 
-	disableList := set.GetShutdownList()
-	log.Infof("Disabling threads on thread pairs.")
-	for _, t := range disableList {
-		log.Infof("Disable thread: %s", t)
-		if m.dryRun {
-			continue
-		}
-		if err := t.Disable(); err != nil {
-			return nil, fmt.Errorf("error disabling thread: %s err: %w", t, err)
-		}
-	}
-	log.Infof("Shutdown successful.")
-	return set, nil
+func (*machineControlImpl) enable() error {
+	return checkFileExistsOnWrite("enable", "on")
 }
 
-func (m *Mitigate) doReverse(data []byte) (mitigate.CPUSet, error) {
-	set, err := mitigate.NewCPUSetFromPossible(data)
-	if err != nil {
-		return nil, err
-	}
+func (*machineControlImpl) disable() error {
+	return checkFileExistsOnWrite("disable", "off")
+}
 
-	log.Infof("Reverse mitigate found the following CPUs...")
-	log.Infof("%s", set)
+// Writes data to SMT control. If file not found, logs file not exist error and returns nil
+// error, which is done because machines without the file pointed to by smtPath only have one
+// thread per core in the first place. Otherwise returns error from ioutil.WriteFile.
+func checkFileExistsOnWrite(op, data string) error {
+	err := ioutil.WriteFile(smtPath, []byte(data), 0644)
+	if err != nil && os.IsExist(err) {
+		log.Infof("File %q does not exist for operation %s. This machine probably has no smt control.", smtPath, op)
+		return nil
+	}
+	return err
+}
 
-	enableList := set.GetRemainingList()
+func (*machineControlImpl) isEnabled() (bool, error) {
+	data, err := ioutil.ReadFile(cpuInfo)
+	return string(data) == "on", err
+}
 
-	log.Infof("Enabling all CPUs...")
-	for _, t := range enableList {
-		log.Infof("Enabling thread: %s", t)
-		if m.dryRun {
-			continue
-		}
-		if err := t.Enable(); err != nil {
-			return nil, fmt.Errorf("error enabling thread: %s err: %w", t, err)
-		}
+func (*machineControlImpl) getCPUs() (mitigate.CPUSet, error) {
+	data, err := ioutil.ReadFile(cpuInfo)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read %s: %w", cpuInfo, err)
+	}
+	set, err := mitigate.NewCPUSet(string(data))
+	if err != nil {
+		return nil, fmt.Errorf("getCPUs: %v", err)
 	}
-	log.Infof("Enable successful.")
 	return set, nil
 }
diff --git a/runsc/cmd/mitigate_test.go b/runsc/cmd/mitigate_test.go
index 51755d9f3..294fc645c 100644
--- a/runsc/cmd/mitigate_test.go
+++ b/runsc/cmd/mitigate_test.go
@@ -18,144 +18,133 @@
 package cmd
 
 import (
-	"fmt"
-	"io/ioutil"
-	"os"
-	"strings"
 	"testing"
 
-	"gvisor.dev/gvisor/runsc/mitigate/mock"
+	"github.com/google/subcommands"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/runsc/mitigate"
 )
 
-type executeTestCase struct {
-	name          string
-	mitigateData  string
-	mitigateError error
-	mitigateCPU   int
-	reverseData   string
-	reverseError  error
-	reverseCPU    int
+type mockMachineControl struct {
+	enabled bool
+	cpus    mitigate.CPUSet
 }
 
-func TestExecute(t *testing.T) {
+func (m *mockMachineControl) enable() error {
+	m.enabled = true
+	return nil
+}
 
-	partial := `processor       : 1
-vendor_id       : AuthenticAMD
-cpu family      : 23
-model           : 49
-model name      : AMD EPYC 7B12
-physical id     : 0
-bugs         : sysret_ss_attrs spectre_v1 spectre_v2 spec_store_bypass
-power management:
-`
+func (m *mockMachineControl) disable() error {
+	if m.cpus.IsVulnerable() {
+		m.enabled = false
+	}
+	return nil
+}
 
+func (m *mockMachineControl) isEnabled() (bool, error) {
+	return m.enabled, nil
+}
+
+func (m *mockMachineControl) getCPUs() (mitigate.CPUSet, error) {
+	set := m.cpus
+	if !m.enabled {
+		set = m.cpus[:len(m.cpus)/2]
+	}
+
+	// Instead of just returning the created CPU set stored in this struct, call
+	// NewCPUSet to exercise that code path as the machineControlImpl would.
+	return mitigate.NewCPUSet(set.String())
+}
+
+type executeTestCase struct {
+	name                string
+	cpu                 mitigate.MockCPU
+	mitigateWantCPUs    int
+	mitigateError       subcommands.ExitStatus
+	mitigateWantEnabled bool
+	reverseWantCPUs     int
+	reverseError        subcommands.ExitStatus
+	reverseWantEnabled  bool
+	dryrun              bool
+}
+
+func TestExecute(t *testing.T) {
 	for _, tc := range []executeTestCase{
 		{
-			name:         "CascadeLake4",
-			mitigateData: mock.CascadeLake4.MakeCPUString(),
-			mitigateCPU:  2,
-			reverseData:  mock.CascadeLake4.MakeSysPossibleString(),
-			reverseCPU:   4,
+			name:                "CascadeLake4",
+			cpu:                 mitigate.CascadeLake4,
+			mitigateWantCPUs:    2,
+			mitigateWantEnabled: false,
+			reverseWantCPUs:     4,
+			reverseWantEnabled:  true,
 		},
 		{
-			name:          "Empty",
-			mitigateData:  "",
-			mitigateError: fmt.Errorf(`mitigate operation failed: no cpus found for: ""`),
-			reverseData:   "",
-			reverseError:  fmt.Errorf(`reverse operation failed: mismatch regex from possible: ""`),
+			name:                "CascadeLake4DryRun",
+			cpu:                 mitigate.CascadeLake4,
+			mitigateWantCPUs:    4,
+			mitigateWantEnabled: true,
+			reverseWantCPUs:     4,
+			reverseWantEnabled:  true,
+			dryrun:              true,
 		},
 		{
-			name: "Partial",
-			mitigateData: `processor       : 0
-vendor_id       : AuthenticAMD
-cpu family      : 23
-model           : 49
-model name      : AMD EPYC 7B12
-physical id     : 0
-core id         : 0
-cpu cores       : 1
-bugs            : sysret_ss_attrs spectre_v1 spectre_v2 spec_store_bypass
-power management::84
-
-` + partial,
-			mitigateError: fmt.Errorf(`mitigate operation failed: failed to match key "core id": %q`, partial),
-			reverseData:   "1-",
-			reverseError:  fmt.Errorf(`reverse operation failed: mismatch regex from possible: %q`, "1-"),
+			name:                "AMD8",
+			cpu:                 mitigate.AMD8,
+			mitigateWantCPUs:    8,
+			mitigateWantEnabled: true,
+			reverseWantCPUs:     8,
+			reverseWantEnabled:  true,
+		},
+		{
+			name:          "Empty",
+			cpu:           mitigate.Empty,
+			mitigateError: Errorf(`mitigate operation failed: no cpus found for: ""`),
+			reverseError:  Errorf(`mitigate operation failed: no cpus found for: ""`),
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
+			set := tc.cpu.MakeCPUSet()
 			m := &Mitigate{
-				dryRun: true,
+				control: &mockMachineControl{
+					enabled: true,
+					cpus:    set,
+				},
+				dryRun: tc.dryrun,
 			}
-			m.doExecuteTest(t, "Mitigate", tc.mitigateData, tc.mitigateCPU, tc.mitigateError)
+			t.Run("Mitigate", func(t *testing.T) {
+				m.doExecuteTest(t, tc.mitigateWantEnabled, tc.mitigateWantCPUs, tc.mitigateError)
+			})
 
 			m.reverse = true
-			m.doExecuteTest(t, "Reverse", tc.reverseData, tc.reverseCPU, tc.reverseError)
+			t.Run("Reverse", func(t *testing.T) {
+				m.doExecuteTest(t, tc.reverseWantEnabled, tc.reverseWantCPUs, tc.reverseError)
+			})
 		})
 	}
 }
 
-func TestExecuteSmoke(t *testing.T) {
-	smokeMitigate, err := ioutil.ReadFile(cpuInfo)
-	if err != nil {
-		t.Fatalf("Failed to read %s: %v", cpuInfo, err)
+// doExecuteTest runs Execute with the mitigate operation and reverse operation.
+func (m *Mitigate) doExecuteTest(t *testing.T, wantEnabled bool, wantCPUs int, wantErr subcommands.ExitStatus) {
+	subError := m.execute()
+	if subError != wantErr {
+		t.Fatalf("Mitigate error mismatch: want: %v got: %v", wantErr, subError)
 	}
 
-	m := &Mitigate{
-		dryRun: true,
+	// case where test should end in error or we don't care
+	// about how many cpus are returned.
+	if wantErr != subcommands.ExitSuccess {
+		log.Infof("return")
+		return
 	}
 
-	m.doExecuteTest(t, "Mitigate", string(smokeMitigate), 0, nil)
-
-	smokeReverse, err := ioutil.ReadFile(allPossibleCPUs)
-	if err != nil {
-		t.Fatalf("Failed to read %s: %v", allPossibleCPUs, err)
+	gotEnabled, _ := m.control.isEnabled()
+	if wantEnabled != gotEnabled {
+		t.Fatalf("Incorrect enabled state: want: %t got: %t", wantEnabled, gotEnabled)
 	}
 
-	m.reverse = true
-	m.doExecuteTest(t, "Reverse", string(smokeReverse), 0, nil)
-}
-
-// doExecuteTest runs Execute with the mitigate operation and reverse operation.
-func (m *Mitigate) doExecuteTest(t *testing.T, name, data string, want int, wantErr error) {
-	t.Run(name, func(t *testing.T) {
-		file, err := ioutil.TempFile("", "outfile.txt")
-		if err != nil {
-			t.Fatalf("Failed to create tmpfile: %v", err)
-		}
-		defer os.Remove(file.Name())
-
-		if _, err := file.WriteString(data); err != nil {
-			t.Fatalf("Failed to write to file: %v", err)
-		}
-
-		// Set fields for mitigate and dryrun to keep test hermetic.
-		m.path = file.Name()
-
-		set, err := m.doExecute()
-		if err = checkErr(wantErr, err); err != nil {
-			t.Fatalf("Mitigate error mismatch: %v", err)
-		}
-
-		// case where test should end in error or we don't care
-		// about how many cpus are returned.
-		if wantErr != nil || want < 1 {
-			return
-		}
-		got := len(set.GetRemainingList())
-		if want != got {
-			t.Fatalf("Failed wrong number of remaining CPUs: want %d, got %d", want, got)
-		}
-
-	})
-}
-
-// checkErr checks error for equality.
-func checkErr(want, got error) error {
-	switch {
-	case want == nil && got == nil:
-	case want == nil || got == nil || want.Error() != strings.Trim(got.Error(), " "):
-		return fmt.Errorf("got: %v want: %v", got, want)
+	gotCPUs, _ := m.control.getCPUs()
+	if len(gotCPUs) != wantCPUs {
+		t.Fatalf("Incorrect number of CPUs: want: %d got: %d", wantCPUs, len(gotCPUs))
 	}
-	return nil
 }
diff --git a/runsc/mitigate/BUILD b/runsc/mitigate/BUILD
index 1238890fc..9f4a7ba8d 100644
--- a/runsc/mitigate/BUILD
+++ b/runsc/mitigate/BUILD
@@ -4,7 +4,10 @@ package(licenses = ["notice"])
 
 go_library(
     name = "mitigate",
-    srcs = ["mitigate.go"],
+    srcs = [
+        "mitigate.go",
+        "mock.go",
+    ],
     visibility = [
         "//runsc:__subpackages__",
     ],
@@ -16,8 +19,5 @@ go_test(
     size = "small",
     srcs = ["mitigate_test.go"],
     library = ":mitigate",
-    deps = [
-        "//runsc/mitigate/mock",
-        "@com_github_google_go_cmp//cmp:go_default_library",
-    ],
+    deps = ["@com_github_google_go_cmp//cmp:go_default_library"],
 )
diff --git a/runsc/mitigate/mitigate.go b/runsc/mitigate/mitigate.go
index 9f29ec873..00e5bf2a9 100644
--- a/runsc/mitigate/mitigate.go
+++ b/runsc/mitigate/mitigate.go
@@ -19,10 +19,7 @@ package mitigate
 
 import (
 	"fmt"
-	"io/ioutil"
-	"os"
 	"regexp"
-	"sort"
 	"strconv"
 	"strings"
 )
@@ -39,128 +36,20 @@ const (
 	physicalIDKey = "physical id"
 	coreIDKey     = "core id"
 	bugsKey       = "bugs"
-
-	// Path to shutdown a CPU.
-	cpuOnlineTemplate = "/sys/devices/system/cpu/cpu%d/online"
 )
 
 // CPUSet contains a map of all CPUs on the system, mapped
 // by Physical ID and CoreIDs. threads with the same
 // Core and Physical ID are Hyperthread pairs.
-type CPUSet map[threadID]*ThreadGroup
+type CPUSet []*CPU
 
 // NewCPUSet creates a CPUSet from data read from /proc/cpuinfo.
-func NewCPUSet(data []byte) (CPUSet, error) {
-	processors, err := getThreads(string(data))
-	if err != nil {
-		return nil, err
-	}
-
-	set := make(CPUSet)
-	for _, p := range processors {
-		// Each ID is of the form physicalID:coreID. Hyperthread pairs
-		// have identical physical and core IDs. We need to match
-		// Hyperthread pairs so that we can shutdown all but one per
-		// pair.
-		core, ok := set[p.id]
-		if !ok {
-			core = &ThreadGroup{}
-			set[p.id] = core
-		}
-		core.isVulnerable = core.isVulnerable || p.IsVulnerable()
-		core.threads = append(core.threads, p)
-	}
-
-	// We need to make sure we shutdown the lowest number processor per
-	// thread group.
-	for _, tg := range set {
-		sort.Slice(tg.threads, func(i, j int) bool {
-			return tg.threads[i].processorNumber < tg.threads[j].processorNumber
-		})
-	}
-	return set, nil
-}
-
-// NewCPUSetFromPossible makes a cpuSet data read from
-// /sys/devices/system/cpu/possible. This is used in enable operations
-// where the caller simply wants to enable all CPUS.
-func NewCPUSetFromPossible(data []byte) (CPUSet, error) {
-	threads, err := GetThreadsFromPossible(data)
-	if err != nil {
-		return nil, err
-	}
-
-	// We don't care if a CPU is vulnerable or not, we just
-	// want to return a list of all CPUs on the host.
-	set := CPUSet{
-		threads[0].id: &ThreadGroup{
-			threads:      threads,
-			isVulnerable: false,
-		},
-	}
-	return set, nil
-}
-
-// String implements the String method for CPUSet.
-func (c CPUSet) String() string {
-	ret := ""
-	for _, tg := range c {
-		ret += fmt.Sprintf("%s\n", tg)
-	}
-	return ret
-}
-
-// GetRemainingList returns the list of threads that will remain active
-// after mitigation.
-func (c CPUSet) GetRemainingList() []Thread {
-	threads := make([]Thread, 0, len(c))
-	for _, core := range c {
-		// If we're vulnerable, take only one thread from the pair.
-		if core.isVulnerable {
-			threads = append(threads, core.threads[0])
-			continue
-		}
-		// Otherwise don't shutdown anything.
-		threads = append(threads, core.threads...)
-	}
-	return threads
-}
-
-// GetShutdownList returns the list of threads that will be shutdown on
-// mitigation.
-func (c CPUSet) GetShutdownList() []Thread {
-	threads := make([]Thread, 0)
-	for _, core := range c {
-		// Only if we're vulnerable do shutdown anything. In this case,
-		// shutdown all but the first entry.
-		if core.isVulnerable && len(core.threads) > 1 {
-			threads = append(threads, core.threads[1:]...)
-		}
-	}
-	return threads
-}
-
-// ThreadGroup represents Hyperthread pairs on the same physical/core ID.
-type ThreadGroup struct {
-	threads      []Thread
-	isVulnerable bool
-}
-
-// String implements the String method for threadGroup.
-func (c ThreadGroup) String() string {
-	ret := fmt.Sprintf("ThreadGroup:\nIsVulnerable: %t\n", c.isVulnerable)
-	for _, processor := range c.threads {
-		ret += fmt.Sprintf("%s\n", processor)
-	}
-	return ret
-}
-
-// getThreads returns threads structs from reading /proc/cpuinfo.
-func getThreads(data string) ([]Thread, error) {
+func NewCPUSet(data string) (CPUSet, error) {
 	// Each processor entry should start with the
 	// processor key. Find the beginings of each.
 	r := buildRegex(processorKey)
 	indices := r.FindAllStringIndex(data, -1)
+
 	if len(indices) < 1 {
 		return nil, fmt.Errorf("no cpus found for: %q", data)
 	}
@@ -172,193 +61,132 @@ func getThreads(data string) ([]Thread, error) {
 	// indexes (e.g. data[index[i], index[i+1]]).
 	// There should be len(indicies) - 1 CPUs
 	// since the last index is the end of the string.
-	cpus := make([]Thread, 0, len(indices))
+	var set CPUSet
 	// Find each string that represents a CPU. These begin "processor".
 	for i := 1; i < len(indices); i++ {
 		start := indices[i-1][0]
 		end := indices[i][0]
 		// Parse the CPU entry, which should be between start/end.
-		c, err := newThread(data[start:end])
+		c, err := newCPU(data[start:end])
 		if err != nil {
 			return nil, err
 		}
-		cpus = append(cpus, c)
+		set = append(set, c)
 	}
-	return cpus, nil
+	return set, nil
 }
 
-// GetThreadsFromPossible makes threads from data read from /sys/devices/system/cpu/possible.
-func GetThreadsFromPossible(data []byte) ([]Thread, error) {
-	possibleRegex := regexp.MustCompile(`(?m)^(\d+)(-(\d+))?$`)
-	matches := possibleRegex.FindStringSubmatch(string(data))
-	if len(matches) != 4 {
-		return nil, fmt.Errorf("mismatch regex from possible: %q", string(data))
-	}
-
-	// If matches[3] is empty, we only have one cpu entry.
-	if matches[3] == "" {
-		matches[3] = matches[1]
-	}
-
-	begin, err := strconv.ParseInt(matches[1], 10, 64)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse begin: %v", err)
-	}
-	end, err := strconv.ParseInt(matches[3], 10, 64)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse end: %v", err)
-	}
-	if begin > end || begin < 0 || end < 0 {
-		return nil, fmt.Errorf("invalid cpu bounds from possible: begin: %d end: %d", begin, end)
-	}
-
-	ret := make([]Thread, 0, end-begin)
-	for i := begin; i <= end; i++ {
-		ret = append(ret, Thread{
-			processorNumber: i,
-			id: threadID{
-				physicalID: 0, // we don't care about id for enable ops.
-				coreID:     0,
-			},
-		})
+// IsVulnerable checks if this CPUSet is vulnerable to MDS.
+func (c CPUSet) IsVulnerable() bool {
+	for _, cpu := range c {
+		if cpu.IsVulnerable() {
+			return true
+		}
 	}
-
-	return ret, nil
+	return false
 }
 
-// threadID for each thread is defined by the physical and
-// core IDs. If equal, two threads are Hyperthread pairs.
-type threadID struct {
-	physicalID int64
-	coreID     int64
+// String implements the String method for CPUSet.
+func (c CPUSet) String() string {
+	parts := make([]string, len(c))
+	for i, cpu := range c {
+		parts[i] = cpu.String()
+	}
+	return strings.Join(parts, "\n")
 }
 
-// Thread represents pertinent info about a single hyperthread in a pair.
-type Thread struct {
+// CPU represents pertinent info about a single hyperthread in a pair.
+type CPU struct {
 	processorNumber int64               // the processor number of this CPU.
 	vendorID        string              // the vendorID of CPU (e.g. AuthenticAMD).
 	cpuFamily       int64               // CPU family number (e.g. 6 for CascadeLake/Skylake).
 	model           int64               // CPU model number (e.g. 85 for CascadeLake/Skylake).
-	id              threadID            // id for this thread
+	physicalID      int64               // Physical ID of this CPU.
+	coreID          int64               // Core ID of this CPU.
 	bugs            map[string]struct{} // map of vulnerabilities parsed from the 'bugs' field.
 }
 
-// newThread parses a CPU from a single cpu entry from /proc/cpuinfo.
-func newThread(data string) (Thread, error) {
-	empty := Thread{}
+func newCPU(data string) (*CPU, error) {
 	processor, err := parseProcessor(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	vendorID, err := parseVendorID(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	cpuFamily, err := parseCPUFamily(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	model, err := parseModel(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	physicalID, err := parsePhysicalID(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	coreID, err := parseCoreID(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
 	bugs, err := parseBugs(data)
 	if err != nil {
-		return empty, err
+		return nil, err
 	}
 
-	return Thread{
+	return &CPU{
 		processorNumber: processor,
 		vendorID:        vendorID,
 		cpuFamily:       cpuFamily,
 		model:           model,
-		id: threadID{
-			physicalID: physicalID,
-			coreID:     coreID,
-		},
-		bugs: bugs,
+		physicalID:      physicalID,
+		coreID:          coreID,
+		bugs:            bugs,
 	}, nil
 }
 
-// String implements the String method for thread.
-func (t Thread) String() string {
-	template := `CPU: %d
-CPU ID: %+v
-Vendor: %s
-Family/Model: %d/%d
-Bugs: %s
+// String implements the String method for CPU.
+func (t *CPU) String() string {
+	template := `%s: %d
+%s: %s
+%s: %d
+%s: %d
+%s: %d
+%s: %d
+%s: %s
 `
-	bugs := make([]string, 0)
+	var bugs []string
 	for bug := range t.bugs {
 		bugs = append(bugs, bug)
 	}
 
-	return fmt.Sprintf(template, t.processorNumber, t.id, t.vendorID, t.cpuFamily, t.model, strings.Join(bugs, ","))
-}
-
-// Enable turns on the CPU by writing 1 to /sys/devices/cpu/cpu{N}/online.
-func (t Thread) Enable() error {
-	// Linux ensures that "cpu0" is always online.
-	if t.processorNumber == 0 {
-		return nil
-	}
-	cpuPath := fmt.Sprintf(cpuOnlineTemplate, t.processorNumber)
-	f, err := os.OpenFile(cpuPath, os.O_WRONLY|os.O_CREATE, 0644)
-	if err != nil {
-		return fmt.Errorf("failed to open file %s: %v", cpuPath, err)
-	}
-	if _, err = f.Write([]byte{'1'}); err != nil {
-		return fmt.Errorf("failed to write '1' to %s: %v", cpuPath, err)
-	}
-	return nil
-}
-
-// Disable turns off the CPU by writing 0 to /sys/devices/cpu/cpu{N}/online.
-func (t Thread) Disable() error {
-	// The core labeled "cpu0" can never be taken offline via this method.
-	// Linux will return EPERM if the user even creates a file at the /sys
-	// path above.
-	if t.processorNumber == 0 {
-		return fmt.Errorf("invalid shutdown operation: cpu0 cannot be disabled")
-	}
-	cpuPath := fmt.Sprintf(cpuOnlineTemplate, t.processorNumber)
-	return ioutil.WriteFile(cpuPath, []byte{'0'}, 0644)
+	return fmt.Sprintf(template,
+		processorKey, t.processorNumber,
+		vendorIDKey, t.vendorID,
+		cpuFamilyKey, t.cpuFamily,
+		modelKey, t.model,
+		physicalIDKey, t.physicalID,
+		coreIDKey, t.coreID,
+		bugsKey, strings.Join(bugs, " "))
 }
 
 // IsVulnerable checks if a CPU is vulnerable to mds.
-func (t Thread) IsVulnerable() bool {
+func (t *CPU) IsVulnerable() bool {
 	_, ok := t.bugs[mds]
 	return ok
 }
 
-// isActive checks if a CPU is active from /sys/devices/system/cpu/cpu{N}/online
-// If the file does not exist (ioutil returns in error), we assume the CPU is on.
-func (t Thread) isActive() bool {
-	cpuPath := fmt.Sprintf(cpuOnlineTemplate, t.processorNumber)
-	data, err := ioutil.ReadFile(cpuPath)
-	if err != nil {
-		return true
-	}
-	return len(data) > 0 && data[0] != '0'
-}
-
 // SimilarTo checks family/model/bugs fields for equality of two
 // processors.
-func (t Thread) SimilarTo(other Thread) bool {
+func (t *CPU) SimilarTo(other *CPU) bool {
 	if t.vendorID != other.vendorID {
 		return false
 	}
diff --git a/runsc/mitigate/mitigate_test.go b/runsc/mitigate/mitigate_test.go
index a1d80581e..e79d879e9 100644
--- a/runsc/mitigate/mitigate_test.go
+++ b/runsc/mitigate/mitigate_test.go
@@ -18,90 +18,53 @@
 package mitigate
 
 import (
-	"fmt"
 	"io/ioutil"
 	"strings"
 	"testing"
-
-	"gvisor.dev/gvisor/runsc/mitigate/mock"
 )
 
 // TestMockCPUSet tests mock cpu test cases against the cpuSet functions.
 func TestMockCPUSet(t *testing.T) {
 	for _, tc := range []struct {
-		testCase     mock.CPU
+		testCase     MockCPU
 		isVulnerable bool
 	}{
 		{
-			testCase:     mock.AMD8,
+			testCase:     AMD8,
 			isVulnerable: false,
 		},
 		{
-			testCase:     mock.Haswell2,
+			testCase:     Haswell2,
 			isVulnerable: true,
 		},
 		{
-			testCase:     mock.Haswell2core,
+			testCase:     Haswell2core,
 			isVulnerable: true,
 		},
 		{
-			testCase:     mock.CascadeLake2,
+			testCase:     CascadeLake2,
 			isVulnerable: true,
 		},
 		{
-			testCase:     mock.CascadeLake4,
+			testCase:     CascadeLake4,
 			isVulnerable: true,
 		},
 	} {
 		t.Run(tc.testCase.Name, func(t *testing.T) {
-			data := tc.testCase.MakeCPUString()
-			set, err := NewCPUSet([]byte(data))
+			data := tc.testCase.MakeCPUSet().String()
+			set, err := NewCPUSet(data)
 			if err != nil {
 				t.Fatalf("Failed to create cpuSet: %v", err)
 			}
 
-			t.Logf("data: %s", data)
-
-			for _, tg := range set {
-				if err := checkSorted(tg.threads); err != nil {
-					t.Fatalf("Failed to sort cpuSet: %v", err)
-				}
-			}
-
-			remaining := set.GetRemainingList()
-			// In the non-vulnerable case, no cores should be shutdown so all should remain.
-			want := tc.testCase.PhysicalCores * tc.testCase.Cores * tc.testCase.ThreadsPerCore
-			if tc.isVulnerable {
-				want = tc.testCase.PhysicalCores * tc.testCase.Cores
-			}
-
-			if want != len(remaining) {
-				t.Fatalf("Failed to shutdown the correct number of cores: want: %d got: %d", want, len(remaining))
-			}
-
-			if !tc.isVulnerable {
-				return
-			}
-
-			// If the set is vulnerable, we expect only 1 thread per hyperthread pair.
-			for _, r := range remaining {
-				if _, ok := set[r.id]; !ok {
-					t.Fatalf("Entry %+v not in map, there must be two entries in the same thread group.", r)
-				}
-				delete(set, r.id)
-			}
-
-			possible := tc.testCase.MakeSysPossibleString()
-			set, err = NewCPUSetFromPossible([]byte(possible))
-			if err != nil {
-				t.Fatalf("Failed to make cpuSet: %v", err)
+			if tc.testCase.NumCPUs() != len(set) {
+				t.Fatalf("Got wrong number of CPUs: want: %d got: %d", tc.testCase.NumCPUs(), len(set))
 			}
 
-			want = tc.testCase.PhysicalCores * tc.testCase.Cores * tc.testCase.ThreadsPerCore
-			got := len(set.GetRemainingList())
-			if got != want {
-				t.Fatalf("Returned the wrong number of CPUs want: %d got: %d", want, got)
+			if set.IsVulnerable() != tc.isVulnerable {
+				t.Fatalf("incorrect vulnerable value: got: %t want: %t", set.IsVulnerable(), tc.isVulnerable)
 			}
+			t.Logf("data: %s", data)
 		})
 	}
 }
@@ -117,15 +80,13 @@ physical id: 0
 core id		: 0
 bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa itlb_multihit
 `
-	want := Thread{
+	want := CPU{
 		processorNumber: 0,
 		vendorID:        "GenuineIntel",
 		cpuFamily:       6,
 		model:           85,
-		id: threadID{
-			physicalID: 0,
-			coreID:     0,
-		},
+		physicalID:      0,
+		coreID:          0,
 		bugs: map[string]struct{}{
 			"cpu_meltdown":      {},
 			"spectre_v1":        {},
@@ -139,7 +100,7 @@ bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
 		},
 	}
 
-	got, err := newThread(data)
+	got, err := newCPU(data)
 	if err != nil {
 		t.Fatalf("getCpu failed with error: %v", err)
 	}
@@ -154,12 +115,12 @@ bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
 }
 
 func TestInvalid(t *testing.T) {
-	result, err := getThreads(`something not a processor`)
+	result, err := newCPU(`something not a processor`)
 	if err == nil {
 		t.Fatalf("getCPU set didn't return an error: %+v", result)
 	}
 
-	if !strings.Contains(err.Error(), "no cpus") {
+	if !strings.Contains(err.Error(), "failed to match key \"processor\"") {
 		t.Fatalf("Incorrect error returned: %v", err)
 	}
 }
@@ -221,7 +182,7 @@ cache_alignment	: 64
 address sizes	: 46 bits physical, 48 bits virtual
 power management:
 `
-	cpuSet, err := getThreads(data)
+	cpuSet, err := NewCPUSet(data)
 	if err != nil {
 		t.Fatalf("getCPUSet failed: %v", err)
 	}
@@ -231,7 +192,7 @@ power management:
 		t.Fatalf("Num CPU mismatch: want: %d, got: %d", wantCPULen, len(cpuSet))
 	}
 
-	wantCPU := Thread{
+	wantCPU := CPU{
 		vendorID:  "GenuineIntel",
 		cpuFamily: 6,
 		model:     63,
@@ -260,17 +221,11 @@ func TestReadFile(t *testing.T) {
 		t.Fatalf("Failed to read cpuinfo: %v", err)
 	}
 
-	set, err := NewCPUSet(data)
+	set, err := NewCPUSet(string(data))
 	if err != nil {
 		t.Fatalf("Failed to parse CPU data %v\n%s", err, data)
 	}
 
-	for _, tg := range set {
-		if err := checkSorted(tg.threads); err != nil {
-			t.Fatalf("Failed to sort cpuSet: %v", err)
-		}
-	}
-
 	if len(set) < 1 {
 		t.Fatalf("Failed to parse any CPUs: %d", len(set))
 	}
@@ -383,7 +338,7 @@ power management:`
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			set, err := getThreads(tc.cpuString)
+			set, err := NewCPUSet(tc.cpuString)
 			if err != nil {
 				t.Fatalf("Failed to getCPUSet:%v\n %s", err, tc.cpuString)
 			}
@@ -404,98 +359,3 @@ power management:`
 		})
 	}
 }
-
-func TestReverse(t *testing.T) {
-	const noParse = "-1-"
-	for _, tc := range []struct {
-		name      string
-		output    string
-		wantErr   error
-		wantCount int
-	}{
-		{
-			name:      "base",
-			output:    "0-7",
-			wantErr:   nil,
-			wantCount: 8,
-		},
-		{
-			name:      "huge",
-			output:    "0-111",
-			wantErr:   nil,
-			wantCount: 112,
-		},
-		{
-			name:      "not zero",
-			output:    "50-53",
-			wantErr:   nil,
-			wantCount: 4,
-		},
-		{
-			name:      "small",
-			output:    "0",
-			wantErr:   nil,
-			wantCount: 1,
-		},
-		{
-			name:    "invalid order",
-			output:  "10-6",
-			wantErr: fmt.Errorf("invalid cpu bounds from possible: begin: %d end: %d", 10, 6),
-		},
-		{
-			name:    "no parse",
-			output:  noParse,
-			wantErr: fmt.Errorf(`mismatch regex from possible: %q`, noParse),
-		},
-	} {
-		t.Run(tc.name, func(t *testing.T) {
-			threads, err := GetThreadsFromPossible([]byte(tc.output))
-
-			switch {
-			case tc.wantErr == nil:
-				if err != nil {
-					t.Fatalf("Wanted nil err, got: %v", err)
-				}
-			case err == nil:
-				t.Fatalf("Want error: %v got: %v", tc.wantErr, err)
-			default:
-				if tc.wantErr.Error() != err.Error() {
-					t.Fatalf("Want error: %v got error: %v", tc.wantErr, err)
-				}
-			}
-
-			if len(threads) != tc.wantCount {
-				t.Fatalf("Want count: %d got: %d", tc.wantCount, len(threads))
-			}
-		})
-	}
-}
-
-func TestReverseSmoke(t *testing.T) {
-	data, err := ioutil.ReadFile("/sys/devices/system/cpu/possible")
-	if err != nil {
-		t.Fatalf("Failed to read from possible: %v", err)
-	}
-	threads, err := GetThreadsFromPossible(data)
-	if err != nil {
-		t.Fatalf("Could not parse possible output: %v", err)
-	}
-
-	if len(threads) <= 0 {
-		t.Fatalf("Didn't get any CPU cores: %d", len(threads))
-	}
-}
-
-func checkSorted(threads []Thread) error {
-	if len(threads) < 2 {
-		return nil
-	}
-	last := threads[0].processorNumber
-	for _, t := range threads[1:] {
-		if last >= t.processorNumber {
-			return fmt.Errorf("threads out of order: thread %d before %d", t.processorNumber, last)
-		}
-		last = t.processorNumber
-	}
-	return nil
-}
diff --git a/runsc/mitigate/mock/mock.go b/runsc/mitigate/mock.go
index 12c59e356..4588ae2ed 100644
--- a/runsc/mitigate/mock/mock.go
+++ b/runsc/mitigate/mock.go
@@ -12,26 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package mock contains mock CPUs for mitigate tests.
-package mock
+package mitigate
 
-import "fmt"
+import "strings"
 
-// CPU represents data from CPUs that will be mitigated.
-type CPU struct {
+// MockCPU represents data from CPUs that will be mitigated.
+type MockCPU struct {
 	Name           string
 	VendorID       string
-	Family         int
-	Model          int
+	Family         int64
+	Model          int64
 	ModelName      string
 	Bugs           string
-	PhysicalCores  int
-	Cores          int
-	ThreadsPerCore int
+	PhysicalCores  int64
+	Cores          int64
+	ThreadsPerCore int64
 }
 
 // CascadeLake2 is a two core Intel CascadeLake machine.
-var CascadeLake2 = CPU{
+var CascadeLake2 = MockCPU{
 	Name:           "CascadeLake",
 	VendorID:       "GenuineIntel",
 	Family:         6,
@@ -44,7 +43,7 @@ var CascadeLake2 = CPU{
 }
 
 // CascadeLake4 is a four core Intel CascadeLake machine.
-var CascadeLake4 = CPU{
+var CascadeLake4 = MockCPU{
 	Name:           "CascadeLake",
 	VendorID:       "GenuineIntel",
 	Family:         6,
@@ -57,7 +56,7 @@ var CascadeLake4 = CPU{
 }
 
 // Haswell2 is a two core Intel Haswell machine.
-var Haswell2 = CPU{
+var Haswell2 = MockCPU{
 	Name:           "Haswell",
 	VendorID:       "GenuineIntel",
 	Family:         6,
@@ -70,7 +69,7 @@ var Haswell2 = CPU{
 }
 
 // Haswell2core is a 2 core Intel Haswell machine with no hyperthread pairs.
-var Haswell2core = CPU{
+var Haswell2core = MockCPU{
 	Name:           "Haswell2Physical",
 	VendorID:       "GenuineIntel",
 	Family:         6,
@@ -83,7 +82,7 @@ var Haswell2core = CPU{
 }
 
 // AMD2 is an two core AMD machine.
-var AMD2 = CPU{
+var AMD2 = MockCPU{
 	Name:           "AMD",
 	VendorID:       "AuthenticAMD",
 	Family:         23,
@@ -96,7 +95,7 @@ var AMD2 = CPU{
 }
 
 // AMD8 is an eight core AMD machine.
-var AMD8 = CPU{
+var AMD8 = MockCPU{
 	Name:           "AMD",
 	VendorID:       "AuthenticAMD",
 	Family:         23,
@@ -108,47 +107,39 @@ var AMD8 = CPU{
 	ThreadsPerCore: 2,
 }
 
-// MakeCPUString makes a string formated like /proc/cpuinfo for each cpuTestCase
-func (tc CPU) MakeCPUString() string {
-	template := `processor	: %d
-vendor_id	: %s
-cpu family	: %d
-model		: %d
-model name	: %s
-physical id  : %d
-core id		: %d
-cpu cores	: %d
-bugs		: %s
-
-`
+// Empty is an empty CPU set.
+var Empty = MockCPU{
+	Name: "Empty",
+}
 
-	ret := ``
-	for i := 0; i < tc.PhysicalCores; i++ {
-		for j := 0; j < tc.Cores; j++ {
-			for k := 0; k < tc.ThreadsPerCore; k++ {
+// MakeCPUSet makes a cpuSet from a MockCPU.
+func (tc MockCPU) MakeCPUSet() CPUSet {
+	bugs := make(map[string]struct{})
+	for _, bug := range strings.Split(tc.Bugs, " ") {
+		bugs[bug] = struct{}{}
+	}
+	var cpus CPUSet = []*CPU{}
+	for i := int64(0); i < tc.PhysicalCores; i++ {
+		for j := int64(0); j < tc.Cores; j++ {
+			for k := int64(0); k < tc.ThreadsPerCore; k++ {
 				processorNum := (i*tc.Cores+j)*tc.ThreadsPerCore + k
-				ret += fmt.Sprintf(template,
-					processorNum, /*processor*/
-					tc.VendorID,  /*vendor_id*/
-					tc.Family,    /*cpu family*/
-					tc.Model,     /*model*/
-					tc.ModelName, /*model name*/
-					i,            /*physical id*/
-					j,            /*core id*/
-					k,            /*cpu cores*/
-					tc.Bugs,      /*bugs*/
-				)
+				cpu := &CPU{
+					processorNumber: processorNum,
+					vendorID:        tc.VendorID,
+					cpuFamily:       tc.Family,
+					model:           tc.Model,
+					physicalID:      i,
+					coreID:          j,
+					bugs:            bugs,
+				}
+				cpus = append(cpus, cpu)
 			}
 		}
 	}
-	return ret
+	return cpus
 }
 
-// MakeSysPossibleString makes a string representing a the contents of /sys/devices/system/cpu/possible.
-func (tc CPU) MakeSysPossibleString() string {
-	max := tc.PhysicalCores * tc.Cores * tc.ThreadsPerCore
-	if max == 1 {
-		return "0"
-	}
-	return fmt.Sprintf("0-%d", max-1)
+// NumCPUs returns the number of CPUs for this CPU.
+func (tc MockCPU) NumCPUs() int {
+	return int(tc.PhysicalCores * tc.Cores * tc.ThreadsPerCore)
 }
diff --git a/runsc/mitigate/mock/BUILD b/runsc/mitigate/mock/BUILD
deleted file mode 100644
index 5019ff9ee..000000000
--- a/runsc/mitigate/mock/BUILD
+++ /dev/null
@@ -1,11 +0,0 @@
-load("//tools:defs.bzl", "go_library")
-
-package(licenses = ["notice"])
-
-go_library(
-    name = "mock",
-    srcs = ["mock.go"],
-    visibility = [
-        "//runsc:__subpackages__",
-    ],
-)