gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/runsc/cmd/mitigate_test.go (about)

     1  // Copyright 2021 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  //go:build amd64
    16  // +build amd64
    17  
    18  package cmd
    19  
    20  import (
    21  	"testing"
    22  
    23  	"github.com/google/subcommands"
    24  	"gvisor.dev/gvisor/pkg/log"
    25  	"gvisor.dev/gvisor/runsc/cmd/util"
    26  	"gvisor.dev/gvisor/runsc/mitigate"
    27  )
    28  
    29  type mockMachineControl struct {
    30  	enabled bool
    31  	cpus    mitigate.CPUSet
    32  }
    33  
    34  func (m *mockMachineControl) enable() error {
    35  	m.enabled = true
    36  	return nil
    37  }
    38  
    39  func (m *mockMachineControl) disable() error {
    40  	if m.cpus.IsVulnerable() {
    41  		m.enabled = false
    42  	}
    43  	return nil
    44  }
    45  
    46  func (m *mockMachineControl) isEnabled() (bool, error) {
    47  	return m.enabled, nil
    48  }
    49  
    50  func (m *mockMachineControl) getCPUs() (mitigate.CPUSet, error) {
    51  	set := m.cpus
    52  	if !m.enabled {
    53  		set = m.cpus[:len(m.cpus)/2]
    54  	}
    55  
    56  	// Instead of just returning the created CPU set stored in this struct, call
    57  	// NewCPUSet to exercise that code path as the machineControlImpl would.
    58  	return mitigate.NewCPUSet(set.String())
    59  }
    60  
    61  type executeTestCase struct {
    62  	name                string
    63  	cpu                 mitigate.MockCPU
    64  	mitigateWantCPUs    int
    65  	mitigateError       subcommands.ExitStatus
    66  	mitigateWantEnabled bool
    67  	reverseWantCPUs     int
    68  	reverseError        subcommands.ExitStatus
    69  	reverseWantEnabled  bool
    70  	dryrun              bool
    71  }
    72  
    73  func TestExecute(t *testing.T) {
    74  	for _, tc := range []executeTestCase{
    75  		{
    76  			name:                "CascadeLake4",
    77  			cpu:                 mitigate.CascadeLake4,
    78  			mitigateWantCPUs:    2,
    79  			mitigateWantEnabled: false,
    80  			reverseWantCPUs:     4,
    81  			reverseWantEnabled:  true,
    82  		},
    83  		{
    84  			name:                "CascadeLake4DryRun",
    85  			cpu:                 mitigate.CascadeLake4,
    86  			mitigateWantCPUs:    4,
    87  			mitigateWantEnabled: true,
    88  			reverseWantCPUs:     4,
    89  			reverseWantEnabled:  true,
    90  			dryrun:              true,
    91  		},
    92  		{
    93  			name:                "AMD8",
    94  			cpu:                 mitigate.AMD8,
    95  			mitigateWantCPUs:    8,
    96  			mitigateWantEnabled: true,
    97  			reverseWantCPUs:     8,
    98  			reverseWantEnabled:  true,
    99  		},
   100  		{
   101  			name:          "Empty",
   102  			cpu:           mitigate.Empty,
   103  			mitigateError: util.Errorf(`mitigate operation failed: no cpus found for: ""`),
   104  			reverseError:  util.Errorf(`mitigate operation failed: no cpus found for: ""`),
   105  		},
   106  	} {
   107  		t.Run(tc.name, func(t *testing.T) {
   108  			set := tc.cpu.MakeCPUSet()
   109  			m := &Mitigate{
   110  				control: &mockMachineControl{
   111  					enabled: true,
   112  					cpus:    set,
   113  				},
   114  				dryRun: tc.dryrun,
   115  			}
   116  			t.Run("Mitigate", func(t *testing.T) {
   117  				m.doExecuteTest(t, tc.mitigateWantEnabled, tc.mitigateWantCPUs, tc.mitigateError)
   118  			})
   119  
   120  			m.reverse = true
   121  			t.Run("Reverse", func(t *testing.T) {
   122  				m.doExecuteTest(t, tc.reverseWantEnabled, tc.reverseWantCPUs, tc.reverseError)
   123  			})
   124  		})
   125  	}
   126  }
   127  
   128  // doExecuteTest runs Execute with the mitigate operation and reverse operation.
   129  func (m *Mitigate) doExecuteTest(t *testing.T, wantEnabled bool, wantCPUs int, wantErr subcommands.ExitStatus) {
   130  	subError := m.execute()
   131  	if subError != wantErr {
   132  		t.Fatalf("Mitigate error mismatch: want: %v got: %v", wantErr, subError)
   133  	}
   134  
   135  	// case where test should end in error or we don't care
   136  	// about how many cpus are returned.
   137  	if wantErr != subcommands.ExitSuccess {
   138  		log.Infof("return")
   139  		return
   140  	}
   141  
   142  	gotEnabled, _ := m.control.isEnabled()
   143  	if wantEnabled != gotEnabled {
   144  		t.Fatalf("Incorrect enabled state: want: %t got: %t", wantEnabled, gotEnabled)
   145  	}
   146  
   147  	gotCPUs, _ := m.control.getCPUs()
   148  	if len(gotCPUs) != wantCPUs {
   149  		t.Fatalf("Incorrect number of CPUs: want: %d got: %d", wantCPUs, len(gotCPUs))
   150  	}
   151  }