src/go.chromium.org/chromiumos/test/provision/cmd/provisionserver/bootstrap/services/firmwareservice/firmwareservice.go - chromiumos/platform/dev-util - Git at Google

 // Copyright 2022 The ChromiumOS Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 package firmwareservice

 import (
 	"context"
 	"errors"
 	"fmt"
 	"log"
 	"strings"
 	"time"

 	"go.chromium.org/chromiumos/test/provision/cmd/provisionserver/bootstrap/services"
 	"go.chromium.org/chromiumos/test/provision/lib/servo_lib"
 	"go.chromium.org/chromiumos/test/provision/lib/servoadapter"

 	conf "go.chromium.org/chromiumos/config/go"
 	"go.chromium.org/chromiumos/config/go/test/api"
 	lab_api "go.chromium.org/chromiumos/config/go/test/lab/api"
 )

 // FirmwareService inherits ServiceInterface
 type FirmwareService struct {
 	// In case of flashing over SSH, |connection| connects to the DUT.
 	// In case of flashing over Servo, |connection| connects to the ServoHost.
 	connection services.ServiceAdapterInterface

 	mainRwPath *conf.StoragePath

 	mainRoPath *conf.StoragePath
 	ecRoPath   *conf.StoragePath

 	board, model string

 	force bool

 	ecChip string

 	// imagesMetadata is a map from gspath -> ImageArchiveMetadata.
 	// Allows to avoid redownloading/reprocessing archives.
 	imagesMetadata map[string]ImageArchiveMetadata

 	useServo bool
 	// servoConfig provides dut-controls and programmer argument for flashing
 	servoConfig *servo_lib.ServoConfig
 	// wrapper around servoClient connection with extra servod-related functions
 	servoConnection servoadapter.ServoHostInterface
 	servoPort       int
 }

 func NewFirmwareService(ctx context.Context, dut *lab_api.Dut, dutClient api.DutServiceClient, servoClient api.ServodServiceClient, req *api.InstallFirmwareRequest) (*FirmwareService, error) {
 	connection := services.NewServiceAdapter(dut, dutClient, false /*noReboot*/)
 	return NewFirmwareServiceFromExistingConnection(ctx, dut, connection, servoClient, req)
 }

 // NewFirmwareServiceFromExistingConnection is equivalent to the above constructor,
 // but recycles a ServiceAdapter. Generally useful for tests.
 func NewFirmwareServiceFromExistingConnection(ctx context.Context, dut *lab_api.Dut, connection services.ServiceAdapterInterface, servoClient api.ServodServiceClient, req *api.InstallFirmwareRequest) (*FirmwareService, error) {
 	if req.FirmwareConfig == nil {
 		return nil, errors.New("request.FirmwareConfig is nil")
 	}

 	// Firmware may be updated in write-protected mode, where only 'rw' regions
 	// would be update, or write-protection may be disabled (dangerous) in order
 	// to update 'ro' regions.
 	//
 	// The only 'rw' firmware is the main one, aka AP firmware, which also syncs
 	// EC 'rw'. It will be flashed with write protection turned on.
 	mainRwPath := req.FirmwareConfig.MainRwPayload.GetFirmwareImagePath()

 	// Read-only firmware images include AP, EC, and PD firmware, and will be
 	// flashed with write protection turned off.
 	mainRoPath := req.FirmwareConfig.MainRoPayload.GetFirmwareImagePath()
 	ecRoPath := req.FirmwareConfig.EcRoPayload.GetFirmwareImagePath()

 	imagesMetadata := make(map[string]ImageArchiveMetadata)

 	board := dut.GetChromeos().GetDutModel().GetBuildTarget()
 	if len(board) == 0 {
 		return nil, errors.New("\"board\" field in Dut proto is required")
 	}
 	model := dut.GetChromeos().GetDutModel().GetModelName()
 	if len(model) == 0 {
 		return nil, errors.New("\"model\" field in Dut proto is required")
 	}

 	force := req.GetForce()
 	useServo := req.GetUseServo()

 	var servoConfig *servo_lib.ServoConfig
 	var servoConnection servoadapter.ServoHostInterface
 	var ecChip string
 	var servoPort int

 	if useServo {
 		if servoClient == nil {
 			return nil, fmt.Errorf("servo use is requested, but servo client not provided")
 		}

 		// Note: dut.GetChromeos().Servo.ServodAddress.Port is the port of cros-servod service,
 		// we seem to be missing port of servod itself.
 		if servoPort == 0 {
 			servoPort = 9999
 		}

 		servoConnection = servoadapter.NewServoHostAdapterFromExecCmder(board, model, servoPort, servoClient)

 		// Ask servod for servo_type. This implicitly checks if servod is running
 		// and servo connection is working.
 		servoTypeStr, err := servoConnection.GetVariable(ctx, "servo_type")
 		if err != nil {
 			return nil, fmt.Errorf("failed to get servo_type: %w. "+
 				"Is servod running on port %v and connected to the DUT?",
 				err, servoPort)
 		}

 		// ask servod for serial number of the connected DUT.
 		servoType := servo_lib.NewServoType(servoTypeStr)

 		if servoType.IsMultipleServos() {
 			// Handle dual servo.
 			// We need CCD if ec_ro is set, otherwise, servo_micro will be faster.
 			preferCCD := false
 			if req.FirmwareConfig.EcRoPayload.GetFirmwareImagePath() != nil {
 				preferCCD = true
 			}
 			servoSubtypeStr := servoType.PickServoSubtype(preferCCD)
 			servoType = servo_lib.NewServoType(servoSubtypeStr)
 		}

 		serial, err := servoConnection.GetVariable(ctx, servoType.GetSerialNumberOption())
 		if err != nil {
 			return nil, err
 		}

 		// finally, get the correct config
 		servoConfig, err = servo_lib.GetServoConfig(board, serial, servoType)
 		if err != nil {
 			return nil, err
 		}

 		ecChip, err = servoConnection.GetVariable(ctx, "ec_chip")
 		if err != nil {
 			return nil, err
 		}
 	}

 	fws := FirmwareService{
 		connection,
 		mainRwPath,
 		mainRoPath,
 		ecRoPath,
 		board,
 		model,
 		force,
 		ecChip,
 		imagesMetadata,
 		useServo,
 		servoConfig,
 		servoConnection,
 		servoPort,
 	}

 	fws.PrintRequestInfo()

 	if !fws.UpdateRo() && !fws.UpdateRw() {
 		return nil, errors.New("request.FirmwareConfig: no paths to images specified")
 	}

 	return &fws, nil
 }

 func (fws *FirmwareService) PrintRequestInfo() {
 	informationString := "provisioning "

 	images := []string{}
 	if fws.mainRwPath != nil {
 		images = append(images, "AP(RW)")
 	}
 	if fws.mainRoPath != nil {
 		images = append(images, "AP(RO)")
 	}
 	if fws.ecRoPath != nil {
 		images = append(images, "EC(RO)")
 	}
 	informationString += strings.Join(images, " and ") + " firmware"

 	flashMode := "SSH"
 	if fws.useServo {
 		flashMode = fmt.Sprintf("Servo %s", (fws.servoConfig.ServoType))
 	}
 	informationString += " over " + flashMode + ". "

 	informationString += "Board: " + fws.board + ". "

 	informationString += "Model: " + fws.model + ". "

 	informationString += fmt.Sprintf("Force: %v.", fws.force)

 	log.Println("[FW Provisioning]", informationString)
 }

 func (fws *FirmwareService) UpdateRw() bool {
 	return fws.mainRwPath != nil
 }

 func (fws *FirmwareService) UpdateRo() bool {
 	return (fws.mainRoPath != nil) || (fws.ecRoPath != nil)
 }

 // GetBoard returns board of the DUT to provision. Returns empty string if board is not known.
 func (fws *FirmwareService) GetBoard() string {
 	return fws.board
 }

 // RestartDut restarts the DUT using one of the available mechanisms.
 // Preferred restart method is to send "power_state:reset" command to servod.
 // If servod restarting failed/not available in the environment,
 // then this function will try to SSH to the DUT and run `restart`,
 // if |requireServoReset| is false. If |requireServoReset| is True, restart over
 // SSH will not be attempted, and the function will return an error.
 func (fws *FirmwareService) RestartDut(ctx context.Context, requireServoReset bool) error {
 	waitForReconnect := func(conn services.ServiceAdapterInterface) error {
 		// Attempts to run `true` on the DUT over SSH until |reconnectRetries| attempts.
 		const reconnectRetries = 10
 		const reconnectAttemptWait = 10 * time.Second
 		const reconnectFailPause = 10 * time.Second
 		var restartErr error
 		for i := 0; i < reconnectRetries; i++ {
 			reconnectCtx, reconnCancel := context.WithTimeout(ctx, reconnectAttemptWait)
 			defer reconnCancel()
 			_, restartErr = conn.RunCmd(reconnectCtx, "true", nil)
 			if restartErr == nil {
 				log.Println("[FW Provisioning: Restart DUT] reestablished connection to the DUT after restart.")
 				return nil
 			}
 			time.Sleep(reconnectFailPause)
 		}
 		log.Printf("[FW Provisioning: Restart DUT] timed out waiting for DUT to restart: %v\n", restartErr)
 		return restartErr
 	}

 	if requireServoReset && fws.servoConnection == nil {
 		return errors.New("servo restart is required but servo connection not available")
 	}
 	// over Servo first
 	if fws.servoConnection != nil {
 		log.Printf("[FW Provisioning: Restart DUT] restarting DUT with \"dut-control power_state:reset\" over servo.\n")
 		servoRestartErr := fws.servoConnection.RunDutControl(ctx, []string{"power_state:reset"})
 		if servoRestartErr == nil {
 			if fws.connection != nil {
 				// If SSH connection to DUT was available, wait until it's back up again.
 				return waitForReconnect(fws.connection)
 			} else {
 				waitDuration := 30 * time.Second
 				log.Printf("[FW Provisioning: Restart DUT] waiting for %v for DUT to finish rebooting.\n", waitDuration.String())
 				time.Sleep(waitDuration)
 				powerState, getPowerStateErr := fws.servoConnection.GetVariable(ctx, "ec_system_powerstate")
 				if getPowerStateErr != nil {
 					log.Printf("[FW Provisioning: Restart DUT] failed to get power state after reboot: %v\n", getPowerStateErr)
 				} else {
 					log.Printf("[FW Provisioning: Restart DUT] DUT power state after reboot: %v\n", powerState)
 				}
 				return getPowerStateErr
 			}
 		}
 		log.Printf("[FW Provisioning: Restart DUT] failed to restart DUT via Servo: %v.\n", servoRestartErr)
 		if requireServoReset {
 			return servoRestartErr
 		}
 	}

 	// over SSH if allowed by |requireServoReset|
 	if fws.connection != nil {
 		log.Printf("[FW Provisioning: Restart DUT] restarting DUT over SSH.")
 		fws.connection.Restart(ctx)
 		return waitForReconnect(fws.connection)
 	}
 	return errors.New("failed to restart: no SSH connection to the DUT")
 }

 // GetModel returns model of the DUT to provision. Returns empty string if model is not known.
 func (fws *FirmwareService) GetModel() string {
 	return fws.model
 }

 // GetFirstState returns the first state of this state machine
 func (fws *FirmwareService) GetFirstState() services.ServiceState {
 	return FirmwarePrepareState{
 		service: *fws,
 	}
 }

 // CleanupOnFailure is called if one of service's states failes to Execute() and
 // currenly cleans up the temporary files.
 func (fws *FirmwareService) CleanupOnFailure(states []services.ServiceState, executionErr error) error {
 	return fws.deleteArchiveDirectories()
 }

 func (fws *FirmwareService) deleteArchiveDirectories() error {
 	var cleanedDevice services.ServiceAdapterInterface
 	if fws.useServo {
 		// If servo is used, the files will be located on the ServoHost.
 		cleanedDevice = fws.servoConnection
 	} else {
 		// If SSH is used, the files will be located on the DUT in /tmp/
 		// It's not strictly necessary to delete them before reboot.
 		cleanedDevice = fws.connection
 	}

 	var allErrors []string
 	for _, imgMetadata := range fws.imagesMetadata {
 		err := cleanedDevice.DeleteDirectory(context.Background(), imgMetadata.archiveDir)
 		if err != nil {
 			allErrors = append(allErrors, fmt.Sprintf("failed to delete %v: %v", imgMetadata.archiveDir, err))
 		}
 	}

 	if len(allErrors) > 0 {
 		return errors.New(strings.Join(allErrors, ". "))
 	}

 	return nil
 }

 // FlashWithFutility flashes the DUT using "futility" tool.
 // futility will be run with "--mode=recovery".
 // if |rwOnly| is true, futility will flash only RW regions.
 // if |rwOnly| is false, futility will flash both RW and RO regions.
 // futilityArgs must include argument(s) that provide path(s) to the images.
 //
 // If flashing over ssh, simply calls runFutility().
 // If flashing over servo, also runs pre- and post-flashing dut-controls.
 func (fws FirmwareService) FlashWithFutility(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
 	if fws.useServo {
 		return fws.servoFlash(ctx, rwOnly, futilityImageArgs)
 	} else {
 		return fws.sshFlash(ctx, rwOnly, futilityImageArgs)
 	}
 }

 func (fws FirmwareService) sshFlash(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
 	return fws.runFutility(ctx, rwOnly, futilityImageArgs)
 }

 func (fws FirmwareService) servoFlash(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
 	dutOnErr := fws.servoConnection.RunAllDutControls(ctx, fws.servoConfig.DutOn)
 	if dutOnErr != nil {
 		return fmt.Errorf("failed to run pre-flashing dut-controls: %w", dutOnErr)
 	}

 	flashErr := fws.runFutility(ctx, rwOnly, futilityImageArgs)

 	// Attempt to run post-flashing dut-controls regardless of flashErr.
 	dutOffErr := fws.servoConnection.RunAllDutControls(ctx, fws.servoConfig.DutOff)
 	if flashErr != nil {
 		return flashErr
 	}
 	if dutOffErr != nil {
 		return fmt.Errorf("failed to run post-flashing dut-controls: %w", dutOffErr)
 	}
 	return nil
 }

 func (fws FirmwareService) runFutility(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
 	if len(futilityImageArgs) == 0 {
 		return fmt.Errorf("unable to flash: no futility Image args provided")
 	}

 	futilityArgs := futilityImageArgs
 	futilityArgs = append([]string{"update", "--mode=recovery"}, futilityArgs...)

 	if rwOnly {
 		futilityArgs = append(futilityArgs, "--wp=1")
 	} else {
 		futilityArgs = append(futilityArgs, "--wp=0")
 	}

 	if fws.force {
 		futilityArgs = append(futilityArgs, "--force")
 	}

 	if fws.useServo {
 		futilityArgs = append(futilityArgs, "-p", fws.servoConfig.Programmer)
 		futilityArgs = append(futilityArgs, fws.servoConfig.ExtraArgs...)
 		// TODO(sfrolov): extra args from fw-config.json
 	}

 	connection := fws.GetConnectionToFlashingDevice()
 	if _, err := connection.RunCmd(ctx, "futility", futilityArgs); err != nil {
 		return err
 	}
 	return nil
 }

 // GetConnectionToFlashingDevice returns connection to the device that stores the
 // firmware image and runs futility.
 // Returns connection to ServoHost if fws.useServo, connection to DUT otherwise.
 func (fws FirmwareService) GetConnectionToFlashingDevice() services.ServiceAdapterInterface {
 	if fws.useServo {
 		return fws.servoConnection
 	} else {
 		return fws.connection
 	}
 }

 // ProvisionWithFlashEC flashes EC image using flash_ec script.
 func (fws FirmwareService) ProvisionWithFlashEC(ctx context.Context, ecImage, flashECScriptPath string) error {
 	customBitbangRate := ""
 	if fws.ecChip == "stm32" {
 		customBitbangRate = "--bitbang_rate=57600"
 	}
 	flashCmdArgs := fmt.Sprintf("--ro --chip=%s --board=%s --image=%s --port=%v %s --verify --verbose",
 		fws.ecChip, fws.model, ecImage, fws.servoPort, customBitbangRate)
 	output, err := fws.GetConnectionToFlashingDevice().RunCmd(ctx, flashECScriptPath, strings.Split(flashCmdArgs, " "))
 	if err != nil {
 		log.Println(output)
 		return err
 	}
 	return nil
 }
	// Copyright 2022 The ChromiumOS Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	package firmwareservice

	import (
	"context"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	"go.chromium.org/chromiumos/test/provision/cmd/provisionserver/bootstrap/services"
	"go.chromium.org/chromiumos/test/provision/lib/servo_lib"
	"go.chromium.org/chromiumos/test/provision/lib/servoadapter"

	conf "go.chromium.org/chromiumos/config/go"
	"go.chromium.org/chromiumos/config/go/test/api"
	lab_api "go.chromium.org/chromiumos/config/go/test/lab/api"
	)

	// FirmwareService inherits ServiceInterface
	type FirmwareService struct {
	// In case of flashing over SSH, \|connection\| connects to the DUT.
	// In case of flashing over Servo, \|connection\| connects to the ServoHost.
	connection services.ServiceAdapterInterface

	mainRwPath *conf.StoragePath

	mainRoPath *conf.StoragePath
	ecRoPath *conf.StoragePath

	board, model string

	force bool

	ecChip string

	// imagesMetadata is a map from gspath -> ImageArchiveMetadata.
	// Allows to avoid redownloading/reprocessing archives.
	imagesMetadata map[string]ImageArchiveMetadata

	useServo bool
	// servoConfig provides dut-controls and programmer argument for flashing
	servoConfig *servo_lib.ServoConfig
	// wrapper around servoClient connection with extra servod-related functions
	servoConnection servoadapter.ServoHostInterface
	servoPort int
	}

	func NewFirmwareService(ctx context.Context, dut lab_api.Dut, dutClient api.DutServiceClient, servoClient api.ServodServiceClient, req api.InstallFirmwareRequest) (*FirmwareService, error) {
	connection := services.NewServiceAdapter(dut, dutClient, false /noReboot/)
	return NewFirmwareServiceFromExistingConnection(ctx, dut, connection, servoClient, req)
	}

	// NewFirmwareServiceFromExistingConnection is equivalent to the above constructor,
	// but recycles a ServiceAdapter. Generally useful for tests.
	func NewFirmwareServiceFromExistingConnection(ctx context.Context, dut lab_api.Dut, connection services.ServiceAdapterInterface, servoClient api.ServodServiceClient, req api.InstallFirmwareRequest) (*FirmwareService, error) {
	if req.FirmwareConfig == nil {
	return nil, errors.New("request.FirmwareConfig is nil")
	}

	// Firmware may be updated in write-protected mode, where only 'rw' regions
	// would be update, or write-protection may be disabled (dangerous) in order
	// to update 'ro' regions.
	//
	// The only 'rw' firmware is the main one, aka AP firmware, which also syncs
	// EC 'rw'. It will be flashed with write protection turned on.
	mainRwPath := req.FirmwareConfig.MainRwPayload.GetFirmwareImagePath()

	// Read-only firmware images include AP, EC, and PD firmware, and will be
	// flashed with write protection turned off.
	mainRoPath := req.FirmwareConfig.MainRoPayload.GetFirmwareImagePath()
	ecRoPath := req.FirmwareConfig.EcRoPayload.GetFirmwareImagePath()

	imagesMetadata := make(map[string]ImageArchiveMetadata)

	board := dut.GetChromeos().GetDutModel().GetBuildTarget()
	if len(board) == 0 {
	return nil, errors.New("\"board\" field in Dut proto is required")
	}
	model := dut.GetChromeos().GetDutModel().GetModelName()
	if len(model) == 0 {
	return nil, errors.New("\"model\" field in Dut proto is required")
	}

	force := req.GetForce()
	useServo := req.GetUseServo()

	var servoConfig *servo_lib.ServoConfig
	var servoConnection servoadapter.ServoHostInterface
	var ecChip string
	var servoPort int

	if useServo {
	if servoClient == nil {
	return nil, fmt.Errorf("servo use is requested, but servo client not provided")
	}

	// Note: dut.GetChromeos().Servo.ServodAddress.Port is the port of cros-servod service,
	// we seem to be missing port of servod itself.
	if servoPort == 0 {
	servoPort = 9999
	}

	servoConnection = servoadapter.NewServoHostAdapterFromExecCmder(board, model, servoPort, servoClient)

	// Ask servod for servo_type. This implicitly checks if servod is running
	// and servo connection is working.
	servoTypeStr, err := servoConnection.GetVariable(ctx, "servo_type")
	if err != nil {
	return nil, fmt.Errorf("failed to get servo_type: %w. "+
	"Is servod running on port %v and connected to the DUT?",
	err, servoPort)
	}

	// ask servod for serial number of the connected DUT.
	servoType := servo_lib.NewServoType(servoTypeStr)

	if servoType.IsMultipleServos() {
	// Handle dual servo.
	// We need CCD if ec_ro is set, otherwise, servo_micro will be faster.
	preferCCD := false
	if req.FirmwareConfig.EcRoPayload.GetFirmwareImagePath() != nil {
	preferCCD = true
	}
	servoSubtypeStr := servoType.PickServoSubtype(preferCCD)
	servoType = servo_lib.NewServoType(servoSubtypeStr)
	}

	serial, err := servoConnection.GetVariable(ctx, servoType.GetSerialNumberOption())
	if err != nil {
	return nil, err
	}

	// finally, get the correct config
	servoConfig, err = servo_lib.GetServoConfig(board, serial, servoType)
	if err != nil {
	return nil, err
	}

	ecChip, err = servoConnection.GetVariable(ctx, "ec_chip")
	if err != nil {
	return nil, err
	}
	}

	fws := FirmwareService{
	connection,
	mainRwPath,
	mainRoPath,
	ecRoPath,
	board,
	model,
	force,
	ecChip,
	imagesMetadata,
	useServo,
	servoConfig,
	servoConnection,
	servoPort,
	}

	fws.PrintRequestInfo()

	if !fws.UpdateRo() && !fws.UpdateRw() {
	return nil, errors.New("request.FirmwareConfig: no paths to images specified")
	}

	return &fws, nil
	}

	func (fws *FirmwareService) PrintRequestInfo() {
	informationString := "provisioning "

	images := []string{}
	if fws.mainRwPath != nil {
	images = append(images, "AP(RW)")
	}
	if fws.mainRoPath != nil {
	images = append(images, "AP(RO)")
	}
	if fws.ecRoPath != nil {
	images = append(images, "EC(RO)")
	}
	informationString += strings.Join(images, " and ") + " firmware"

	flashMode := "SSH"
	if fws.useServo {
	flashMode = fmt.Sprintf("Servo %s", (fws.servoConfig.ServoType))
	}
	informationString += " over " + flashMode + ". "

	informationString += "Board: " + fws.board + ". "

	informationString += "Model: " + fws.model + ". "

	informationString += fmt.Sprintf("Force: %v.", fws.force)

	log.Println("[FW Provisioning]", informationString)
	}

	func (fws *FirmwareService) UpdateRw() bool {
	return fws.mainRwPath != nil
	}

	func (fws *FirmwareService) UpdateRo() bool {
	return (fws.mainRoPath != nil) \|\| (fws.ecRoPath != nil)
	}

	// GetBoard returns board of the DUT to provision. Returns empty string if board is not known.
	func (fws *FirmwareService) GetBoard() string {
	return fws.board
	}

	// RestartDut restarts the DUT using one of the available mechanisms.
	// Preferred restart method is to send "power_state:reset" command to servod.
	// If servod restarting failed/not available in the environment,
	// then this function will try to SSH to the DUT and run `restart`,
	// if \|requireServoReset\| is false. If \|requireServoReset\| is True, restart over
	// SSH will not be attempted, and the function will return an error.
	func (fws *FirmwareService) RestartDut(ctx context.Context, requireServoReset bool) error {
	waitForReconnect := func(conn services.ServiceAdapterInterface) error {
	// Attempts to run `true` on the DUT over SSH until \|reconnectRetries\| attempts.
	const reconnectRetries = 10
	const reconnectAttemptWait = 10 * time.Second
	const reconnectFailPause = 10 * time.Second
	var restartErr error
	for i := 0; i < reconnectRetries; i++ {
	reconnectCtx, reconnCancel := context.WithTimeout(ctx, reconnectAttemptWait)
	defer reconnCancel()
	_, restartErr = conn.RunCmd(reconnectCtx, "true", nil)
	if restartErr == nil {
	log.Println("[FW Provisioning: Restart DUT] reestablished connection to the DUT after restart.")
	return nil
	}
	time.Sleep(reconnectFailPause)
	}
	log.Printf("[FW Provisioning: Restart DUT] timed out waiting for DUT to restart: %v\n", restartErr)
	return restartErr
	}

	if requireServoReset && fws.servoConnection == nil {
	return errors.New("servo restart is required but servo connection not available")
	}
	// over Servo first
	if fws.servoConnection != nil {
	log.Printf("[FW Provisioning: Restart DUT] restarting DUT with \"dut-control power_state:reset\" over servo.\n")
	servoRestartErr := fws.servoConnection.RunDutControl(ctx, []string{"power_state:reset"})
	if servoRestartErr == nil {
	if fws.connection != nil {
	// If SSH connection to DUT was available, wait until it's back up again.
	return waitForReconnect(fws.connection)
	} else {
	waitDuration := 30 * time.Second
	log.Printf("[FW Provisioning: Restart DUT] waiting for %v for DUT to finish rebooting.\n", waitDuration.String())
	time.Sleep(waitDuration)
	powerState, getPowerStateErr := fws.servoConnection.GetVariable(ctx, "ec_system_powerstate")
	if getPowerStateErr != nil {
	log.Printf("[FW Provisioning: Restart DUT] failed to get power state after reboot: %v\n", getPowerStateErr)
	} else {
	log.Printf("[FW Provisioning: Restart DUT] DUT power state after reboot: %v\n", powerState)
	}
	return getPowerStateErr
	}
	}
	log.Printf("[FW Provisioning: Restart DUT] failed to restart DUT via Servo: %v.\n", servoRestartErr)
	if requireServoReset {
	return servoRestartErr
	}
	}

	// over SSH if allowed by \|requireServoReset\|
	if fws.connection != nil {
	log.Printf("[FW Provisioning: Restart DUT] restarting DUT over SSH.")
	fws.connection.Restart(ctx)
	return waitForReconnect(fws.connection)
	}
	return errors.New("failed to restart: no SSH connection to the DUT")
	}

	// GetModel returns model of the DUT to provision. Returns empty string if model is not known.
	func (fws *FirmwareService) GetModel() string {
	return fws.model
	}

	// GetFirstState returns the first state of this state machine
	func (fws *FirmwareService) GetFirstState() services.ServiceState {
	return FirmwarePrepareState{
	service: *fws,
	}
	}

	// CleanupOnFailure is called if one of service's states failes to Execute() and
	// currenly cleans up the temporary files.
	func (fws *FirmwareService) CleanupOnFailure(states []services.ServiceState, executionErr error) error {
	return fws.deleteArchiveDirectories()
	}

	func (fws *FirmwareService) deleteArchiveDirectories() error {
	var cleanedDevice services.ServiceAdapterInterface
	if fws.useServo {
	// If servo is used, the files will be located on the ServoHost.
	cleanedDevice = fws.servoConnection
	} else {
	// If SSH is used, the files will be located on the DUT in /tmp/
	// It's not strictly necessary to delete them before reboot.
	cleanedDevice = fws.connection
	}

	var allErrors []string
	for _, imgMetadata := range fws.imagesMetadata {
	err := cleanedDevice.DeleteDirectory(context.Background(), imgMetadata.archiveDir)
	if err != nil {
	allErrors = append(allErrors, fmt.Sprintf("failed to delete %v: %v", imgMetadata.archiveDir, err))
	}
	}

	if len(allErrors) > 0 {
	return errors.New(strings.Join(allErrors, ". "))
	}

	return nil
	}

	// FlashWithFutility flashes the DUT using "futility" tool.
	// futility will be run with "--mode=recovery".
	// if \|rwOnly\| is true, futility will flash only RW regions.
	// if \|rwOnly\| is false, futility will flash both RW and RO regions.
	// futilityArgs must include argument(s) that provide path(s) to the images.
	//
	// If flashing over ssh, simply calls runFutility().
	// If flashing over servo, also runs pre- and post-flashing dut-controls.
	func (fws FirmwareService) FlashWithFutility(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
	if fws.useServo {
	return fws.servoFlash(ctx, rwOnly, futilityImageArgs)
	} else {
	return fws.sshFlash(ctx, rwOnly, futilityImageArgs)
	}
	}

	func (fws FirmwareService) sshFlash(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
	return fws.runFutility(ctx, rwOnly, futilityImageArgs)
	}

	func (fws FirmwareService) servoFlash(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
	dutOnErr := fws.servoConnection.RunAllDutControls(ctx, fws.servoConfig.DutOn)
	if dutOnErr != nil {
	return fmt.Errorf("failed to run pre-flashing dut-controls: %w", dutOnErr)
	}

	flashErr := fws.runFutility(ctx, rwOnly, futilityImageArgs)

	// Attempt to run post-flashing dut-controls regardless of flashErr.
	dutOffErr := fws.servoConnection.RunAllDutControls(ctx, fws.servoConfig.DutOff)
	if flashErr != nil {
	return flashErr
	}
	if dutOffErr != nil {
	return fmt.Errorf("failed to run post-flashing dut-controls: %w", dutOffErr)
	}
	return nil
	}

	func (fws FirmwareService) runFutility(ctx context.Context, rwOnly bool, futilityImageArgs []string) error {
	if len(futilityImageArgs) == 0 {
	return fmt.Errorf("unable to flash: no futility Image args provided")
	}

	futilityArgs := futilityImageArgs
	futilityArgs = append([]string{"update", "--mode=recovery"}, futilityArgs...)

	if rwOnly {
	futilityArgs = append(futilityArgs, "--wp=1")
	} else {
	futilityArgs = append(futilityArgs, "--wp=0")
	}

	if fws.force {
	futilityArgs = append(futilityArgs, "--force")
	}

	if fws.useServo {
	futilityArgs = append(futilityArgs, "-p", fws.servoConfig.Programmer)
	futilityArgs = append(futilityArgs, fws.servoConfig.ExtraArgs...)
	// TODO(sfrolov): extra args from fw-config.json
	}

	connection := fws.GetConnectionToFlashingDevice()
	if _, err := connection.RunCmd(ctx, "futility", futilityArgs); err != nil {
	return err
	}
	return nil
	}

	// GetConnectionToFlashingDevice returns connection to the device that stores the
	// firmware image and runs futility.
	// Returns connection to ServoHost if fws.useServo, connection to DUT otherwise.
	func (fws FirmwareService) GetConnectionToFlashingDevice() services.ServiceAdapterInterface {
	if fws.useServo {
	return fws.servoConnection
	} else {
	return fws.connection
	}
	}

	// ProvisionWithFlashEC flashes EC image using flash_ec script.
	func (fws FirmwareService) ProvisionWithFlashEC(ctx context.Context, ecImage, flashECScriptPath string) error {
	customBitbangRate := ""
	if fws.ecChip == "stm32" {
	customBitbangRate = "--bitbang_rate=57600"
	}
	flashCmdArgs := fmt.Sprintf("--ro --chip=%s --board=%s --image=%s --port=%v %s --verify --verbose",
	fws.ecChip, fws.model, ecImage, fws.servoPort, customBitbangRate)
	output, err := fws.GetConnectionToFlashingDevice().RunCmd(ctx, flashECScriptPath, strings.Split(flashCmdArgs, " "))
	if err != nil {
	log.Println(output)
	return err
	}
	return nil
	}