[SMT] Refactor runsc mititgate

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

PiperOrigin-RevId: 389215975

3 files changed, 178 insertions, 182 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
 }