1 files changed, 80 insertions, 73 deletions
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
 }