Add reverse flag to mitigate.

Add reverse operation to mitigate that just enables
all CPUs.

PiperOrigin-RevId: 360511215

1 files changed, 71 insertions, 11 deletions

diff --git a/runsc/mitigate/mitigate.go b/runsc/mitigate/mitigate.go
index 3ea58454f..91de623e3 100644
--- a/runsc/mitigate/mitigate.go
+++ b/runsc/mitigate/mitigate.go
@@ -21,15 +21,23 @@ package mitigate
 
 import (
 	"fmt"
+	"io/ioutil"
 
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/runsc/flag"
 )
 
+const (
+	cpuInfo         = "/proc/cpuinfo"
+	allPossibleCPUs = "/sys/devices/system/cpu/possible"
+)
+
 // Mitigate handles high level mitigate operations provided to runsc.
 type Mitigate struct {
-	dryRun bool     // Run the command without changing the underlying system.
-	other  mitigate // Struct holds extra mitigate logic.
+	dryRun  bool     // Run the command without changing the underlying system.
+	reverse bool     // Reverse mitigate by turning on all CPU cores.
+	other   mitigate // Struct holds extra mitigate logic.
+	path    string   // path to read for each operation (e.g. /proc/cpuinfo).
 }
 
 // Usage implments Usage for cmd.Mitigate.
@@ -37,6 +45,8 @@ func (m Mitigate) Usage() string {
 	usageString := `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.
 `
 	return usageString + m.other.usage()
 }
@@ -44,31 +54,81 @@ Mitigate mitigates a system to the "MDS" vulnerability by implementing a manual
 // SetFlags sets flags for the command Mitigate.
 func (m Mitigate) SetFlags(f *flag.FlagSet) {
 	f.BoolVar(&m.dryRun, "dryrun", false, "run the command without changing system")
+	f.BoolVar(&m.reverse, "reverse", false, "reverse mitigate by enabling all CPUs")
 	m.other.setFlags(f)
+	m.path = cpuInfo
+	if m.reverse {
+		m.path = allPossibleCPUs
+	}
 }
 
 // Execute executes the Mitigate command.
-func (m Mitigate) Execute(data []byte) error {
+func (m Mitigate) Execute() error {
+	data, err := ioutil.ReadFile(m.path)
+	if err != nil {
+		return fmt.Errorf("failed to read %s: %v", m.path, err)
+	}
+
+	if m.reverse {
+		err := m.doReverse(data)
+		if err != nil {
+			return fmt.Errorf("reverse operation failed: %v", err)
+		}
+		return nil
+	}
+
+	set, err := m.doMitigate(data)
+	if err != nil {
+		return fmt.Errorf("mitigate operation failed: %v", err)
+	}
+	return m.other.execute(set, m.dryRun)
+}
+
+func (m Mitigate) doMitigate(data []byte) (cpuSet, error) {
 	set, err := newCPUSet(data, m.other.vulnerable)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	log.Infof("Mitigate found the following CPUs...")
 	log.Infof("%s", set)
 
-	shutdownList := set.getShutdownList()
-	log.Infof("Shutting down threads on thread pairs.")
-	for _, t := range shutdownList {
-		log.Infof("Shutting down thread: %s", t)
+	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.shutdown(); err != nil {
-			return fmt.Errorf("error shutting down thread: %s err: %v", t, err)
+		if err := t.disable(); err != nil {
+			return nil, fmt.Errorf("error disabling thread: %s err: %v", t, err)
 		}
 	}
 	log.Infof("Shutdown successful.")
-	m.other.execute(set, m.dryRun)
+	return set, nil
+}
+
+func (m Mitigate) doReverse(data []byte) error {
+	set, err := newCPUSetFromPossible(data)
+	if err != nil {
+		return 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 fmt.Errorf("error enabling thread: %s err: %v", t, err)
+		}
+	}
+	log.Infof("Enable successful.")
 	return nil
 }