src/chromiumos/tast/remote/bundles/cros/firmware/write_protect.go - chromiumos/platform/tast-tests - Git at Google

 // Copyright 2021 The Chromium OS Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 package firmware

 import (
 	"context"
 	"fmt"
 	"path/filepath"
 	"regexp"
 	"strconv"
 	"time"

 	"chromiumos/tast/common/servo"
 	"chromiumos/tast/ctxutil"
 	"chromiumos/tast/errors"
 	"chromiumos/tast/remote/firmware"
 	"chromiumos/tast/remote/firmware/fixture"
 	"chromiumos/tast/remote/firmware/reporters"
 	pb "chromiumos/tast/services/cros/firmware"
 	"chromiumos/tast/ssh"
 	"chromiumos/tast/testing"
 	"chromiumos/tast/testing/hwdep"
 )

 // Create enum to specify which tests need to be run
 type wpTestRebootMethod int

 const (
 	rebootWithModeAwareReboot wpTestRebootMethod = iota
 	rebootWithShutdownCmd
 	rebootWithRebootCmd
 	rebootWithPowerBtn
 )

 func init() {
 	testing.AddTest(&testing.Test{
 		Func:         WriteProtect,
 		Desc:         "Verify enabling and disabling write protect works as expected",
 		Contacts:     []string{"tij@google.com", "cros-fw-engprod@google.com"},
 		Attr:         []string{"group:firmware", "firmware_unstable"},
 		SoftwareDeps: []string{"crossystem", "flashrom"},
 		ServiceDeps:  []string{"tast.cros.firmware.BiosService"},
 		HardwareDeps: hwdep.D(hwdep.ChromeEC()),
 		Timeout:      20 * time.Minute,
 		Params: []testing.Param{
 			{
 				Name:    "dev_mode_with_mode_aware_reboot",
 				Fixture: fixture.DevMode,
 				Val:     rebootWithModeAwareReboot,
 			},
 			{
 				Name:    "normal_mode_with_mode_aware_reboot",
 				Fixture: fixture.NormalMode,
 				Val:     rebootWithModeAwareReboot,
 			},
 			{
 				Name:    "normal_mode_with_shutdown_cmd",
 				Fixture: fixture.NormalMode,
 				Val:     rebootWithShutdownCmd,
 			},
 			{
 				Name:    "normal_mode_with_reboot_cmd",
 				Fixture: fixture.NormalMode,
 				Val:     rebootWithRebootCmd,
 			},
 			{
 				Name:    "normal_mode_with_power_btn",
 				Fixture: fixture.NormalMode,
 				Val:     rebootWithPowerBtn,
 			},
 		},
 	})
 }

 const (
 	shutdownTimeout time.Duration = 2 * time.Second
 	rebootTimeout   time.Duration = 5 * time.Second
 )

 var (
 	reFlashromSuccess = regexp.MustCompile(`SUCCESS`)
 	reFlashromFail    = regexp.MustCompile(`FAILED`)
 	reWPROfmap        = regexp.MustCompile(`WP_RO\s+(\d+)\s+(\d+)`)
 	reROFRIDfmap      = regexp.MustCompile(`RO_FRID\s+(\d+)\s+(\d+)`)
 )

 type wpTarget string

 const (
 	targetBIOS wpTarget = "bios"
 	targetEC   wpTarget = "ec"
 )

 var flashromTargets = map[wpTarget]string{
 	targetBIOS: "host",
 	targetEC:   "ec",
 }

 var tmpDirPath = filepath.Join("/", "mnt", "stateful_partition", fmt.Sprintf("flashrom_%d", time.Now().Unix()))

 func WriteProtect(ctx context.Context, s *testing.State) {
 	h := s.FixtValue().(*fixture.Value).Helper
 	rebootMethod := s.Param().(wpTestRebootMethod)
 	if err := h.RequireServo(ctx); err != nil {
 		s.Fatal("Failed to connect to servo: ", err)
 	}
 	if err := h.RequireBiosServiceClient(ctx); err != nil {
 		s.Fatal("Requiring BiosServiceClient: ", err)
 	}

 	// Back up EC_RW.
 	s.Log("Back up current EC_RW region")
 	ecrwPath, err := h.BiosServiceClient.BackupImageSection(ctx, &pb.FWBackUpSection{Section: pb.ImageSection_ECRWImageSection, Programmer: pb.Programmer_ECProgrammer})
 	if err != nil {
 		s.Fatal("Failed to backup current EC_RW region: ", err)
 	}
 	s.Log("EC_RW region backup is stored at: ", ecrwPath.Path)

 	s.Log("Create temp dir in DUT")
 	if _, err = h.DUT.Conn().CommandContext(ctx, "mkdir", "-p", tmpDirPath).Output(ssh.DumpLogOnError); err != nil {
 		s.Fatal("Failed to create temp dir: ", err)
 	}

 	cleanupContext := ctx
 	ctx, cancel := ctxutil.Shorten(ctx, 5*time.Minute)
 	defer cancel()
 	defer func(ctx context.Context) {
 		s.Log("Reset write protect to false")
 		if err := setWriteProtect(ctx, h, targetEC, false); err != nil {
 			s.Fatal("Failed to set FW write protect state: ", err)
 		}

 		// Require again here since reboots in test cause nil pointer errors otherwise.
 		if err := h.RequireBiosServiceClient(ctx); err != nil {
 			s.Fatal("Requiring BiosServiceClient: ", err)
 		}

 		// Restore EC_RW.
 		s.Log("Restore EC_RW region with backup from: ", ecrwPath.Path)
 		if _, err := h.BiosServiceClient.RestoreImageSection(ctx, ecrwPath); err != nil {
 			s.Fatal("Failed to restore EC_RW image: ", err)
 		}

 		s.Log("Delete EC backup")
 		if _, err := h.DUT.Conn().CommandContext(ctx, "rm", ecrwPath.Path).Output(ssh.DumpLogOnError); err != nil {
 			s.Fatal("Failed to delete ec backup: ", err)
 		}
 		// Clean up temp directory.
 		s.Log("Delete temp dir and contained files from DUT")
 		if _, err := h.DUT.Conn().CommandContext(ctx, "rm", "-r", tmpDirPath).Output(ssh.DumpLogOnError); err != nil {
 			s.Fatal("Failed to delete temp dir: ", err)
 		}
 	}(cleanupContext)

 	// This call takes ~= 10 mins to complete.
 	s.Log("Test wp state over reboot on target EC")
 	if err := testWPOverReboot(ctx, h, targetEC, rebootMethod); err != nil {
 		s.Fatal("Failed to preserve wp state over reboots: ", err)
 	}

 	// This call takes ~= 10 mins to complete.
 	s.Log("Test wp state over reboot on target BIOS")
 	if err := testWPOverReboot(ctx, h, targetBIOS, rebootMethod); err != nil {
 		s.Fatal("Failed to preserve wp state over reboots: ", err)
 	}

 	// This call takes ~= 2 mins to complete.
 	s.Log("Test flashrom read/write with wp on target EC")
 	if err := testReadWrite(ctx, h, targetEC); err != nil {
 		s.Fatal("Read/write behaved unexpectedly: ", err)
 	}

 	// This call takes ~= 2 mins to complete.
 	s.Log("Test flashrom read/write with wp on target BIOS")
 	if err := testReadWrite(ctx, h, targetBIOS); err != nil {
 		s.Fatal("Read/write behaved unexpectedly: ", err)
 	}

 }

 func testWPOverReboot(ctx context.Context, h *firmware.Helper, target wpTarget, testRebootMethod wpTestRebootMethod) error {
 	var rebootMethod string
 	var rebootFunc func(context.Context, *firmware.Helper) error
 	switch testRebootMethod {
 	case rebootWithModeAwareReboot:
 		rebootMethod = "mode aware reboot"
 		rebootFunc = performModeAwareReboot
 	case rebootWithShutdownCmd:
 		rebootMethod = "mode aware reboot"
 		rebootFunc = performRebootWithShutdownCmd
 	case rebootWithRebootCmd:
 		rebootMethod = "mode aware reboot"
 		rebootFunc = performRebootWithRebootCmd
 	case rebootWithPowerBtn:
 		rebootMethod = "mode aware reboot"
 		rebootFunc = performRebootWithPowerBtn
 	}

 	testing.ContextLog(ctx, "Enable Write Protect")
 	if err := setWriteProtect(ctx, h, target, true); err != nil {
 		return errors.Wrap(err, "failed to enable FW write protect state")
 	}
 	testing.ContextLog(ctx, "Reboot DUT using ", rebootMethod)
 	if err := rebootFunc(ctx, h); err != nil {
 		return errors.Wrapf(err, "failed to reboot with %q", rebootMethod)
 	}
 	testing.ContextLog(ctx, "Expect write protect state to be enabled")
 	if err := checkCrossystem(ctx, h, 1); err != nil {
 		return errors.Wrap(err, "failed to check crossystem")
 	}
 	testing.ContextLog(ctx, "Disable Write Protect")
 	if err := setWriteProtect(ctx, h, targetEC, false); err != nil {
 		return errors.Wrap(err, "failed to disable FW write protect state")
 	}
 	testing.ContextLog(ctx, "Reboot DUT using ", rebootMethod)
 	if err := rebootFunc(ctx, h); err != nil {
 		return errors.Wrapf(err, "failed to reboot with %q", rebootMethod)
 	}
 	testing.ContextLog(ctx, "Expect write protect state to be disabled")
 	if err := checkCrossystem(ctx, h, 0); err != nil {
 		return errors.Wrap(err, "failed to check crossystem")
 	}
 	return nil
 }

 func performRebootWithRebootCmd(ctx context.Context, h *firmware.Helper) error {
 	cmd := h.DUT.Conn().CommandContext(ctx, "reboot")
 	if err := cmd.Start(); err != nil {
 		return errors.Wrap(err, "failed to reboot DUT")
 	}

 	testing.ContextLog(ctx, "Sleep, wait for DUT to reboot")
 	if err := testing.Sleep(ctx, rebootTimeout); err != nil {
 		return errors.Wrapf(err, "failed to sleep for %s after initiating reboot", rebootTimeout)
 	}

 	testing.ContextLog(ctx, "Check for S0 powerstate")
 	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
 		return errors.Wrap(err, "failed to get S0 powerstate")
 	}
 	return h.WaitConnect(ctx)
 }

 func performRebootWithShutdownCmd(ctx context.Context, h *firmware.Helper) error {
 	cmd := h.DUT.Conn().CommandContext(ctx, "/sbin/shutdown", "-P", "now")
 	if err := cmd.Start(); err != nil {
 		return errors.Wrap(err, "failed to shut down DUT")
 	}

 	testing.ContextLog(ctx, "Sleep, wait for DUT to shutdown")
 	if err := testing.Sleep(ctx, shutdownTimeout); err != nil {
 		return errors.Wrapf(err, "failed to sleep for %s after initiating shutdown", shutdownTimeout)
 	}

 	testing.ContextLog(ctx, "Check for G3 or S5 powerstate")
 	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "G3", "S5"); err != nil {
 		return errors.Wrap(err, "failed to get G3 or S5 powerstate")
 	}

 	testing.ContextLog(ctx, "Power DUT back on with short press of the power button")
 	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurPress); err != nil {
 		return errors.Wrap(err, "failed to power on DUT with short press of the power button")
 	}

 	testing.ContextLog(ctx, "Check for S0 powerstate")
 	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
 		return errors.Wrap(err, "failed to get S0 powerstate")
 	}
 	return h.WaitConnect(ctx)
 }

 func performRebootWithPowerBtn(ctx context.Context, h *firmware.Helper) error {
 	testing.ContextLog(ctx, "Power DUT off with long press of the power button")
 	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurLongPress); err != nil {
 		return errors.Wrap(err, "failed to power on DUT with short press of the power button")
 	}

 	testing.ContextLog(ctx, "Sleep, wait for DUT to power off")
 	if err := testing.Sleep(ctx, rebootTimeout); err != nil {
 		return errors.Wrapf(err, "failed to sleep for %s after pressing powerbutton", rebootTimeout)
 	}

 	testing.ContextLog(ctx, "Check for G3 or S5 powerstate")
 	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "G3", "S5"); err != nil {
 		return errors.Wrap(err, "failed to get G3 or S5 powerstate")
 	}

 	testing.ContextLog(ctx, "Power DUT back on with short press of the power button")
 	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurPress); err != nil {
 		return errors.Wrap(err, "failed to power on DUT with short press of the power button")
 	}

 	testing.ContextLog(ctx, "Check for S0 powerstate")
 	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
 		return errors.Wrap(err, "failed to get S0 powerstate")
 	}
 	return h.WaitConnect(ctx)
 }

 func testReadWrite(ctx context.Context, h *firmware.Helper, target wpTarget) error {
 	// Current firmware image as read from flash.
 	roBefore := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_before.bin", target))
 	// Current firmware image with modification to test writing.
 	roTest := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_test.bin", target))
 	// Firmware as read after writing flash.
 	roAfter := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_after.bin", target))

 	testing.ContextLog(ctx, "Enable Write Protect")
 	if err := setWriteProtect(ctx, h, target, true); err != nil {
 		return errors.Wrap(err, "failed to set FW write protect state")
 	}

 	testing.ContextLogf(ctx, "Save current fw image to %q in DUT", roBefore)
 	out, err := h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roBefore)).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrap(err, "failed to run command flashrom")
 	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
 		return errors.Errorf("flashrom did not produce sucess message: %s", string(out))
 	}

 	testing.ContextLog(ctx, "Checking fmap")
 	out, err = h.DUT.Conn().CommandContext(ctx, "dump_fmap", "-p", roBefore).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrap(err, "failed to dump fmap")
 	}

 	// Format for the dumped fmap is "Name offset size".
 	wproMatch := reWPROfmap.FindSubmatch(out)
 	if wproMatch == nil {
 		return errors.New("didn't find WP_RO in fmap")
 	}
 	wproOffset, err := strconv.Atoi(string(wproMatch[1]))
 	if err != nil {
 		return errors.Wrap(err, "failed to get WP_RO offset as integer")
 	}

 	if rofridMatch := reROFRIDfmap.FindSubmatch(out); rofridMatch != nil {
 		rofridOffset, err := strconv.Atoi(string(rofridMatch[1]))
 		if err != nil {
 			return errors.Wrap(err, "failed to get offset int for RO_FRID")
 		}

 		rofridSize, err := strconv.Atoi(string(rofridMatch[2]))
 		if err != nil {
 			return errors.Wrap(err, "failed to get size int for RO_FRID")
 		}

 		if _, err := h.DUT.Conn().CommandContext(ctx, "cp", roBefore, roTest).CombinedOutput(ssh.DumpLogOnError); err != nil {
 			return errors.Wrapf(err, "failed to copy %q to %q", roBefore, roTest)
 		}

 		dd1Args := []string{
 			fmt.Sprintf("if=%s", roTest), "bs=1",
 			fmt.Sprintf("count=%d", rofridSize),
 			fmt.Sprintf("skip=%d", rofridOffset-wproOffset),
 		}
 		trArgs := []string{`"[a-zA-Z]"`, `"[A-Za-z]"`}
 		dd2Args := []string{
 			fmt.Sprintf("of=%s", roTest), "bs=1",
 			fmt.Sprintf("count=%d", rofridSize),
 			fmt.Sprintf("seek=%d", rofridOffset-wproOffset), "conv=notrunc",
 		}
 		dd1Cmd := h.DUT.Conn().CommandContext(ctx, "dd", dd1Args...)
 		trCmd := h.DUT.Conn().CommandContext(ctx, "tr", trArgs...)
 		dd2Cmd := h.DUT.Conn().CommandContext(ctx, "dd", dd2Args...)
 		trCmd.Stdin, _ = dd1Cmd.StdoutPipe()
 		dd2Cmd.Stdin, _ = trCmd.StdoutPipe()
 		if err := dd2Cmd.Start(); err != nil {
 			return errors.Wrap(err, "failed to start second dd cmd")
 		}
 		if err := trCmd.Start(); err != nil {
 			return errors.Wrap(err, "failed to start tr cmd")
 		}
 		if err := dd1Cmd.Run(); err != nil {
 			return errors.Wrap(err, "failed to run first dd cmd")
 		}
 		if err := trCmd.Wait(); err != nil {
 			return errors.Wrap(err, "failed to wait for tr cmd to complete")
 		}
 		if err := dd2Cmd.Wait(); err != nil {
 			return errors.Wrap(err, "failed to wait for second dd cmd to complete")
 		}
 	} else {
 		return errors.New("could not find RO_FRID in fmap")
 	}

 	testing.ContextLog(ctx, "Attempt to write fw with wp enabled")
 	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-w", "-i", fmt.Sprintf("WP_RO:%s", roTest)).CombinedOutput(ssh.DumpLogOnError)
 	// We expect an error, so err shouldn't be nil, but neither should out. If out is nil, then the error is from an unrelated source.
 	if out == nil {
 		return errors.New("flashrom command did not produce any output")
 	} else if match := reFlashromFail.FindSubmatch(out); match == nil {
 		return errors.Errorf("flashrom did not produce failure message when trying to write: %s", string(out))
 	}

 	testing.ContextLog(ctx, "Read fw, make sure write didn't succeed with wp enabled")
 	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roAfter)).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrap(err, "failed to run command flashrom")
 	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
 		return errors.Errorf("flashrom did not produce success message when trying to read: %s", string(out))
 	}

 	if err := performModeAwareReboot(ctx, h); err != nil {
 		return errors.Wrap(err, "failed to do a mode aware reboot")
 	}

 	out, err = h.DUT.Conn().CommandContext(ctx, "cmp", roBefore, roAfter).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrapf(err, "failed to compare %q and %q using 'cmp'", roBefore, roAfter)
 	} else if string(out) != "" {
 		return errors.Wrapf(err, "files %q and %q were not identical, so either write protect or read failed", roBefore, roAfter)
 	}

 	testing.ContextLog(ctx, "Disable Write Protect")
 	if err := setWriteProtect(ctx, h, target, false); err != nil {
 		return errors.Wrap(err, "failed to set FW write protect state")
 	}

 	testing.ContextLog(ctx, "Attempt to write fw with write protect disabled")
 	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-w", "-i", fmt.Sprintf("WP_RO:%s", roTest)).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrap(err, "failed to run command flashrom")
 	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
 		return errors.Errorf("flashrom did not produce success message when trying to write: %s", string(out))
 	}

 	testing.ContextLog(ctx, "Read fw, make sure write succeeded with wp disabled")
 	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roAfter)).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrap(err, "failed to run command flashrom")
 	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
 		return errors.Errorf("flashrom did not produce success message when trying to read: %s", string(out))
 	}

 	out, err = h.DUT.Conn().CommandContext(ctx, "cmp", roTest, roAfter).CombinedOutput(ssh.DumpLogOnError)
 	if err != nil {
 		return errors.Wrapf(err, "failed to compare %q and %q using 'cmp'", roTest, roAfter)
 	} else if string(out) != "" {
 		return errors.Errorf("Files %q and %q were not identical", roTest, roAfter)
 	}
 	return nil
 }

 func setWriteProtect(ctx context.Context, h *firmware.Helper, target wpTarget, enable bool) error {
 	enableStr := "enable"
 	fwwpState := servo.FWWPStateOn
 	if !enable {
 		enableStr = "disable"
 		fwwpState = servo.FWWPStateOff
 	}

 	if target == targetBIOS {
 		// Disable hardware wp for now so flashrom cmd can run.
 		if err := h.Servo.SetFWWPState(ctx, servo.FWWPStateOff); err != nil {
 			return errors.Wrap(err, "failed to disable firmware write protect")
 		}

 		// This file will get removed when tmpDirPath is removed.
 		tmpImagePath := filepath.Join(tmpDirPath, "tmp_file.bin")
 		flashromArgs := []string{
 			"-p", flashromTargets[target],
 			"-i", fmt.Sprintf("WP_RO:%s", tmpImagePath),
 			"--wp-region", "WP_RO",
 			fmt.Sprintf("--wp-%s", enableStr),
 		}
 		out, err := h.DUT.Conn().CommandContext(ctx, "flashrom", flashromArgs...).Output(ssh.DumpLogOnError)
 		if err != nil {
 			return errors.Wrapf(err, "failed run flashrom cmd: %s", string(out))
 		} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
 			return errors.Errorf("flashrom did not produce success message when trying to write: %s", string(out))
 		}

 		if err := h.Servo.SetFWWPState(ctx, fwwpState); err != nil {
 			return errors.Wrapf(err, "failed to %s firmware write protect", enableStr)
 		}

 	} else {
 		// Enable software wp before hardware wp if enabling.
 		if enable {
 			if err := h.Servo.RunECCommand(ctx, "flashwp enable"); err != nil {
 				return errors.Wrap(err, "failed to enable flashwp")
 			}
 		}

 		if err := h.Servo.SetFWWPState(ctx, fwwpState); err != nil {
 			return errors.Wrapf(err, "failed to %s firmware write protect", enableStr)
 		}

 		// Disable software wp after hardware wp so its allowed.
 		if !enable {
 			if err := h.Servo.RunECCommand(ctx, "flashwp disable"); err != nil {
 				return errors.Wrap(err, "failed to disable flashwp")
 			}
 		}
 	}

 	return performModeAwareReboot(ctx, h)
 }

 func performModeAwareReboot(ctx context.Context, h *firmware.Helper) error {
 	// Create new mode switcher every time to prevent nil pointer errors.
 	ms, err := firmware.NewModeSwitcher(ctx, h)
 	if err != nil {
 		return errors.Wrap(err, "failed to create mode switcher")
 	}
 	testing.ContextLog(ctx, "Performing mode aware reboot")
 	return ms.ModeAwareReboot(ctx, firmware.ColdReset)
 }

 func checkCrossystem(ctx context.Context, h *firmware.Helper, expectedWpsw int) error {
 	r := reporters.New(h.DUT)
 	testing.ContextLog(ctx, "Check crossystem for write protect state param")
 	paramMap, err := r.Crossystem(ctx, reporters.CrossystemParamWpswCur)
 	if err != nil {
 		return errors.Wrapf(err, "failed to get crossystem %v value", reporters.CrossystemParamWpswCur)
 	}
 	currWpsw, err := strconv.Atoi(paramMap[reporters.CrossystemParamWpswCur])
 	if err != nil {
 		return errors.Wrap(err, "failed to convert crossystem wpsw value to integer value")
 	}
 	testing.ContextLogf(ctx, "Current write protect state: %v, Expected state: %v", currWpsw, expectedWpsw)
 	if currWpsw != expectedWpsw {
 		return errors.Errorf("expected WP state to %v, is actually %v", expectedWpsw, currWpsw)
 	}
 	return nil
 }
	// Copyright 2021 The Chromium OS Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	package firmware

	import (
	"context"
	"fmt"
	"path/filepath"
	"regexp"
	"strconv"
	"time"

	"chromiumos/tast/common/servo"
	"chromiumos/tast/ctxutil"
	"chromiumos/tast/errors"
	"chromiumos/tast/remote/firmware"
	"chromiumos/tast/remote/firmware/fixture"
	"chromiumos/tast/remote/firmware/reporters"
	pb "chromiumos/tast/services/cros/firmware"
	"chromiumos/tast/ssh"
	"chromiumos/tast/testing"
	"chromiumos/tast/testing/hwdep"
	)

	// Create enum to specify which tests need to be run
	type wpTestRebootMethod int

	const (
	rebootWithModeAwareReboot wpTestRebootMethod = iota
	rebootWithShutdownCmd
	rebootWithRebootCmd
	rebootWithPowerBtn
	)

	func init() {
	testing.AddTest(&testing.Test{
	Func: WriteProtect,
	Desc: "Verify enabling and disabling write protect works as expected",
	Contacts: []string{"tij@google.com", "cros-fw-engprod@google.com"},
	Attr: []string{"group:firmware", "firmware_unstable"},
	SoftwareDeps: []string{"crossystem", "flashrom"},
	ServiceDeps: []string{"tast.cros.firmware.BiosService"},
	HardwareDeps: hwdep.D(hwdep.ChromeEC()),
	Timeout: 20 * time.Minute,
	Params: []testing.Param{
	{
	Name: "dev_mode_with_mode_aware_reboot",
	Fixture: fixture.DevMode,
	Val: rebootWithModeAwareReboot,
	},
	{
	Name: "normal_mode_with_mode_aware_reboot",
	Fixture: fixture.NormalMode,
	Val: rebootWithModeAwareReboot,
	},
	{
	Name: "normal_mode_with_shutdown_cmd",
	Fixture: fixture.NormalMode,
	Val: rebootWithShutdownCmd,
	},
	{
	Name: "normal_mode_with_reboot_cmd",
	Fixture: fixture.NormalMode,
	Val: rebootWithRebootCmd,
	},
	{
	Name: "normal_mode_with_power_btn",
	Fixture: fixture.NormalMode,
	Val: rebootWithPowerBtn,
	},
	},
	})
	}

	const (
	shutdownTimeout time.Duration = 2 * time.Second
	rebootTimeout time.Duration = 5 * time.Second
	)

	var (
	reFlashromSuccess = regexp.MustCompile(`SUCCESS`)
	reFlashromFail = regexp.MustCompile(`FAILED`)
	reWPROfmap = regexp.MustCompile(`WP_RO\s+(\d+)\s+(\d+)`)
	reROFRIDfmap = regexp.MustCompile(`RO_FRID\s+(\d+)\s+(\d+)`)
	)

	type wpTarget string

	const (
	targetBIOS wpTarget = "bios"
	targetEC wpTarget = "ec"
	)

	var flashromTargets = map[wpTarget]string{
	targetBIOS: "host",
	targetEC: "ec",
	}

	var tmpDirPath = filepath.Join("/", "mnt", "stateful_partition", fmt.Sprintf("flashrom_%d", time.Now().Unix()))

	func WriteProtect(ctx context.Context, s *testing.State) {
	h := s.FixtValue().(*fixture.Value).Helper
	rebootMethod := s.Param().(wpTestRebootMethod)
	if err := h.RequireServo(ctx); err != nil {
	s.Fatal("Failed to connect to servo: ", err)
	}
	if err := h.RequireBiosServiceClient(ctx); err != nil {
	s.Fatal("Requiring BiosServiceClient: ", err)
	}

	// Back up EC_RW.
	s.Log("Back up current EC_RW region")
	ecrwPath, err := h.BiosServiceClient.BackupImageSection(ctx, &pb.FWBackUpSection{Section: pb.ImageSection_ECRWImageSection, Programmer: pb.Programmer_ECProgrammer})
	if err != nil {
	s.Fatal("Failed to backup current EC_RW region: ", err)
	}
	s.Log("EC_RW region backup is stored at: ", ecrwPath.Path)

	s.Log("Create temp dir in DUT")
	if _, err = h.DUT.Conn().CommandContext(ctx, "mkdir", "-p", tmpDirPath).Output(ssh.DumpLogOnError); err != nil {
	s.Fatal("Failed to create temp dir: ", err)
	}

	cleanupContext := ctx
	ctx, cancel := ctxutil.Shorten(ctx, 5*time.Minute)
	defer cancel()
	defer func(ctx context.Context) {
	s.Log("Reset write protect to false")
	if err := setWriteProtect(ctx, h, targetEC, false); err != nil {
	s.Fatal("Failed to set FW write protect state: ", err)
	}

	// Require again here since reboots in test cause nil pointer errors otherwise.
	if err := h.RequireBiosServiceClient(ctx); err != nil {
	s.Fatal("Requiring BiosServiceClient: ", err)
	}

	// Restore EC_RW.
	s.Log("Restore EC_RW region with backup from: ", ecrwPath.Path)
	if _, err := h.BiosServiceClient.RestoreImageSection(ctx, ecrwPath); err != nil {
	s.Fatal("Failed to restore EC_RW image: ", err)
	}

	s.Log("Delete EC backup")
	if _, err := h.DUT.Conn().CommandContext(ctx, "rm", ecrwPath.Path).Output(ssh.DumpLogOnError); err != nil {
	s.Fatal("Failed to delete ec backup: ", err)
	}
	// Clean up temp directory.
	s.Log("Delete temp dir and contained files from DUT")
	if _, err := h.DUT.Conn().CommandContext(ctx, "rm", "-r", tmpDirPath).Output(ssh.DumpLogOnError); err != nil {
	s.Fatal("Failed to delete temp dir: ", err)
	}
	}(cleanupContext)

	// This call takes ~= 10 mins to complete.
	s.Log("Test wp state over reboot on target EC")
	if err := testWPOverReboot(ctx, h, targetEC, rebootMethod); err != nil {
	s.Fatal("Failed to preserve wp state over reboots: ", err)
	}

	// This call takes ~= 10 mins to complete.
	s.Log("Test wp state over reboot on target BIOS")
	if err := testWPOverReboot(ctx, h, targetBIOS, rebootMethod); err != nil {
	s.Fatal("Failed to preserve wp state over reboots: ", err)
	}

	// This call takes ~= 2 mins to complete.
	s.Log("Test flashrom read/write with wp on target EC")
	if err := testReadWrite(ctx, h, targetEC); err != nil {
	s.Fatal("Read/write behaved unexpectedly: ", err)
	}

	// This call takes ~= 2 mins to complete.
	s.Log("Test flashrom read/write with wp on target BIOS")
	if err := testReadWrite(ctx, h, targetBIOS); err != nil {
	s.Fatal("Read/write behaved unexpectedly: ", err)
	}

	}

	func testWPOverReboot(ctx context.Context, h *firmware.Helper, target wpTarget, testRebootMethod wpTestRebootMethod) error {
	var rebootMethod string
	var rebootFunc func(context.Context, *firmware.Helper) error
	switch testRebootMethod {
	case rebootWithModeAwareReboot:
	rebootMethod = "mode aware reboot"
	rebootFunc = performModeAwareReboot
	case rebootWithShutdownCmd:
	rebootMethod = "mode aware reboot"
	rebootFunc = performRebootWithShutdownCmd
	case rebootWithRebootCmd:
	rebootMethod = "mode aware reboot"
	rebootFunc = performRebootWithRebootCmd
	case rebootWithPowerBtn:
	rebootMethod = "mode aware reboot"
	rebootFunc = performRebootWithPowerBtn
	}

	testing.ContextLog(ctx, "Enable Write Protect")
	if err := setWriteProtect(ctx, h, target, true); err != nil {
	return errors.Wrap(err, "failed to enable FW write protect state")
	}
	testing.ContextLog(ctx, "Reboot DUT using ", rebootMethod)
	if err := rebootFunc(ctx, h); err != nil {
	return errors.Wrapf(err, "failed to reboot with %q", rebootMethod)
	}
	testing.ContextLog(ctx, "Expect write protect state to be enabled")
	if err := checkCrossystem(ctx, h, 1); err != nil {
	return errors.Wrap(err, "failed to check crossystem")
	}
	testing.ContextLog(ctx, "Disable Write Protect")
	if err := setWriteProtect(ctx, h, targetEC, false); err != nil {
	return errors.Wrap(err, "failed to disable FW write protect state")
	}
	testing.ContextLog(ctx, "Reboot DUT using ", rebootMethod)
	if err := rebootFunc(ctx, h); err != nil {
	return errors.Wrapf(err, "failed to reboot with %q", rebootMethod)
	}
	testing.ContextLog(ctx, "Expect write protect state to be disabled")
	if err := checkCrossystem(ctx, h, 0); err != nil {
	return errors.Wrap(err, "failed to check crossystem")
	}
	return nil
	}

	func performRebootWithRebootCmd(ctx context.Context, h *firmware.Helper) error {
	cmd := h.DUT.Conn().CommandContext(ctx, "reboot")
	if err := cmd.Start(); err != nil {
	return errors.Wrap(err, "failed to reboot DUT")
	}

	testing.ContextLog(ctx, "Sleep, wait for DUT to reboot")
	if err := testing.Sleep(ctx, rebootTimeout); err != nil {
	return errors.Wrapf(err, "failed to sleep for %s after initiating reboot", rebootTimeout)
	}

	testing.ContextLog(ctx, "Check for S0 powerstate")
	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
	return errors.Wrap(err, "failed to get S0 powerstate")
	}
	return h.WaitConnect(ctx)
	}

	func performRebootWithShutdownCmd(ctx context.Context, h *firmware.Helper) error {
	cmd := h.DUT.Conn().CommandContext(ctx, "/sbin/shutdown", "-P", "now")
	if err := cmd.Start(); err != nil {
	return errors.Wrap(err, "failed to shut down DUT")
	}

	testing.ContextLog(ctx, "Sleep, wait for DUT to shutdown")
	if err := testing.Sleep(ctx, shutdownTimeout); err != nil {
	return errors.Wrapf(err, "failed to sleep for %s after initiating shutdown", shutdownTimeout)
	}

	testing.ContextLog(ctx, "Check for G3 or S5 powerstate")
	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "G3", "S5"); err != nil {
	return errors.Wrap(err, "failed to get G3 or S5 powerstate")
	}

	testing.ContextLog(ctx, "Power DUT back on with short press of the power button")
	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurPress); err != nil {
	return errors.Wrap(err, "failed to power on DUT with short press of the power button")
	}

	testing.ContextLog(ctx, "Check for S0 powerstate")
	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
	return errors.Wrap(err, "failed to get S0 powerstate")
	}
	return h.WaitConnect(ctx)
	}

	func performRebootWithPowerBtn(ctx context.Context, h *firmware.Helper) error {
	testing.ContextLog(ctx, "Power DUT off with long press of the power button")
	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurLongPress); err != nil {
	return errors.Wrap(err, "failed to power on DUT with short press of the power button")
	}

	testing.ContextLog(ctx, "Sleep, wait for DUT to power off")
	if err := testing.Sleep(ctx, rebootTimeout); err != nil {
	return errors.Wrapf(err, "failed to sleep for %s after pressing powerbutton", rebootTimeout)
	}

	testing.ContextLog(ctx, "Check for G3 or S5 powerstate")
	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "G3", "S5"); err != nil {
	return errors.Wrap(err, "failed to get G3 or S5 powerstate")
	}

	testing.ContextLog(ctx, "Power DUT back on with short press of the power button")
	if err := h.Servo.KeypressWithDuration(ctx, servo.PowerKey, servo.DurPress); err != nil {
	return errors.Wrap(err, "failed to power on DUT with short press of the power button")
	}

	testing.ContextLog(ctx, "Check for S0 powerstate")
	if err := h.WaitForPowerStates(ctx, firmware.PowerStateInterval, firmware.PowerStateTimeout, "S0"); err != nil {
	return errors.Wrap(err, "failed to get S0 powerstate")
	}
	return h.WaitConnect(ctx)
	}

	func testReadWrite(ctx context.Context, h *firmware.Helper, target wpTarget) error {
	// Current firmware image as read from flash.
	roBefore := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_before.bin", target))
	// Current firmware image with modification to test writing.
	roTest := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_test.bin", target))
	// Firmware as read after writing flash.
	roAfter := filepath.Join(tmpDirPath, fmt.Sprintf("%s_ro_after.bin", target))

	testing.ContextLog(ctx, "Enable Write Protect")
	if err := setWriteProtect(ctx, h, target, true); err != nil {
	return errors.Wrap(err, "failed to set FW write protect state")
	}

	testing.ContextLogf(ctx, "Save current fw image to %q in DUT", roBefore)
	out, err := h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roBefore)).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrap(err, "failed to run command flashrom")
	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce sucess message: %s", string(out))
	}

	testing.ContextLog(ctx, "Checking fmap")
	out, err = h.DUT.Conn().CommandContext(ctx, "dump_fmap", "-p", roBefore).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrap(err, "failed to dump fmap")
	}

	// Format for the dumped fmap is "Name offset size".
	wproMatch := reWPROfmap.FindSubmatch(out)
	if wproMatch == nil {
	return errors.New("didn't find WP_RO in fmap")
	}
	wproOffset, err := strconv.Atoi(string(wproMatch[1]))
	if err != nil {
	return errors.Wrap(err, "failed to get WP_RO offset as integer")
	}

	if rofridMatch := reROFRIDfmap.FindSubmatch(out); rofridMatch != nil {
	rofridOffset, err := strconv.Atoi(string(rofridMatch[1]))
	if err != nil {
	return errors.Wrap(err, "failed to get offset int for RO_FRID")
	}

	rofridSize, err := strconv.Atoi(string(rofridMatch[2]))
	if err != nil {
	return errors.Wrap(err, "failed to get size int for RO_FRID")
	}

	if _, err := h.DUT.Conn().CommandContext(ctx, "cp", roBefore, roTest).CombinedOutput(ssh.DumpLogOnError); err != nil {
	return errors.Wrapf(err, "failed to copy %q to %q", roBefore, roTest)
	}

	dd1Args := []string{
	fmt.Sprintf("if=%s", roTest), "bs=1",
	fmt.Sprintf("count=%d", rofridSize),
	fmt.Sprintf("skip=%d", rofridOffset-wproOffset),
	}
	trArgs := []string{`"[a-zA-Z]"`, `"[A-Za-z]"`}
	dd2Args := []string{
	fmt.Sprintf("of=%s", roTest), "bs=1",
	fmt.Sprintf("count=%d", rofridSize),
	fmt.Sprintf("seek=%d", rofridOffset-wproOffset), "conv=notrunc",
	}
	dd1Cmd := h.DUT.Conn().CommandContext(ctx, "dd", dd1Args...)
	trCmd := h.DUT.Conn().CommandContext(ctx, "tr", trArgs...)
	dd2Cmd := h.DUT.Conn().CommandContext(ctx, "dd", dd2Args...)
	trCmd.Stdin, _ = dd1Cmd.StdoutPipe()
	dd2Cmd.Stdin, _ = trCmd.StdoutPipe()
	if err := dd2Cmd.Start(); err != nil {
	return errors.Wrap(err, "failed to start second dd cmd")
	}
	if err := trCmd.Start(); err != nil {
	return errors.Wrap(err, "failed to start tr cmd")
	}
	if err := dd1Cmd.Run(); err != nil {
	return errors.Wrap(err, "failed to run first dd cmd")
	}
	if err := trCmd.Wait(); err != nil {
	return errors.Wrap(err, "failed to wait for tr cmd to complete")
	}
	if err := dd2Cmd.Wait(); err != nil {
	return errors.Wrap(err, "failed to wait for second dd cmd to complete")
	}
	} else {
	return errors.New("could not find RO_FRID in fmap")
	}

	testing.ContextLog(ctx, "Attempt to write fw with wp enabled")
	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-w", "-i", fmt.Sprintf("WP_RO:%s", roTest)).CombinedOutput(ssh.DumpLogOnError)
	// We expect an error, so err shouldn't be nil, but neither should out. If out is nil, then the error is from an unrelated source.
	if out == nil {
	return errors.New("flashrom command did not produce any output")
	} else if match := reFlashromFail.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce failure message when trying to write: %s", string(out))
	}

	testing.ContextLog(ctx, "Read fw, make sure write didn't succeed with wp enabled")
	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roAfter)).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrap(err, "failed to run command flashrom")
	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce success message when trying to read: %s", string(out))
	}

	if err := performModeAwareReboot(ctx, h); err != nil {
	return errors.Wrap(err, "failed to do a mode aware reboot")
	}

	out, err = h.DUT.Conn().CommandContext(ctx, "cmp", roBefore, roAfter).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrapf(err, "failed to compare %q and %q using 'cmp'", roBefore, roAfter)
	} else if string(out) != "" {
	return errors.Wrapf(err, "files %q and %q were not identical, so either write protect or read failed", roBefore, roAfter)
	}

	testing.ContextLog(ctx, "Disable Write Protect")
	if err := setWriteProtect(ctx, h, target, false); err != nil {
	return errors.Wrap(err, "failed to set FW write protect state")
	}

	testing.ContextLog(ctx, "Attempt to write fw with write protect disabled")
	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-w", "-i", fmt.Sprintf("WP_RO:%s", roTest)).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrap(err, "failed to run command flashrom")
	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce success message when trying to write: %s", string(out))
	}

	testing.ContextLog(ctx, "Read fw, make sure write succeeded with wp disabled")
	out, err = h.DUT.Conn().CommandContext(ctx, "flashrom", "-p", flashromTargets[target], "-r", "-i", fmt.Sprintf("WP_RO:%s", roAfter)).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrap(err, "failed to run command flashrom")
	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce success message when trying to read: %s", string(out))
	}

	out, err = h.DUT.Conn().CommandContext(ctx, "cmp", roTest, roAfter).CombinedOutput(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrapf(err, "failed to compare %q and %q using 'cmp'", roTest, roAfter)
	} else if string(out) != "" {
	return errors.Errorf("Files %q and %q were not identical", roTest, roAfter)
	}
	return nil
	}

	func setWriteProtect(ctx context.Context, h *firmware.Helper, target wpTarget, enable bool) error {
	enableStr := "enable"
	fwwpState := servo.FWWPStateOn
	if !enable {
	enableStr = "disable"
	fwwpState = servo.FWWPStateOff
	}

	if target == targetBIOS {
	// Disable hardware wp for now so flashrom cmd can run.
	if err := h.Servo.SetFWWPState(ctx, servo.FWWPStateOff); err != nil {
	return errors.Wrap(err, "failed to disable firmware write protect")
	}

	// This file will get removed when tmpDirPath is removed.
	tmpImagePath := filepath.Join(tmpDirPath, "tmp_file.bin")
	flashromArgs := []string{
	"-p", flashromTargets[target],
	"-i", fmt.Sprintf("WP_RO:%s", tmpImagePath),
	"--wp-region", "WP_RO",
	fmt.Sprintf("--wp-%s", enableStr),
	}
	out, err := h.DUT.Conn().CommandContext(ctx, "flashrom", flashromArgs...).Output(ssh.DumpLogOnError)
	if err != nil {
	return errors.Wrapf(err, "failed run flashrom cmd: %s", string(out))
	} else if match := reFlashromSuccess.FindSubmatch(out); match == nil {
	return errors.Errorf("flashrom did not produce success message when trying to write: %s", string(out))
	}

	if err := h.Servo.SetFWWPState(ctx, fwwpState); err != nil {
	return errors.Wrapf(err, "failed to %s firmware write protect", enableStr)
	}

	} else {
	// Enable software wp before hardware wp if enabling.
	if enable {
	if err := h.Servo.RunECCommand(ctx, "flashwp enable"); err != nil {
	return errors.Wrap(err, "failed to enable flashwp")
	}
	}

	if err := h.Servo.SetFWWPState(ctx, fwwpState); err != nil {
	return errors.Wrapf(err, "failed to %s firmware write protect", enableStr)
	}

	// Disable software wp after hardware wp so its allowed.
	if !enable {
	if err := h.Servo.RunECCommand(ctx, "flashwp disable"); err != nil {
	return errors.Wrap(err, "failed to disable flashwp")
	}
	}
	}

	return performModeAwareReboot(ctx, h)
	}

	func performModeAwareReboot(ctx context.Context, h *firmware.Helper) error {
	// Create new mode switcher every time to prevent nil pointer errors.
	ms, err := firmware.NewModeSwitcher(ctx, h)
	if err != nil {
	return errors.Wrap(err, "failed to create mode switcher")
	}
	testing.ContextLog(ctx, "Performing mode aware reboot")
	return ms.ModeAwareReboot(ctx, firmware.ColdReset)
	}

	func checkCrossystem(ctx context.Context, h *firmware.Helper, expectedWpsw int) error {
	r := reporters.New(h.DUT)
	testing.ContextLog(ctx, "Check crossystem for write protect state param")
	paramMap, err := r.Crossystem(ctx, reporters.CrossystemParamWpswCur)
	if err != nil {
	return errors.Wrapf(err, "failed to get crossystem %v value", reporters.CrossystemParamWpswCur)
	}
	currWpsw, err := strconv.Atoi(paramMap[reporters.CrossystemParamWpswCur])
	if err != nil {
	return errors.Wrap(err, "failed to convert crossystem wpsw value to integer value")
	}
	testing.ContextLogf(ctx, "Current write protect state: %v, Expected state: %v", currWpsw, expectedWpsw)
	if currWpsw != expectedWpsw {
	return errors.Errorf("expected WP state to %v, is actually %v", expectedWpsw, currWpsw)
	}
	return nil
	}