1 files changed, 115 insertions, 7 deletions
diff --git a/runsc/cmd/mitigate.go b/runsc/cmd/mitigate.go
index 822af1917..fddf0e0dd 100644
--- a/runsc/cmd/mitigate.go
+++ b/runsc/cmd/mitigate.go
@@ -16,6 +16,8 @@ package cmd
 
 import (
 	"context"
+	"fmt"
+	"io/ioutil"
 
 	"github.com/google/subcommands"
 	"gvisor.dev/gvisor/pkg/log"
@@ -23,9 +25,23 @@ import (
 	"gvisor.dev/gvisor/runsc/mitigate"
 )
 
+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"
+)
+
 // Mitigate implements subcommands.Command for the "mitigate" command.
 type Mitigate struct {
-	mitigate mitigate.Mitigate
+	// Run the command without changing the underlying system.
+	dryRun bool
+	// Reverse mitigate by turning on all CPU cores.
+	reverse bool
+	// Path to file to read to create CPUSet.
+	path string
+	// Callback to check if a given thread is vulnerable.
+	vulnerable func(other mitigate.Thread) bool
 }
 
 // Name implements subcommands.command.name.
@@ -38,14 +54,19 @@ func (*Mitigate) Synopsis() string {
 	return "mitigate mitigates the underlying system against side channel attacks"
 }
 
-// Usage implements subcommands.Command.Usage.
-func (m *Mitigate) Usage() string {
-	return m.mitigate.Usage()
+// Usage implments Usage for cmd.Mitigate.
+func (m Mitigate) Usage() string {
+	return `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.
+
+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.`
 }
 
-// SetFlags implements subcommands.Command.SetFlags.
+// SetFlags sets flags for the command Mitigate.
 func (m *Mitigate) SetFlags(f *flag.FlagSet) {
-	m.mitigate.SetFlags(f)
+	f.BoolVar(&m.dryRun, "dryrun", false, "run the command without changing system")
+	f.BoolVar(&m.reverse, "reverse", false, "reverse mitigate by enabling all CPUs")
 }
 
 // Execute implements subcommands.Command.Execute.
@@ -55,10 +76,97 @@ func (m *Mitigate) Execute(_ context.Context, f *flag.FlagSet, args ...interface
 		return subcommands.ExitUsageError
 	}
 
-	if err := m.mitigate.Execute(); err != nil {
+	m.path = cpuInfo
+	if m.reverse {
+		m.path = allPossibleCPUs
+	}
+
+	m.vulnerable = func(other mitigate.Thread) bool {
+		return other.IsVulnerable()
+	}
+
+	if _, err := m.doExecute(); err != nil {
 		log.Warningf("Execute failed: %v", err)
 		return subcommands.ExitFailure
 	}
 
 	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.")
+	}
+	if m.reverse {
+		data, err := ioutil.ReadFile(m.path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read %s: %v", m.path, err)
+		}
+
+		set, err := m.doReverse(data)
+		if err != nil {
+			return nil, fmt.Errorf("reverse operation failed: %v", err)
+		}
+		return set, nil
+	}
+
+	data, err := ioutil.ReadFile(m.path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read %s: %v", m.path, err)
+	}
+	set, err := m.doMitigate(data)
+	if err != nil {
+		return nil, fmt.Errorf("mitigate operation failed: %v", err)
+	}
+	return set, nil
+}
+
+func (m *Mitigate) doMitigate(data []byte) (mitigate.CPUSet, error) {
+	set, err := mitigate.NewCPUSet(data, m.vulnerable)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("Mitigate found the following CPUs...")
+	log.Infof("%s", set)
+
+	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: %v", t, err)
+		}
+	}
+	log.Infof("Shutdown successful.")
+	return set, nil
+}
+
+func (m *Mitigate) doReverse(data []byte) (mitigate.CPUSet, error) {
+	set, err := mitigate.NewCPUSetFromPossible(data)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("Reverse mitigate found the following CPUs...")
+	log.Infof("%s", set)
+
+	enableList := set.GetRemainingList()
+
+	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: %v", t, err)
+		}
+	}
+	log.Infof("Enable successful.")
+	return set, nil
+}