Add CPUSet for runsc mitigate.

PiperOrigin-RevId: 355242055

2 files changed, 336 insertions, 58 deletions

diff --git a/runsc/mitigate/cpu.go b/runsc/mitigate/cpu.go
index 113b98159..ae4ce9579 100644
--- a/runsc/mitigate/cpu.go
+++ b/runsc/mitigate/cpu.go
@@ -16,6 +16,7 @@ package mitigate
 
 import (
 	"fmt"
+	"io/ioutil"
 	"regexp"
 	"strconv"
 	"strings"
@@ -31,16 +32,104 @@ const (
 )
 
 const (
-	processorKey = "processor"
-	vendorIDKey  = "vendor_id"
-	cpuFamilyKey = "cpu family"
-	modelKey     = "model"
-	coreIDKey    = "core id"
-	bugsKey      = "bugs"
+	processorKey  = "processor"
+	vendorIDKey   = "vendor_id"
+	cpuFamilyKey  = "cpu family"
+	modelKey      = "model"
+	physicalIDKey = "physical id"
+	coreIDKey     = "core id"
+	bugsKey       = "bugs"
 )
 
-// getCPUSet returns cpu structs from reading /proc/cpuinfo.
-func getCPUSet(data string) ([]*cpu, error) {
+const (
+	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[cpuID]*threadGroup
+
+// newCPUSet creates a CPUSet from data read from /proc/cpuinfo.
+func newCPUSet(data []byte, vulnerable func(*thread) bool) (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 || vulnerable(p)
+		core.threads = append(core.threads, p)
+	}
+	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) {
 	// Each processor entry should start with the
 	// processor key. Find the beginings of each.
 	r := buildRegex(processorKey, `\d+`)
@@ -56,13 +145,13 @@ func getCPUSet(data string) ([]*cpu, 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.
-	var cpus = make([]*cpu, 0, len(indices)-1)
+	var cpus = make([]*thread, 0, len(indices)-1)
 	// 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 := getCPU(data[start:end])
+		c, err := newThread(data[start:end])
 		if err != nil {
 			return nil, err
 		}
@@ -71,18 +160,25 @@ func getCPUSet(data string) ([]*cpu, error) {
 	return cpus, nil
 }
 
+// cpuID for each thread is defined by the physical and
+// core IDs. If equal, two threads are Hyperthread pairs.
+type cpuID struct {
+	physicalID int64
+	coreID     int64
+}
+
 // type cpu represents pertinent info about a cpu.
-type cpu struct {
+type thread 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).
-	coreID          int64               // This CPU's core id to match Hyperthread Pairs
+	id              cpuID               // id for this thread
 	bugs            map[string]struct{} // map of vulnerabilities parsed from the 'bugs' field.
 }
 
-// getCPU parses a CPU from a single cpu entry from /proc/cpuinfo.
-func getCPU(data string) (*cpu, error) {
+// newThread parses a CPU from a single cpu entry from /proc/cpuinfo.
+func newThread(data string) (*thread, error) {
 	processor, err := parseProcessor(data)
 	if err != nil {
 		return nil, err
@@ -103,6 +199,11 @@ func getCPU(data string) (*cpu, error) {
 		return nil, err
 	}
 
+	physicalID, err := parsePhysicalID(data)
+	if err != nil {
+		return nil, err
+	}
+
 	coreID, err := parseCoreID(data)
 	if err != nil {
 		return nil, err
@@ -113,16 +214,41 @@ func getCPU(data string) (*cpu, error) {
 		return nil, err
 	}
 
-	return &cpu{
+	return &thread{
 		processorNumber: processor,
 		vendorID:        vendorID,
 		cpuFamily:       cpuFamily,
 		model:           model,
-		coreID:          coreID,
-		bugs:            bugs,
+		id: cpuID{
+			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
+`
+	bugs := make([]string, 0)
+	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, ","))
+}
+
+// shutdown turns off the CPU by writing 0 to /sys/devices/cpu/cpu{N}/online.
+func (t *thread) shutdown() error {
+	cpuPath := fmt.Sprintf(cpuOnlineTemplate, t.processorNumber)
+	return ioutil.WriteFile(cpuPath, []byte{'0'}, 0644)
+}
+
 // List of pertinent side channel vulnerablilites.
 // For mds, see: https://www.kernel.org/doc/html/latest/admin-guide/hw-vuln/mds.html.
 var vulnerabilities = []string{
@@ -134,35 +260,46 @@ var vulnerabilities = []string{
 }
 
 // isVulnerable checks if a CPU is vulnerable to pertinent bugs.
-func (c *cpu) isVulnerable() bool {
+func (t *thread) isVulnerable() bool {
 	for _, bug := range vulnerabilities {
-		if _, ok := c.bugs[bug]; ok {
+		if _, ok := t.bugs[bug]; ok {
 			return true
 		}
 	}
 	return false
 }
 
+// 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 (c *cpu) similarTo(other *cpu) bool {
-	if c.vendorID != other.vendorID {
+func (t *thread) similarTo(other *thread) bool {
+	if t.vendorID != other.vendorID {
 		return false
 	}
 
-	if other.cpuFamily != c.cpuFamily {
+	if other.cpuFamily != t.cpuFamily {
 		return false
 	}
 
-	if other.model != c.model {
+	if other.model != t.model {
 		return false
 	}
 
-	if len(other.bugs) != len(c.bugs) {
+	if len(other.bugs) != len(t.bugs) {
 		return false
 	}
 
-	for bug := range c.bugs {
+	for bug := range t.bugs {
 		if _, ok := other.bugs[bug]; !ok {
 			return false
 		}
@@ -190,6 +327,11 @@ func parseModel(data string) (int64, error) {
 	return parseIntegerResult(data, modelKey)
 }
 
+// parsePhysicalID parses the physical id field.
+func parsePhysicalID(data string) (int64, error) {
+	return parseIntegerResult(data, physicalIDKey)
+}
+
 // parseCoreID parses the core id field.
 func parseCoreID(data string) (int64, error) {
 	return parseIntegerResult(data, coreIDKey)
diff --git a/runsc/mitigate/cpu_test.go b/runsc/mitigate/cpu_test.go
index 77b714a02..21c12f586 100644
--- a/runsc/mitigate/cpu_test.go
+++ b/runsc/mitigate/cpu_test.go
@@ -15,26 +15,163 @@
 package mitigate
 
 import (
+	"fmt"
 	"io/ioutil"
 	"strings"
 	"testing"
 )
 
-// CPU info for a Intel CascadeLake processor. Both Skylake and CascadeLake have
-// the same family/model numbers, but with different bugs (e.g. skylake has
-// cpu_meltdown).
-var cascadeLake = &cpu{
-	vendorID:  "GenuineIntel",
-	cpuFamily: 6,
-	model:     85,
-	bugs: map[string]struct{}{
-		"spectre_v1":        struct{}{},
-		"spectre_v2":        struct{}{},
-		"spec_store_bypass": struct{}{},
-		mds:                 struct{}{},
-		swapgs:              struct{}{},
-		taa:                 struct{}{},
-	},
+// cpuTestCase represents data from CPUs that will be mitigated.
+type cpuTestCase struct {
+	name           string
+	vendorID       string
+	family         int
+	model          int
+	modelName      string
+	bugs           string
+	physicalCores  int
+	cores          int
+	threadsPerCore int
+}
+
+var cascadeLake4 = cpuTestCase{
+	name:           "CascadeLake",
+	vendorID:       "GenuineIntel",
+	family:         6,
+	model:          85,
+	modelName:      "Intel(R) Xeon(R) CPU",
+	bugs:           "spectre_v1 spectre_v2 spec_store_bypass mds swapgs taa",
+	physicalCores:  1,
+	cores:          2,
+	threadsPerCore: 2,
+}
+
+var haswell2 = cpuTestCase{
+	name:           "Haswell",
+	vendorID:       "GenuineIntel",
+	family:         6,
+	model:          63,
+	modelName:      "Intel(R) Xeon(R) CPU",
+	bugs:           "cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs",
+	physicalCores:  1,
+	cores:          1,
+	threadsPerCore: 2,
+}
+
+var haswell2core = cpuTestCase{
+	name:           "Haswell2Physical",
+	vendorID:       "GenuineIntel",
+	family:         6,
+	model:          63,
+	modelName:      "Intel(R) Xeon(R) CPU",
+	bugs:           "cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs",
+	physicalCores:  2,
+	cores:          1,
+	threadsPerCore: 1,
+}
+
+var amd8 = cpuTestCase{
+	name:           "AMD",
+	vendorID:       "AuthenticAMD",
+	family:         23,
+	model:          49,
+	modelName:      "AMD EPYC 7B12",
+	bugs:           "sysret_ss_attrs spectre_v1 spectre_v2 spec_store_bypass",
+	physicalCores:  4,
+	cores:          1,
+	threadsPerCore: 2,
+}
+
+// makeCPUString makes a string formated like /proc/cpuinfo for each cpuTestCase
+func (tc cpuTestCase) 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
+`
+	ret := ``
+	for i := 0; i < tc.physicalCores; i++ {
+		for j := 0; j < tc.cores; j++ {
+			for k := 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*/
+					tc.cores*tc.physicalCores, /*cpu cores*/
+					tc.bugs /*bugs*/)
+			}
+		}
+	}
+	return ret
+}
+
+// TestMockCPUSet tests mock cpu test cases against the cpuSet functions.
+func TestMockCPUSet(t *testing.T) {
+	for _, tc := range []struct {
+		testCase     cpuTestCase
+		isVulnerable bool
+	}{
+		{
+			testCase:     amd8,
+			isVulnerable: false,
+		},
+		{
+			testCase:     haswell2,
+			isVulnerable: true,
+		},
+		{
+			testCase:     haswell2core,
+			isVulnerable: true,
+		},
+
+		{
+			testCase:     cascadeLake4,
+			isVulnerable: true,
+		},
+	} {
+		t.Run(tc.testCase.name, func(t *testing.T) {
+			data := tc.testCase.makeCPUString()
+			vulnerable := func(t *thread) bool {
+				return t.isVulnerable()
+			}
+			set, err := newCPUSet([]byte(data), vulnerable)
+			if err != nil {
+				t.Fatalf("Failed to ")
+			}
+			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)
+			}
+		})
+	}
 }
 
 // TestGetCPU tests basic parsing of single CPU strings from reading
@@ -44,15 +181,19 @@ func TestGetCPU(t *testing.T) {
 vendor_id	: GenuineIntel
 cpu family	: 6
 model		: 85
+physical id: 0
 core id		: 0
 bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa itlb_multihit
 `
-	want := cpu{
+	want := thread{
 		processorNumber: 0,
 		vendorID:        "GenuineIntel",
 		cpuFamily:       6,
 		model:           85,
-		coreID:          0,
+		id: cpuID{
+			physicalID: 0,
+			coreID:     0,
+		},
 		bugs: map[string]struct{}{
 			"cpu_meltdown":      struct{}{},
 			"spectre_v1":        struct{}{},
@@ -66,7 +207,7 @@ bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
 		},
 	}
 
-	got, err := getCPU(data)
+	got, err := newThread(data)
 	if err != nil {
 		t.Fatalf("getCpu failed with error: %v", err)
 	}
@@ -81,7 +222,7 @@ bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
 }
 
 func TestInvalid(t *testing.T) {
-	result, err := getCPUSet(`something not a processor`)
+	result, err := getThreads(`something not a processor`)
 	if err == nil {
 		t.Fatalf("getCPU set didn't return an error: %+v", result)
 	}
@@ -148,7 +289,7 @@ cache_alignment	: 64
 address sizes	: 46 bits physical, 48 bits virtual
 power management:
 `
-	cpuSet, err := getCPUSet(data)
+	cpuSet, err := getThreads(data)
 	if err != nil {
 		t.Fatalf("getCPUSet failed: %v", err)
 	}
@@ -158,7 +299,7 @@ power management:
 		t.Fatalf("Num CPU mismatch: want: %d, got: %d", wantCPULen, len(cpuSet))
 	}
 
-	wantCPU := cpu{
+	wantCPU := thread{
 		vendorID:  "GenuineIntel",
 		cpuFamily: 6,
 		model:     63,
@@ -187,7 +328,11 @@ func TestReadFile(t *testing.T) {
 		t.Fatalf("Failed to read cpuinfo: %v", err)
 	}
 
-	set, err := getCPUSet(string(data))
+	vulnerable := func(t *thread) bool {
+		return t.isVulnerable()
+	}
+
+	set, err := newCPUSet(data, vulnerable)
 	if err != nil {
 		t.Fatalf("Failed to parse CPU data %v\n%s", err, data)
 	}
@@ -196,9 +341,7 @@ func TestReadFile(t *testing.T) {
 		t.Fatalf("Failed to parse any CPUs: %d", len(set))
 	}
 
-	for _, c := range set {
-		t.Logf("CPU: %+v: %t", c, c.isVulnerable())
-	}
+	t.Log(set)
 }
 
 // TestVulnerable tests if the isVulnerable method is correct
@@ -332,17 +475,13 @@ power management:`
 			cpuString:  skylake,
 			vulnerable: true,
 		}, {
-			name:       "cascadeLake",
-			cpuString:  cascade,
-			vulnerable: false,
-		}, {
 			name:       "amd",
 			cpuString:  amd,
 			vulnerable: false,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			set, err := getCPUSet(tc.cpuString)
+			set, err := getThreads(tc.cpuString)
 			if err != nil {
 				t.Fatalf("Failed to getCPUSet:%v\n %s", err, tc.cpuString)
 			}
@@ -353,9 +492,6 @@ power management:`
 
 			for _, c := range set {
 				got := func() bool {
-					if cascadeLake.similarTo(c) {
-						return false
-					}
 					return c.isVulnerable()
 				}()